1use crate::ciphertext::Ciphertext;
 2use crate::elem::ElemCommon;
 3use base2k::{Module, VecZnx, VecZnxBigOps, VecZnxDftOps, VmpPMat, VmpPMatOps, assert_alignement};
 4use std::cmp::min;
 5
 6pub fn key_switch_tmp_bytes(module: &Module, log_base2k: usize, res_logq: usize, in_logq: usize, gct_logq: usize) -> usize {
 7    let gct_cols: usize = (gct_logq + log_base2k - 1) / log_base2k;
 8    let in_cols: usize = (in_logq + log_base2k - 1) / log_base2k;
 9    let res_cols: usize = (res_logq + log_base2k - 1) / log_base2k;
10    return module.vmp_apply_dft_to_dft_tmp_bytes(res_cols, in_cols, in_cols, gct_cols)
11        + module.bytes_of_vec_znx_dft(std::cmp::min(res_cols, in_cols))
12        + module.bytes_of_vec_znx_dft(gct_cols);
13}
14
15pub fn key_switch_rlwe(
16    module: &Module,
17    c: &mut Ciphertext<VecZnx>,
18    a: &Ciphertext<VecZnx>,
19    b: &Ciphertext<VmpPMat>,
20    b_cols: usize,
21    tmp_bytes: &mut [u8],
22) {
23    key_switch_rlwe_core(module, c, a, b, b_cols, tmp_bytes);
24}
25
26pub fn key_switch_rlwe_inplace(
27    module: &Module,
28    a: &mut Ciphertext<VecZnx>,
29    b: &Ciphertext<VmpPMat>,
30    b_cols: usize,
31    tmp_bytes: &mut [u8],
32) {
33    key_switch_rlwe_core(module, a, a, b, b_cols, tmp_bytes);
34}
35
36fn key_switch_rlwe_core(
37    module: &Module,
38    c: *mut Ciphertext<VecZnx>,
39    a: *const Ciphertext<VecZnx>,
40    b: &Ciphertext<VmpPMat>,
41    b_cols: usize,
42    tmp_bytes: &mut [u8],
43) {
44    // SAFETY WARNING: must ensure `c` and `a` are valid for read/write
45    let c: &mut Ciphertext<VecZnx> = unsafe { &mut *c };
46    let a: &Ciphertext<VecZnx> = unsafe { &*a };
47
48    let cols: usize = min(min(c.cols(), a.cols()), b.rows());
49
50    #[cfg(debug_assertions)]
51    {
52        assert!(b_cols <= b.cols());
53        assert!(tmp_bytes.len() >= key_switch_tmp_bytes(module, c.cols(), a.cols(), b.rows(), b.cols()));
54        assert_alignement(tmp_bytes.as_ptr());
55    }
56
57    let (tmp_bytes_a1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
58    let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
59
60    let mut a1_dft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_a1_dft);
61    let mut res_dft = module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_res_dft);
62    let mut res_big = res_dft.as_vec_znx_big();
63
64    module.vec_znx_dft(&mut a1_dft, a.at(1));
65    module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.at(0), tmp_bytes);
66    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
67
68    module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
69    module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(0), &mut res_big, tmp_bytes);
70
71    module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.at(1), tmp_bytes);
72    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
73
74    module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(1), &mut res_big, tmp_bytes);
75}
76
77pub fn key_switch_grlwe(module: &Module, c: &mut Ciphertext<VecZnx>, a: &Ciphertext<VecZnx>, b: &Ciphertext<VmpPMat>) {}
78
79pub fn key_switch_rgsw(module: &Module, c: &mut Ciphertext<VecZnx>, a: &Ciphertext<VecZnx>, b: &Ciphertext<VmpPMat>) {}