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        bound: f64,
126        scratch: &mut Scratch,
127    ) where
128        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
129    {
130        let (vec_znx_tmp, scratch_1) = scratch.tmp_vec_znx(module, self.rank() + 1, self.size());
131        let mut ct_idft = GLWECiphertext {
132            data: vec_znx_tmp,
133            basek: self.basek,
134            k: self.k,
135        };
136        ct_idft.encrypt_zero_sk(
137            module, sk_dft, source_xa, source_xe, sigma, bound, scratch_1,
138        );
139
140        ct_idft.dft(module, self);
141    }
142
143    pub fn keyswitch<DataLhs, DataRhs>(
144        &mut self,
145        module: &Module<FFT64>,
146        lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
147        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
148        scratch: &mut Scratch,
149    ) where
150        VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
151        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
152    {
153        rhs.0.prod_with_glwe_fourier(module, self, lhs, scratch);
154    }
155
156    pub fn keyswitch_inplace<DataRhs>(
157        &mut self,
158        module: &Module<FFT64>,
159        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
160        scratch: &mut Scratch,
161    ) where
162        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
163    {
164        rhs.0.prod_with_glwe_fourier_inplace(module, self, scratch);
165    }
166
167    pub fn external_product<DataLhs, DataRhs>(
168        &mut self,
169        module: &Module<FFT64>,
170        lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
171        rhs: &GGSWCiphertext<DataRhs, FFT64>,
172        scratch: &mut Scratch,
173    ) where
174        VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
175        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
176    {
177        rhs.prod_with_glwe_fourier(module, self, lhs, scratch);
178    }
179
180    pub fn external_product_inplace<DataRhs>(
181        &mut self,
182        module: &Module<FFT64>,
183        rhs: &GGSWCiphertext<DataRhs, FFT64>,
184        scratch: &mut Scratch,
185    ) where
186        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
187    {
188        rhs.prod_with_glwe_fourier_inplace(module, self, scratch);
189    }
190}
191
192impl<DataSelf> GLWECiphertextFourier<DataSelf, FFT64>
193where
194    VecZnxDft<DataSelf, FFT64>: VecZnxDftToRef<FFT64>,
195{
196    pub fn decrypt<DataPt, DataSk>(
197        &self,
198        module: &Module<FFT64>,
199        pt: &mut GLWEPlaintext<DataPt>,
200        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
201        scratch: &mut Scratch,
202    ) where
203        VecZnx<DataPt>: VecZnxToMut + VecZnxToRef,
204        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
205    {
206        #[cfg(debug_assertions)]
207        {
208            assert_eq!(self.rank(), sk_dft.rank());
209            assert_eq!(self.n(), module.n());
210            assert_eq!(pt.n(), module.n());
211            assert_eq!(sk_dft.n(), module.n());
212        }
213
214        let cols = self.rank() + 1;
215
216        let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
217        pt_big.zero();
218
219        {
220            (1..cols).for_each(|i| {
221                let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
222                module.svp_apply(&mut ci_dft, 0, sk_dft, i - 1, self, i);
223                let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
224                module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
225            });
226        }
227
228        {
229            let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
230            // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
231            module.vec_znx_idft(&mut c0_big, 0, self, 0, scratch_2);
232            module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
233        }
234
235        // pt = norm(BIG(m + e))
236        module.vec_znx_big_normalize(self.basek(), pt, 0, &mut pt_big, 0, scratch_1);
237
238        pt.basek = self.basek();
239        pt.k = pt.k().min(self.k());
240    }
241
242    pub(crate) fn idft<DataRes>(&self, module: &Module<FFT64>, res: &mut GLWECiphertext<DataRes>, scratch: &mut Scratch)
243    where
244        GLWECiphertext<DataRes>: VecZnxToMut,
245    {
246        #[cfg(debug_assertions)]
247        {
248            assert_eq!(self.rank(), res.rank());
249            assert_eq!(self.basek(), res.basek())
250        }
251
252        let min_size: usize = self.size().min(res.size());
253
254        let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
255
256        (0..self.rank() + 1).for_each(|i| {
257            module.vec_znx_idft(&mut res_big, 0, self, i, scratch1);
258            module.vec_znx_big_normalize(self.basek(), res, i, &res_big, 0, scratch1);
259        });
260    }
261}