1use base2k::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
  3    ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxOps,
  4    ZnxZero,
  5};
  6use sampling::source::Source;
  7
  8use crate::{
  9    elem::Infos,
 10    elem_rlwe::{RLWECt, RLWECtDft, RLWEPt},
 11    keys::SecretKeyDft,
 12    utils::derive_size,
 13};
 14
 15pub struct GRLWECt<C, B: Backend> {
 16    pub data: MatZnxDft<C, B>,
 17    pub log_base2k: usize,
 18    pub log_k: usize,
 19}
 20
 21impl<B: Backend> GRLWECt<Vec<u8>, B> {
 22    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
 23        Self {
 24            data: module.new_mat_znx_dft(rows, 1, 2, derive_size(log_base2k, log_k)),
 25            log_base2k: log_base2k,
 26            log_k: log_k,
 27        }
 28    }
 29}
 30
 31impl<C> GRLWECt<C, FFT64>
 32where
 33    MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
 34{
 35    pub fn get_row(&self, module: &Module<FFT64>, i: usize, res: &mut RLWECtDft<C, FFT64>)
 36    where
 37        VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64>,
 38    {
 39        module.vmp_extract_row(res, self, i, 0);
 40    }
 41}
 42
 43impl<T, B: Backend> Infos for GRLWECt<T, B> {
 44    type Inner = MatZnxDft<T, B>;
 45
 46    fn inner(&self) -> &Self::Inner {
 47        &self.data
 48    }
 49
 50    fn log_base2k(&self) -> usize {
 51        self.log_base2k
 52    }
 53
 54    fn log_k(&self) -> usize {
 55        self.log_k
 56    }
 57}
 58
 59impl<C, B: Backend> MatZnxDftToMut<B> for GRLWECt<C, B>
 60where
 61    MatZnxDft<C, B>: MatZnxDftToMut<B>,
 62{
 63    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
 64        self.data.to_mut()
 65    }
 66}
 67
 68impl<C, B: Backend> MatZnxDftToRef<B> for GRLWECt<C, B>
 69where
 70    MatZnxDft<C, B>: MatZnxDftToRef<B>,
 71{
 72    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
 73        self.data.to_ref()
 74    }
 75}
 76
 77impl GRLWECt<Vec<u8>, FFT64> {
 78    pub fn encrypt_sk_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
 79        RLWECt::encrypt_sk_scratch_bytes(module, size)
 80            + module.bytes_of_vec_znx(2, size)
 81            + module.bytes_of_vec_znx(1, size)
 82            + module.bytes_of_vec_znx_dft(2, size)
 83    }
 84
 85    // pub fn encrypt_pk_scratch_bytes(module: &Module<FFT64>, pk_size: usize) -> usize {
 86    // RLWECt::encrypt_pk_scratch_bytes(module, pk_size)
 87    // }
 88}
 89
 90pub fn encrypt_grlwe_sk<C, P, S>(
 91    module: &Module<FFT64>,
 92    ct: &mut GRLWECt<C, FFT64>,
 93    pt: &ScalarZnx<P>,
 94    sk_dft: &SecretKeyDft<S, FFT64>,
 95    source_xa: &mut Source,
 96    source_xe: &mut Source,
 97    sigma: f64,
 98    bound: f64,
 99    scratch: &mut Scratch,
100) where
101    MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
102    ScalarZnx<P>: ScalarZnxToRef,
103    ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
104{
105    let rows: usize = ct.rows();
106    let size: usize = ct.size();
107    let log_base2k: usize = ct.log_base2k();
108
109    let (tmp_znx_pt, scrach_1) = scratch.tmp_vec_znx(module, 1, size);
110    let (tmp_znx_ct, scrach_2) = scrach_1.tmp_vec_znx(module, 2, size);
111    let (mut vec_znx_dft_ct, scratch_3) = scrach_2.tmp_vec_znx_dft(module, 2, size);
112
113    let mut vec_znx_pt: RLWEPt<&mut [u8]> = RLWEPt {
114        data: tmp_znx_pt,
115        log_base2k: log_base2k,
116        log_k: ct.log_k(),
117    };
118
119    let mut vec_znx_ct: RLWECt<&mut [u8]> = RLWECt {
120        data: tmp_znx_ct,
121        log_base2k: log_base2k,
122        log_k: ct.log_k(),
123    };
124
125    (0..rows).for_each(|row_i| {
126        // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
127        module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_i, pt, 0);
128        module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scratch_3);
129
130        // rlwe encrypt of vec_znx_pt into vec_znx_ct
131        vec_znx_ct.encrypt_sk(
132            module,
133            Some(&vec_znx_pt),
134            sk_dft,
135            source_xa,
136            source_xe,
137            sigma,
138            bound,
139            scratch_3,
140        );
141
142        vec_znx_pt.data.zero(); // zeroes for next iteration
143
144        // Switch vec_znx_ct into DFT domain
145        module.vec_znx_dft(&mut vec_znx_dft_ct, 0, &vec_znx_ct, 0);
146        module.vec_znx_dft(&mut vec_znx_dft_ct, 1, &vec_znx_ct, 1);
147
148        // Stores vec_znx_dft_ct into thw i-th row of the MatZnxDft
149        module.vmp_prepare_row(ct, row_i, 0, &vec_znx_dft_ct);
150    });
151}
152
153impl<C> GRLWECt<C, FFT64> {
154    pub fn encrypt_sk<P, S>(
155        &mut self,
156        module: &Module<FFT64>,
157        pt: &ScalarZnx<P>,
158        sk_dft: &SecretKeyDft<S, FFT64>,
159        source_xa: &mut Source,
160        source_xe: &mut Source,
161        sigma: f64,
162        bound: f64,
163        scratch: &mut Scratch,
164    ) where
165        MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
166        ScalarZnx<P>: ScalarZnxToRef,
167        ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
168    {
169        encrypt_grlwe_sk(
170            module, self, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
171        )
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use base2k::{FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps};
178    use sampling::source::Source;
179
180    use crate::{
181        elem::Infos,
182        elem_rlwe::{RLWECtDft, RLWEPt},
183        keys::{SecretKey, SecretKeyDft},
184    };
185
186    use super::GRLWECt;
187
188    #[test]
189    fn encrypt_sk_fft64() {
190        let module: Module<FFT64> = Module::<FFT64>::new(2048);
191        let log_base2k: usize = 8;
192        let log_k_ct: usize = 54;
193        let rows: usize = 4;
194
195        let sigma: f64 = 3.2;
196        let bound: f64 = sigma * 6.0;
197
198        let mut ct: GRLWECt<Vec<u8>, FFT64> = GRLWECt::new(&module, log_base2k, log_k_ct, rows);
199        let mut pt: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
200        let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
201
202        let mut source_xs: Source = Source::new([0u8; 32]);
203        let mut source_xe: Source = Source::new([0u8; 32]);
204        let mut source_xa: Source = Source::new([0u8; 32]);
205
206        pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
207
208        let mut scratch: ScratchOwned = ScratchOwned::new(
209            GRLWECt::encrypt_sk_scratch_bytes(&module, ct.size()) | RLWECtDft::decrypt_scratch_bytes(&module, ct.size()),
210        );
211
212        let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
213        sk.fill_ternary_prob(0.5, &mut source_xs);
214
215        let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
216        sk_dft.dft(&module, &sk);
217
218        ct.encrypt_sk(
219            &module,
220            &pt_scalar,
221            &sk_dft,
222            &mut source_xa,
223            &mut source_xe,
224            sigma,
225            bound,
226            scratch.borrow(),
227        );
228
229        let mut ct_rlwe_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_ct, 2);
230
231        (0..ct.rows()).for_each(|row_i| {
232            ct.get_row(&module, row_i, &mut ct_rlwe_dft);
233            ct_rlwe_dft.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
234            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &pt_scalar, 0);
235            let std_pt: f64 = pt.data.std(0, log_base2k) * (log_k_ct as f64).exp2();
236            assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
237        });
238
239        module.free();
240    }
241}