1use backend::{
  2    Backend, FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc,
  3    VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    elem::Infos, ggsw_ciphertext::GGSWCiphertext, glwe_ciphertext::GLWECiphertext, glwe_plaintext::GLWEPlaintext,
  9    keys::SecretKeyFourier, keyswitch_key::GLWESwitchingKey, utils::derive_size,
 10};
 11
 12pub struct GLWECiphertextFourier<C, B: Backend> {
 13    pub data: VecZnxDft<C, B>,
 14    pub basek: usize,
 15    pub k: usize,
 16}
 17
 18impl<B: Backend> GLWECiphertextFourier<Vec<u8>, B> {
 19    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
 20        Self {
 21            data: module.new_vec_znx_dft(rank + 1, derive_size(basek, k)),
 22            basek: basek,
 23            k: k,
 24        }
 25    }
 26}
 27
 28impl<T, B: Backend> Infos for GLWECiphertextFourier<T, B> {
 29    type Inner = VecZnxDft<T, B>;
 30
 31    fn inner(&self) -> &Self::Inner {
 32        &self.data
 33    }
 34
 35    fn basek(&self) -> usize {
 36        self.basek
 37    }
 38
 39    fn k(&self) -> usize {
 40        self.k
 41    }
 42}
 43
 44impl<T, B: Backend> GLWECiphertextFourier<T, B> {
 45    pub fn rank(&self) -> usize {
 46        self.cols() - 1
 47    }
 48}
 49
 50impl GLWECiphertextFourier<Vec<u8>, FFT64> {
 51    #[allow(dead_code)]
 52    pub(crate) fn idft_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
 53        module.bytes_of_vec_znx(1, size) + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())
 54    }
 55
 56    pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, ct_size: usize) -> usize {
 57        module.bytes_of_vec_znx(rank + 1, ct_size) + GLWECiphertext::encrypt_sk_scratch_space(module, ct_size)
 58    }
 59
 60    pub fn decrypt_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
 61        (module.vec_znx_big_normalize_tmp_bytes()
 62            | module.bytes_of_vec_znx_dft(1, ct_size)
 63            | (module.bytes_of_vec_znx_big(1, ct_size) + module.vec_znx_idft_tmp_bytes()))
 64            + module.bytes_of_vec_znx_big(1, ct_size)
 65    }
 66
 67    pub fn keyswitch_scratch_space(
 68        module: &Module<FFT64>,
 69        out_size: usize,
 70        out_rank: usize,
 71        in_size: usize,
 72        in_rank: usize,
 73        ksk_size: usize,
 74    ) -> usize {
 75        module.bytes_of_vec_znx(out_rank + 1, out_size)
 76            + GLWECiphertext::keyswitch_from_fourier_scratch_space(module, out_size, out_rank, in_size, in_rank, ksk_size)
 77    }
 78
 79    pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
 80        Self::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, ksk_size)
 81    }
 82
 83    pub fn external_product_scratch_space(
 84        module: &Module<FFT64>,
 85        out_size: usize,
 86        in_size: usize,
 87        ggsw_size: usize,
 88        rank: usize,
 89    ) -> usize {
 90        let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
 91        let vmp: usize = module.vmp_apply_tmp_bytes(out_size, in_size, in_size, rank + 1, rank + 1, ggsw_size);
 92        let res_small: usize = module.bytes_of_vec_znx(rank + 1, out_size);
 93        let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
 94
 95        res_dft + (vmp | (res_small + normalize))
 96    }
 97
 98    pub fn external_product_inplace_scratch_space(
 99        module: &Module<FFT64>,
100        out_size: usize,
101        ggsw_size: usize,
102        rank: usize,
103    ) -> usize {
104        Self::external_product_scratch_space(module, out_size, out_size, ggsw_size, rank)
105    }
106}
107
108impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
109    pub fn encrypt_zero_sk<DataSk: AsRef<[u8]>>(
110        &mut self,
111        module: &Module<FFT64>,
112        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
113        source_xa: &mut Source,
114        source_xe: &mut Source,
115        sigma: f64,
116        scratch: &mut Scratch,
117    ) {
118        let (vec_znx_tmp, scratch_1) = scratch.tmp_vec_znx(module, self.rank() + 1, self.size());
119        let mut ct_idft = GLWECiphertext {
120            data: vec_znx_tmp,
121            basek: self.basek,
122            k: self.k,
123        };
124        ct_idft.encrypt_zero_sk(module, sk_dft, source_xa, source_xe, sigma, scratch_1);
125
126        ct_idft.dft(module, self);
127    }
128
129    pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
130        &mut self,
131        module: &Module<FFT64>,
132        lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
133        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
134        scratch: &mut Scratch,
135    ) {
136        let cols_out: usize = rhs.rank_out() + 1;
137
138        // Space fr normalized VMP result outside of DFT domain
139        let (res_idft_data, scratch1) = scratch.tmp_vec_znx(module, cols_out, lhs.size());
140
141        let mut res_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
142            data: res_idft_data,
143            basek: lhs.basek,
144            k: lhs.k,
145        };
146
147        res_idft.keyswitch_from_fourier(module, lhs, rhs, scratch1);
148
149        (0..cols_out).for_each(|i| {
150            module.vec_znx_dft(&mut self.data, i, &res_idft.data, i);
151        });
152    }
153
154    pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
155        &mut self,
156        module: &Module<FFT64>,
157        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
158        scratch: &mut Scratch,
159    ) {
160        unsafe {
161            let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
162            self.keyswitch(&module, &*self_ptr, rhs, scratch);
163        }
164    }
165
166    pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
167        &mut self,
168        module: &Module<FFT64>,
169        lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
170        rhs: &GGSWCiphertext<DataRhs, FFT64>,
171        scratch: &mut Scratch,
172    ) {
173        let basek: usize = self.basek();
174
175        #[cfg(debug_assertions)]
176        {
177            assert_eq!(rhs.rank(), lhs.rank());
178            assert_eq!(rhs.rank(), self.rank());
179            assert_eq!(self.basek(), basek);
180            assert_eq!(lhs.basek(), basek);
181            assert_eq!(rhs.n(), module.n());
182            assert_eq!(self.n(), module.n());
183            assert_eq!(lhs.n(), module.n());
184        }
185
186        let cols: usize = rhs.rank() + 1;
187
188        // Space for VMP result in DFT domain and high precision
189        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
190
191        {
192            module.vmp_apply(&mut res_dft, &lhs.data, &rhs.data, scratch1);
193        }
194
195        // VMP result in high precision
196        let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
197
198        // Space for VMP result normalized
199        let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
200        (0..cols).for_each(|i| {
201            module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
202            module.vec_znx_dft(&mut self.data, i, &res_small, i);
203        });
204    }
205
206    pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
207        &mut self,
208        module: &Module<FFT64>,
209        rhs: &GGSWCiphertext<DataRhs, FFT64>,
210        scratch: &mut Scratch,
211    ) {
212        unsafe {
213            let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
214            self.external_product(&module, &*self_ptr, rhs, scratch);
215        }
216    }
217}
218
219impl<DataSelf: AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
220    pub fn decrypt<DataPt: AsRef<[u8]> + AsMut<[u8]>, DataSk: AsRef<[u8]>>(
221        &self,
222        module: &Module<FFT64>,
223        pt: &mut GLWEPlaintext<DataPt>,
224        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
225        scratch: &mut Scratch,
226    ) {
227        #[cfg(debug_assertions)]
228        {
229            assert_eq!(self.rank(), sk_dft.rank());
230            assert_eq!(self.n(), module.n());
231            assert_eq!(pt.n(), module.n());
232            assert_eq!(sk_dft.n(), module.n());
233        }
234
235        let cols = self.rank() + 1;
236
237        let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
238        pt_big.zero();
239
240        {
241            (1..cols).for_each(|i| {
242                let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
243                module.svp_apply(&mut ci_dft, 0, &sk_dft.data, i - 1, &self.data, i);
244                let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
245                module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
246            });
247        }
248
249        {
250            let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
251            // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
252            module.vec_znx_idft(&mut c0_big, 0, &self.data, 0, scratch_2);
253            module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
254        }
255
256        // pt = norm(BIG(m + e))
257        module.vec_znx_big_normalize(self.basek(), &mut pt.data, 0, &mut pt_big, 0, scratch_1);
258
259        pt.basek = self.basek();
260        pt.k = pt.k().min(self.k());
261    }
262
263    #[allow(dead_code)]
264    pub(crate) fn idft<DataRes: AsRef<[u8]> + AsMut<[u8]>>(
265        &self,
266        module: &Module<FFT64>,
267        res: &mut GLWECiphertext<DataRes>,
268        scratch: &mut Scratch,
269    ) {
270        #[cfg(debug_assertions)]
271        {
272            assert_eq!(self.rank(), res.rank());
273            assert_eq!(self.basek(), res.basek())
274        }
275
276        let min_size: usize = self.size().min(res.size());
277
278        let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
279
280        (0..self.rank() + 1).for_each(|i| {
281            module.vec_znx_idft(&mut res_big, 0, &self.data, i, scratch1);
282            module.vec_znx_big_normalize(self.basek(), &mut res.data, i, &res_big, 0, scratch1);
283        });
284    }
285}