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