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