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