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