1use crate::{ciphertext::Ciphertext, elem::ElemCommon, parameters::Parameters};
  2use base2k::{
  3    Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VmpPMat, VmpPMatOps,
  4    assert_alignement,
  5};
  6use std::cmp::min;
  7
  8impl Parameters {
  9    pub fn rgsw_product_tmp_bytes(
 10        &self,
 11        res_logq: usize,
 12        in_logq: usize,
 13        gct_logq: usize,
 14    ) -> usize {
 15        rgsw_product_tmp_bytes(
 16            self.module(),
 17            self.log_base2k(),
 18            res_logq,
 19            in_logq,
 20            gct_logq,
 21        )
 22    }
 23}
 24pub fn rgsw_product_tmp_bytes(
 25    module: &Module,
 26    log_base2k: usize,
 27    res_logq: usize,
 28    in_logq: usize,
 29    gct_logq: usize,
 30) -> usize {
 31    let gct_cols: usize = (gct_logq + log_base2k - 1) / log_base2k;
 32    let in_cols: usize = (in_logq + log_base2k - 1) / log_base2k;
 33    let res_cols: usize = (res_logq + log_base2k - 1) / log_base2k;
 34    return module.vmp_apply_dft_to_dft_tmp_bytes(res_cols, in_cols, in_cols, gct_cols)
 35        + module.bytes_of_vec_znx_dft(std::cmp::min(res_cols, in_cols))
 36        + 2 * module.bytes_of_vec_znx_dft(gct_cols);
 37}
 38
 39pub fn rgsw_product(
 40    module: &Module,
 41    c: &mut Ciphertext<VecZnx>,
 42    a: &Ciphertext<VecZnx>,
 43    b: &Ciphertext<VmpPMat>,
 44    b_cols: usize,
 45    tmp_bytes: &mut [u8],
 46) {
 47    #[cfg(debug_assertions)]
 48    {
 49        assert!(b_cols <= b.cols());
 50        assert_eq!(c.size(), 2);
 51        assert_eq!(a.size(), 2);
 52        assert_eq!(b.size(), 4);
 53        assert!(
 54            tmp_bytes.len()
 55                >= rgsw_product_tmp_bytes(
 56                    module,
 57                    c.cols(),
 58                    a.cols(),
 59                    min(b.rows(), a.cols()),
 60                    b_cols
 61                )
 62        );
 63        assert_alignement(tmp_bytes.as_ptr());
 64    }
 65
 66    let (tmp_bytes_ai_dft, tmp_bytes) =
 67        tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(a.cols()));
 68    let (tmp_bytes_c0_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
 69    let (tmp_bytes_c1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
 70
 71    let mut ai_dft: VecZnxDft =
 72        module.new_vec_znx_dft_from_bytes_borrow(a.cols(), tmp_bytes_ai_dft);
 73    let mut c0_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_c0_dft);
 74    let mut c1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_c1_dft);
 75
 76    let mut c0_big: VecZnxBig = c0_dft.as_vec_znx_big();
 77    let mut c1_big: VecZnxBig = c1_dft.as_vec_znx_big();
 78
 79    module.vec_znx_dft(&mut ai_dft, a.at(0));
 80    module.vmp_apply_dft_to_dft(&mut c0_dft, &ai_dft, b.at(0), tmp_bytes);
 81    module.vmp_apply_dft_to_dft(&mut c1_dft, &ai_dft, b.at(1), tmp_bytes);
 82
 83    module.vec_znx_dft(&mut ai_dft, a.at(1));
 84    module.vmp_apply_dft_to_dft_add(&mut c0_dft, &ai_dft, b.at(2), tmp_bytes);
 85    module.vmp_apply_dft_to_dft_add(&mut c1_dft, &ai_dft, b.at(3), tmp_bytes);
 86
 87    module.vec_znx_idft_tmp_a(&mut c0_big, &mut c0_dft);
 88    module.vec_znx_idft_tmp_a(&mut c1_big, &mut c1_dft);
 89
 90    module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(0), &mut c0_big, tmp_bytes);
 91    module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(1), &mut c1_big, tmp_bytes);
 92}
 93
 94pub fn rgsw_product_inplace(
 95    module: &Module,
 96    a: &mut Ciphertext<VecZnx>,
 97    b: &Ciphertext<VmpPMat>,
 98    b_cols: usize,
 99    tmp_bytes: &mut [u8],
100) {
101    #[cfg(debug_assertions)]
102    {
103        assert!(b_cols <= b.cols());
104        assert_eq!(a.size(), 2);
105        assert_eq!(b.size(), 4);
106        assert!(
107            tmp_bytes.len()
108                >= rgsw_product_tmp_bytes(
109                    module,
110                    a.cols(),
111                    a.cols(),
112                    min(b.rows(), a.cols()),
113                    b_cols
114                )
115        );
116        assert_alignement(tmp_bytes.as_ptr());
117    }
118
119    let (tmp_bytes_ai_dft, tmp_bytes) =
120        tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(a.cols()));
121    let (tmp_bytes_c0_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
122    let (tmp_bytes_c1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
123
124    let mut ai_dft: VecZnxDft =
125        module.new_vec_znx_dft_from_bytes_borrow(a.cols(), tmp_bytes_ai_dft);
126    let mut c0_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_c0_dft);
127    let mut c1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_c1_dft);
128
129    let mut c0_big: VecZnxBig = c0_dft.as_vec_znx_big();
130    let mut c1_big: VecZnxBig = c1_dft.as_vec_znx_big();
131
132    module.vec_znx_dft(&mut ai_dft, a.at(0));
133    module.vmp_apply_dft_to_dft(&mut c0_dft, &ai_dft, b.at(0), tmp_bytes);
134    module.vmp_apply_dft_to_dft(&mut c1_dft, &ai_dft, b.at(1), tmp_bytes);
135
136    module.vec_znx_dft(&mut ai_dft, a.at(1));
137    module.vmp_apply_dft_to_dft_add(&mut c0_dft, &ai_dft, b.at(2), tmp_bytes);
138    module.vmp_apply_dft_to_dft_add(&mut c1_dft, &ai_dft, b.at(3), tmp_bytes);
139
140    module.vec_znx_idft_tmp_a(&mut c0_big, &mut c0_dft);
141    module.vec_znx_idft_tmp_a(&mut c1_big, &mut c1_dft);
142
143    module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(0), &mut c0_big, tmp_bytes);
144    module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(1), &mut c1_big, tmp_bytes);
145}
146
147#[cfg(test)]
148mod test {
149    use crate::{
150        ciphertext::{Ciphertext, new_rgsw_ciphertext},
151        decryptor::decrypt_rlwe,
152        elem::ElemCommon,
153        encryptor::{encrypt_rgsw_sk, encrypt_rlwe_sk},
154        keys::SecretKey,
155        parameters::{DEFAULT_SIGMA, Parameters, ParametersLiteral},
156        plaintext::Plaintext,
157        rgsw_product::rgsw_product_inplace,
158    };
159    use base2k::{
160        BACKEND, Encoding, Module, Scalar, SvpPPol, SvpPPolOps, VecZnx, VecZnxOps, VmpPMat,
161        alloc_aligned,
162    };
163    use sampling::source::{Source, new_seed};
164
165    #[test]
166    fn test_rgsw_product() {
167        let log_base2k: usize = 10;
168        let log_q: usize = 50;
169        let log_p: usize = 15;
170
171        // Basic parameters with enough limbs to test edge cases
172        let params_lit: ParametersLiteral = ParametersLiteral {
173            backend: BACKEND::FFT64,
174            log_n: 12,
175            log_q: log_q,
176            log_p: log_p,
177            log_base2k: log_base2k,
178            log_scale: 20,
179            xe: 3.2,
180            xs: 1 << 11,
181        };
182
183        let params: Parameters = Parameters::new(&params_lit);
184
185        let module: &Module = params.module();
186        let log_q: usize = params.log_q();
187        let log_qp: usize = params.log_qp();
188        let gct_rows: usize = params.cols_q();
189        let gct_cols: usize = params.cols_qp();
190
191        // scratch space
192        let mut tmp_bytes: Vec<u8> = alloc_aligned(
193            params.decrypt_rlwe_tmp_byte(log_q)
194                | params.encrypt_rlwe_sk_tmp_bytes(log_q)
195                | params.rgsw_product_tmp_bytes(log_q, log_q, log_qp)
196                | params.encrypt_rgsw_sk_tmp_bytes(gct_rows, log_qp),
197        );
198
199        // Samplers for public and private randomness
200        let mut source_xe: Source = Source::new(new_seed());
201        let mut source_xa: Source = Source::new(new_seed());
202        let mut source_xs: Source = Source::new(new_seed());
203
204        let mut sk: SecretKey = SecretKey::new(module);
205        sk.fill_ternary_hw(params.xs(), &mut source_xs);
206        let mut sk_svp_ppol: SvpPPol = module.new_svp_ppol();
207        module.svp_prepare(&mut sk_svp_ppol, &sk.0);
208
209        let mut ct_rgsw: Ciphertext<VmpPMat> =
210            new_rgsw_ciphertext(module, log_base2k, gct_rows, log_qp);
211
212        let k: i64 = 3;
213
214        // X^k
215        let m: Scalar = module.new_scalar();
216        let data: &mut [i64] = m.raw_mut();
217        data[k as usize] = 1;
218
219        encrypt_rgsw_sk(
220            module,
221            &mut ct_rgsw,
222            &m,
223            &sk_svp_ppol,
224            &mut source_xa,
225            &mut source_xe,
226            DEFAULT_SIGMA,
227            &mut tmp_bytes,
228        );
229
230        let log_k: usize = 2 * log_base2k;
231
232        let mut ct: Ciphertext<VecZnx> = params.new_ciphertext(log_q);
233        let mut pt: Plaintext = params.new_plaintext(log_q);
234        let mut pt_rotate: Plaintext = params.new_plaintext(log_q);
235
236        pt.at_mut(0).encode_vec_i64(log_base2k, log_k, &data, 32);
237
238        module.vec_znx_rotate(k, pt_rotate.at_mut(0), pt.at_mut(0));
239
240        encrypt_rlwe_sk(
241            module,
242            &mut ct.elem_mut(),
243            Some(pt.at(0)),
244            &sk_svp_ppol,
245            &mut source_xa,
246            &mut source_xe,
247            params.xe(),
248            &mut tmp_bytes,
249        );
250
251        rgsw_product_inplace(module, &mut ct, &ct_rgsw, gct_cols, &mut tmp_bytes);
252
253        decrypt_rlwe(
254            module,
255            pt.elem_mut(),
256            ct.elem(),
257            &sk_svp_ppol,
258            &mut tmp_bytes,
259        );
260
261        module.vec_znx_sub_ba_inplace(pt.at_mut(0), pt_rotate.at(0));
262
263        //pt.at(0).print(pt.cols(), 16);
264
265        let noise_have: f64 = pt.at(0).std(log_base2k).log2();
266
267        let var_msg: f64 = 1f64 / params.n() as f64; // X^{k}
268        let var_a0_err: f64 = params.xe() * params.xe();
269        let var_a1_err: f64 = 1f64 / 12f64;
270
271        let noise_pred: f64 =
272            params.noise_rgsw_product(var_msg, var_a0_err, var_a1_err, ct.log_q(), ct_rgsw.log_q());
273
274        println!("noise_pred: {}", noise_pred);
275        println!("noise_have: {}", noise_have);
276
277        assert!(noise_have <= noise_pred + 1.0);
278    }
279}
280
281impl Parameters {
282    pub fn noise_rgsw_product(
283        &self,
284        var_msg: f64,
285        var_a0_err: f64,
286        var_a1_err: f64,
287        a_logq: usize,
288        b_logq: usize,
289    ) -> f64 {
290        let n: f64 = self.n() as f64;
291        let var_xs: f64 = self.xs() as f64;
292
293        let var_gct_err_lhs: f64;
294        let var_gct_err_rhs: f64;
295        if b_logq < self.log_qp() {
296            let var_round: f64 = 1f64 / 12f64;
297            var_gct_err_lhs = var_round;
298            var_gct_err_rhs = var_round;
299        } else {
300            var_gct_err_lhs = self.xe() * self.xe();
301            var_gct_err_rhs = 0f64;
302        }
303
304        noise_rgsw_product(
305            n,
306            self.log_base2k(),
307            var_xs,
308            var_msg,
309            var_a0_err,
310            var_a1_err,
311            var_gct_err_lhs,
312            var_gct_err_rhs,
313            a_logq,
314            b_logq,
315        )
316    }
317}
318
319pub fn noise_rgsw_product(
320    n: f64,
321    log_base2k: usize,
322    var_xs: f64,
323    var_msg: f64,
324    var_a0_err: f64,
325    var_a1_err: f64,
326    var_gct_err_lhs: f64,
327    var_gct_err_rhs: f64,
328    a_logq: usize,
329    b_logq: usize,
330) -> f64 {
331    let a_logq: usize = min(a_logq, b_logq);
332    let a_cols: usize = (a_logq + log_base2k - 1) / log_base2k;
333
334    let b_scale = 2.0f64.powi(b_logq as i32);
335    let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
336
337    let base: f64 = (1 << (log_base2k)) as f64;
338    let var_base: f64 = base * base / 12f64;
339
340    // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
341    // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
342    let mut noise: f64 =
343        2.0 * (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
344    noise += var_msg * var_a0_err * a_scale * a_scale * n;
345    noise += var_msg * var_a1_err * a_scale * a_scale * n * var_xs;
346    noise = noise.sqrt();
347    noise /= b_scale;
348    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
349}