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