1use crate::{
  2    ciphertext::Ciphertext,
  3    elem::{Elem, ElemVecZnx, VecZnxCommon},
  4};
  5use base2k::{Module, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VmpPMat, VmpPMatOps};
  6use std::cmp::min;
  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<const OVERWRITE: bool, T>(
 24    module: &Module,
 25    res: &mut Elem<T>,
 26    b: &Ciphertext<VmpPMat>,
 27    tmp_bytes: &mut [u8],
 28) where
 29    T: VecZnxCommon<Owned = T>,
 30    Elem<T>: ElemVecZnx<T>,
 31{
 32    unsafe {
 33        let a_ptr: *const T = res.at(1) as *const T;
 34        gadget_product::<OVERWRITE, T>(module, res, &*a_ptr, b, tmp_bytes);
 35    }
 36}
 37
 38/// Evaluates the gadget product res <- a x b.
 39///
 40/// # Arguments
 41///
 42/// * `module`: backend support for operations mod (X^N + 1).
 43/// * `res`: an [Elem] to store (-cs + m * a + e, c) with res_ncols limbs.
 44/// * `a`: a [VecZnx] of a_ncols limbs.
 45/// * `b`: a [GadgetCiphertext] as a vector of (-Bs + m * 2^{-k} + E, B)
 46///       containing b_nrows [VecZnx], each of b_ncols limbs.
 47///
 48/// # Computation
 49///
 50/// res = sum[min(a_ncols, b_nrows)] decomp(a, i) * (-B[i]s + m * 2^{-k*i} + E[i], B[i])
 51///     = (cs + m * a + e, c) with min(res_limbs, b_cols) limbs.
 52pub fn gadget_product<const OVERWRITE: bool, T>(
 53    module: &Module,
 54    res: &mut Elem<T>,
 55    a: &T,
 56    b: &Ciphertext<VmpPMat>,
 57    tmp_bytes: &mut [u8],
 58) where
 59    T: VecZnxCommon<Owned = T>,
 60    Elem<T>: ElemVecZnx<T>,
 61{
 62    let log_base2k: usize = b.log_base2k();
 63    let rows: usize = min(b.rows(), a.cols());
 64    let cols: usize = b.cols();
 65
 66    let bytes_vmp_apply_dft: usize =
 67        module.vmp_apply_dft_to_dft_tmp_bytes(cols, a.cols(), rows, cols);
 68    let bytes_vec_znx_dft: usize = module.bytes_of_vec_znx_dft(cols);
 69
 70    let (tmp_bytes_vmp_apply_dft, tmp_bytes) = tmp_bytes.split_at_mut(bytes_vmp_apply_dft);
 71    let (tmp_bytes_c1_dft, tmp_bytes_res_dft) = tmp_bytes.split_at_mut(bytes_vec_znx_dft);
 72
 73    let mut tmp_a_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_c1_dft);
 74    let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_res_dft);
 75    let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
 76
 77    // Alias c0 and c1 part of res_big
 78    let (tmp_bytes_res_dft_c0, tmp_bytes_res_dft_c1) =
 79        tmp_bytes_res_dft.split_at_mut(bytes_vec_znx_dft >> 1);
 80    let res_big_c0: VecZnxBig = module.new_vec_znx_big_from_bytes(cols >> 1, tmp_bytes_res_dft_c0);
 81    let mut res_big_c1: VecZnxBig =
 82        module.new_vec_znx_big_from_bytes(cols >> 1, tmp_bytes_res_dft_c1);
 83
 84    // tmp_a_dft <- DFT(a)
 85    // (n x cols) <- (n x limbs=rows) x (rows x cols)
 86    // res_dft[a * (G0|G1)] <- sum[rows] tmp_a_dft x (DFT(G0)|DFT(G1))
 87    gadget_product_core(module, &mut res_dft, a, b.at(0), tmp_bytes_vmp_apply_dft);
 88
 89    // res_big[a * (G0|G1)] <- IDFT(res_dft[a * (G0|G1)])
 90    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
 91
 92    // res_big <- res[0] + res_big[a*G0]
 93    module.vec_znx_big_add_small_inplace(&mut res_big, res.at(0));
 94    module.vec_znx_big_normalize(log_base2k, res.at_mut(0), &res_big_c0, tmp_bytes_c1_dft);
 95
 96    if OVERWRITE {
 97        // res[1] = normalize(res_big[a*G1])
 98        module.vec_znx_big_normalize(log_base2k, res.at_mut(1), &res_big_c1, tmp_bytes_c1_dft);
 99    } else {
100        // res[1] = normalize(res_big[a*G1] + res[1])
101        module.vec_znx_big_add_small_inplace(&mut res_big_c1, res.at(1));
102        module.vec_znx_big_normalize(log_base2k, res.at_mut(1), &res_big_c1, tmp_bytes_c1_dft);
103    }
104}
105
106pub fn gadget_product_core<T>(
107    module: &Module,
108    res_dft: &mut VecZnxDft,
109    a: &T,
110    b: &VmpPMat,
111    tmp_bytes_vmp_apply_dft: &mut [u8],
112) where
113    T: VecZnxCommon<Owned = T>,
114    Elem<T>: ElemVecZnx<T>,
115{
116    // res_dft <- DFT(a)
117    module.vec_znx_dft(res_dft, a, a.cols());
118
119    // (n x cols) <- (n x limbs=rows) x (rows x cols)
120    // res_dft[a * (G0|G1)] <- sum[rows] res_dft x (DFT(G0)|DFT(G1))
121    module.vmp_apply_dft_to_dft_inplace(res_dft, b, tmp_bytes_vmp_apply_dft);
122}
123
124pub fn rgsw_product<T>(
125    module: &Module,
126    res: &mut Elem<T>,
127    a: &Ciphertext<T>,
128    b: &Ciphertext<VmpPMat>,
129    tmp_bytes: &mut [u8],
130) where
131    T: VecZnxCommon<Owned = T>,
132    Elem<T>: ElemVecZnx<T>,
133{
134    let log_base2k: usize = b.log_base2k();
135    let rows: usize = min(b.rows(), a.cols());
136    let cols: usize = b.cols();
137    let in_cols = a.cols();
138    let out_cols: usize = a.cols();
139
140    let bytes_of_vec_znx_dft = module.bytes_of_vec_znx_dft(cols);
141    let bytes_of_vmp_apply_dft_to_dft =
142        module.vmp_apply_dft_to_dft_tmp_bytes(out_cols, in_cols, rows, cols);
143
144    let (tmp_bytes_c0_dft, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_vec_znx_dft);
145    let (tmp_bytes_c1_dft, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_vec_znx_dft);
146    let (tmp_bytes_tmp_dft, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_vec_znx_dft);
147    let (tmp_bytes_r1_dft, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_vec_znx_dft);
148    let (tmp_bytes_r2_dft, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_vec_znx_dft);
149    let (bytes_of_vmp_apply_dft_to_dft, tmp_bytes) =
150        tmp_bytes.split_at_mut(bytes_of_vmp_apply_dft_to_dft);
151
152    let mut c0_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_c0_dft);
153    let mut c1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_c1_dft);
154    let mut tmp_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_tmp_dft);
155    let mut r1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_r1_dft);
156    let mut r2_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes(cols, tmp_bytes_r2_dft);
157
158    // c0_dft <- DFT(a[0])
159    module.vec_znx_dft(&mut c0_dft, a.at(0), in_cols);
160
161    // r_dft <- sum[rows] c0_dft[cols] x RGSW[0][cols]
162    module.vmp_apply_dft_to_dft(
163        &mut r1_dft,
164        &c1_dft,
165        &b.0.value[0],
166        bytes_of_vmp_apply_dft_to_dft,
167    );
168
169    // c1_dft <- DFT(a[1])
170    module.vec_znx_dft(&mut c1_dft, a.at(1), in_cols);
171}