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