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