1use backend::{
  2    Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScratchOwned, VecZnxDft,
  3    ZnxInfos, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{elem::Infos, glwe_ciphertext_fourier::GLWECiphertextFourier};
  8
  9#[derive(Clone, Copy, Debug)]
 10pub(crate) enum SecretDistribution {
 11    TernaryFixed(usize), // Ternary with fixed Hamming weight
 12    TernaryProb(f64),    // Ternary with probabilistic Hamming weight
 13    ZERO,                // Debug mod
 14    NONE,
 15}
 16
 17pub struct SecretKey<T> {
 18    pub(crate) data: ScalarZnx<T>,
 19    pub(crate) dist: SecretDistribution,
 20}
 21
 22impl SecretKey<Vec<u8>> {
 23    pub fn alloc<B: Backend>(module: &Module<B>, rank: usize) -> Self {
 24        Self {
 25            data: module.new_scalar_znx(rank),
 26            dist: SecretDistribution::NONE,
 27        }
 28    }
 29
 30    pub fn bytes_of(module: &Module<FFT64>, rank: usize) -> usize {
 31        module.bytes_of_scalar_znx(rank + 1)
 32    }
 33}
 34
 35impl<DataSelf> SecretKey<DataSelf> {
 36    pub fn n(&self) -> usize {
 37        self.data.n()
 38    }
 39
 40    pub fn log_n(&self) -> usize {
 41        self.data.log_n()
 42    }
 43
 44    pub fn rank(&self) -> usize {
 45        self.data.cols()
 46    }
 47}
 48
 49impl<S: AsMut<[u8]> + AsRef<[u8]>> SecretKey<S> {
 50    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 51        (0..self.rank()).for_each(|i| {
 52            self.data.fill_ternary_prob(i, prob, source);
 53        });
 54        self.dist = SecretDistribution::TernaryProb(prob);
 55    }
 56
 57    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 58        (0..self.rank()).for_each(|i| {
 59            self.data.fill_ternary_hw(i, hw, source);
 60        });
 61        self.dist = SecretDistribution::TernaryFixed(hw);
 62    }
 63
 64    pub fn fill_zero(&mut self) {
 65        self.data.zero();
 66        self.dist = SecretDistribution::ZERO;
 67    }
 68}
 69
 70pub struct SecretKeyFourier<T, B: Backend> {
 71    pub(crate) data: ScalarZnxDft<T, B>,
 72    pub(crate) dist: SecretDistribution,
 73}
 74
 75impl<DataSelf, B: Backend> SecretKeyFourier<DataSelf, B> {
 76    pub fn n(&self) -> usize {
 77        self.data.n()
 78    }
 79
 80    pub fn log_n(&self) -> usize {
 81        self.data.log_n()
 82    }
 83
 84    pub fn rank(&self) -> usize {
 85        self.data.cols()
 86    }
 87}
 88
 89impl<B: Backend> SecretKeyFourier<Vec<u8>, B> {
 90    pub fn alloc(module: &Module<B>, rank: usize) -> Self {
 91        Self {
 92            data: module.new_scalar_znx_dft(rank),
 93            dist: SecretDistribution::NONE,
 94        }
 95    }
 96
 97    pub fn bytes_of(module: &Module<B>, rank: usize) -> usize {
 98        module.bytes_of_scalar_znx_dft(rank + 1)
 99    }
100}
101
102impl<D: AsRef<[u8]> + AsMut<[u8]>> SecretKeyFourier<D, FFT64> {
103    pub fn dft<S: AsRef<[u8]>>(&mut self, module: &Module<FFT64>, sk: &SecretKey<S>) {
104        #[cfg(debug_assertions)]
105        {
106            match sk.dist {
107                SecretDistribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
108                _ => {}
109            }
110
111            assert_eq!(self.n(), module.n());
112            assert_eq!(sk.n(), module.n());
113            assert_eq!(self.rank(), sk.rank());
114        }
115
116        (0..self.rank()).for_each(|i| {
117            module.svp_prepare(&mut self.data, i, &sk.data, i);
118        });
119        self.dist = sk.dist;
120    }
121}
122
123pub struct GLWEPublicKey<D, B: Backend> {
124    pub(crate) data: GLWECiphertextFourier<D, B>,
125    pub(crate) dist: SecretDistribution,
126}
127
128impl<B: Backend> GLWEPublicKey<Vec<u8>, B> {
129    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
130        Self {
131            data: GLWECiphertextFourier::alloc(module, basek, k, rank),
132            dist: SecretDistribution::NONE,
133        }
134    }
135
136    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
137        GLWECiphertextFourier::<Vec<u8>, B>::bytes_of(module, basek, k, rank)
138    }
139}
140
141impl<T, B: Backend> Infos for GLWEPublicKey<T, B> {
142    type Inner = VecZnxDft<T, B>;
143
144    fn inner(&self) -> &Self::Inner {
145        &self.data.data
146    }
147
148    fn basek(&self) -> usize {
149        self.data.basek
150    }
151
152    fn k(&self) -> usize {
153        self.data.k
154    }
155}
156
157impl<T, B: Backend> GLWEPublicKey<T, B> {
158    pub fn rank(&self) -> usize {
159        self.cols() - 1
160    }
161}
162
163impl<C: AsRef<[u8]> + AsMut<[u8]>> GLWEPublicKey<C, FFT64> {
164    pub fn generate_from_sk<S: AsRef<[u8]>>(
165        &mut self,
166        module: &Module<FFT64>,
167        sk_dft: &SecretKeyFourier<S, FFT64>,
168        source_xa: &mut Source,
169        source_xe: &mut Source,
170        sigma: f64,
171    ) {
172        #[cfg(debug_assertions)]
173        {
174            match sk_dft.dist {
175                SecretDistribution::NONE => panic!("invalid sk_dft: SecretDistribution::NONE"),
176                _ => {}
177            }
178        }
179
180        // Its ok to allocate scratch space here since pk is usually generated only once.
181        let mut scratch: ScratchOwned = ScratchOwned::new(GLWECiphertextFourier::encrypt_sk_scratch_space(
182            module,
183            self.basek(),
184            self.k(),
185            self.rank(),
186        ));
187
188        self.data.encrypt_zero_sk(
189            module,
190            sk_dft,
191            source_xa,
192            source_xe,
193            sigma,
194            scratch.borrow(),
195        );
196        self.dist = sk_dft.dist;
197    }
198}