1use crate::{
 2    ciphertext::{Ciphertext, GadgetCiphertext},
 3    elem::Elem,
 4    keys::SwitchingKey,
 5};
 6use base2k::{Infos, Module, VecZnx, VecZnxBig, VecZnxDft, VecZnxOps, VmpPMatOps};
 7
 8pub fn gadget_product_tmp_bytes(
 9    module: &Module,
10    log_base2k: usize,
11    out_log_q: usize,
12    in_log_q: usize,
13    gct_rows: usize,
14    gct_log_q: usize,
15) -> usize {
16    let gct_cols: usize = (gct_log_q + log_base2k - 1) / log_base2k;
17    let in_limbs: usize = (in_log_q + log_base2k - 1) / log_base2k;
18    let out_limbs: usize = (out_log_q + log_base2k - 1) / log_base2k;
19    module.vmp_apply_dft_to_dft_tmp_bytes(out_limbs, in_limbs, gct_rows, gct_cols)
20        + 2 * module.bytes_of_vec_znx_dft(gct_cols)
21}
22
23pub fn gadget_product_inplace_thread_safe<const OVERWRITE: bool>(
24    module: &Module,
25    res: &mut Elem,
26    b: &GadgetCiphertext,
27    tmp_bytes: &mut [u8],
28) {
29    unsafe {
30        let a_ptr: *const VecZnx = res.at(1) as *const VecZnx;
31        gadget_product_thread_safe::<OVERWRITE>(module, res, &*a_ptr, b, tmp_bytes);
32    }
33}
34
35/// Evaluates the gadget product res <- a x b.
36///
37/// # Arguments
38///
39/// * `module`: backend support for operations mod (X^N + 1).
40/// * `res`: an [Elem] to store (-cs + m * a + e, c) with res_ncols limbs.
41/// * `a`: a [VecZnx] of a_ncols limbs.
42/// * `b`: a [GadgetCiphertext] as a vector of (-Bs + m * 2^{-k} + E, B)
43///       containing b_nrows [VecZnx], each of b_ncols limbs.
44///
45/// # Computation
46///
47/// res = sum[min(a_ncols, b_nrows)] decomp(a, i) * (-B[i]s + m * 2^{-k*i} + E[i], B[i])
48///     = (cs + m * a + e, c) with min(res_limbs, b_cols) limbs.
49pub fn gadget_product_thread_safe<const OVERWRITE: bool>(
50    module: &Module,
51    res: &mut Elem,
52    a: &VecZnx,
53    b: &GadgetCiphertext,
54    tmp_bytes: &mut [u8],
55) {
56    let log_base2k: usize = b.log_base2k();
57    let cols: usize = b.cols();
58
59    let (tmp_bytes_vmp_apply_dft, tmp_bytes) =
60        tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
61
62    let (tmp_bytes_c1_dft, tmp_bytes_res_dft) = tmp_bytes.split_at_mut(tmp_bytes.len() >> 1);
63
64    let mut c1_dft: VecZnxDft = module.new_vec_znx_from_bytes(cols, tmp_bytes_c1_dft);
65    let mut res_dft: VecZnxDft = module.new_vec_znx_from_bytes(cols, tmp_bytes_res_dft);
66    let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
67
68    // a_dft <- DFT(a) [cols]
69    module.vec_znx_dft(&mut c1_dft, a, a.limbs());
70
71    // >>>>>>>> RES[0]
72
73    // res_dft <- sum[rows] DFT(a)[cols] x GadgetCiphertext[0][cols]
74    module.vmp_apply_dft_to_dft(&mut res_dft, &c1_dft, &b.value[0], tmp_bytes_vmp_apply_dft);
75
76    // res_big <- IDFT(DFT(a) x GadgetCiphertext[0])
77    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
78
79    // res_big <- res[0] + a_dft x GadgetCiphertext[0]
80    module.vec_znx_big_add_small_inplace(&mut res_big, res.at(0));
81    module.vec_znx_big_normalize(log_base2k, res.at_mut(0), &res_big, tmp_bytes_vmp_apply_dft);
82
83    // >>>>>>>> RES[1]
84
85    // res_dft <- DFT(c1) x GadgetCiphertext[1]
86    module.vmp_apply_dft_to_dft(&mut res_dft, &c1_dft, &b.value[1], tmp_bytes_vmp_apply_dft);
87
88    // res_big <- IDFT(DFT(c1) x GadgetCiphertext[1])
89    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
90
91    if OVERWRITE {
92        // res[1] = normalize(a_dft x GadgetCiphertext[1])
93        module.vec_znx_big_normalize(log_base2k, res.at_mut(1), &res_big, tmp_bytes_vmp_apply_dft);
94    } else {
95        // res[1] = normalize(a_dft x GadgetCiphertext[1] + res[1])
96        module.vec_znx_big_add_small_inplace(&mut res_big, res.at(0));
97        module.vec_znx_big_normalize(log_base2k, res.at_mut(0), &res_big, tmp_bytes_vmp_apply_dft);
98    }
99}