1use crate::ciphertext::{Ciphertext, GadgetCiphertext};
 2use base2k::{Module, VecZnxBig, VecZnxDft, VmpPMatOps};
 3
 4pub fn gadget_product_tmp_bytes(
 5    module: &Module,
 6    log_base2k: usize,
 7    out_log_q: usize,
 8    in_log_q: usize,
 9    gct_rows: usize,
10    gct_log_q: usize,
11) -> usize {
12    let gct_cols: usize = (gct_log_q + log_base2k - 1) / log_base2k;
13    let in_limbs: usize = (in_log_q + log_base2k - 1) / log_base2k;
14    let out_limbs: usize = (out_log_q + log_base2k - 1) / log_base2k;
15    module.vmp_apply_dft_to_dft_tmp_bytes(out_limbs, in_limbs, gct_rows, gct_cols)
16        + 2 * module.bytes_of_vec_znx_dft(gct_cols)
17}
18
19pub fn gadget_product_inplace(
20    module: &Module,
21    a: &mut Ciphertext,
22    b: &GadgetCiphertext,
23    tmp_bytes: &mut [u8],
24) {
25    // This is safe to do because the relevant values of a are copied to a buffer before being
26    // overwritten.
27    unsafe {
28        let a_ptr: *mut Ciphertext = a;
29        gadget_product(module, a, &*a_ptr, b, tmp_bytes)
30    }
31}
32
33pub fn gadget_product(
34    module: &Module,
35    res: &mut Ciphertext,
36    a: &Ciphertext,
37    b: &GadgetCiphertext,
38    tmp_bytes: &mut [u8],
39) {
40    assert!(
41        a.log_base2k() == b.log_base2k(),
42        "invalid inputs: a.log_base2k={} != b.log_base2k={}",
43        a.log_base2k(),
44        b.log_base2k()
45    );
46
47    let log_base2k: usize = b.log_base2k();
48    let cols: usize = b.cols();
49
50    let (tmp_bytes_vmp_apply_dft, tmp_bytes) =
51        tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
52
53    let (tmp_bytes_c1_dft, tmp_bytes_res_dft) = tmp_bytes.split_at_mut(tmp_bytes.len() >> 1);
54
55    let mut c1_dft: VecZnxDft = module.new_vec_znx_from_bytes(cols, tmp_bytes_c1_dft);
56    let mut res_dft: VecZnxDft = module.new_vec_znx_from_bytes(cols, tmp_bytes_res_dft);
57    let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
58
59    // c1_dft <- DFT(b[1])
60    module.vec_znx_dft(&mut c1_dft, a.at(1), a.limbs());
61
62    // res_dft <- DFT(c1) x GadgetCiphertext[0]
63    module.vmp_apply_dft_to_dft(&mut res_dft, &c1_dft, &b.value[0], tmp_bytes_vmp_apply_dft);
64
65    // res_big <- IDFT(DFT(c1) x GadgetCiphertext[0])
66    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
67
68    // res_big <- c0 + c1_dft x GadgetCiphertext[0]
69    module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0), cols);
70
71    // res[0] = normalize(c0 + c1_dft x GadgetCiphertext[0])
72    module.vec_znx_big_normalize(log_base2k, res.at_mut(0), &res_big, tmp_bytes_vmp_apply_dft);
73
74    // res_dft <- DFT(c1) x GadgetCiphertext[1]
75    module.vmp_apply_dft_to_dft(&mut res_dft, &c1_dft, &b.value[1], tmp_bytes_vmp_apply_dft);
76
77    // res_big <- IDFT(DFT(c1) x GadgetCiphertext[1])
78    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
79
80    // res[1] = normalize(c1_dft x GadgetCiphertext[1])
81    module.vec_znx_big_normalize(log_base2k, res.at_mut(1), &res_big, tmp_bytes_vmp_apply_dft);
82}