1use crate::{
 2    ciphertext::Ciphertext,
 3    elem::{Elem, ElemVecZnx, VecZnxCommon},
 4    keys::SecretKey,
 5    parameters::Parameters,
 6    plaintext::Plaintext,
 7};
 8use base2k::{
 9    Infos, Module, SvpPPol, SvpPPolOps, VecZnx, VecZnxApi, VecZnxBigOps, VecZnxDft, VecZnxDftOps,
10};
11use std::cmp::min;
12
13pub struct Decryptor {
14    sk: SvpPPol,
15}
16
17impl Decryptor {
18    pub fn new(params: &Parameters, sk: &SecretKey) -> Self {
19        let mut sk_svp_ppol: SvpPPol = params.module().new_svp_ppol();
20        sk.prepare(params.module(), &mut sk_svp_ppol);
21        Self { sk: sk_svp_ppol }
22    }
23}
24
25pub fn decrypt_rlwe_thread_safe_tmp_byte(module: &Module, limbs: usize) -> usize {
26    module.bytes_of_vec_znx_dft(limbs) + module.vec_znx_big_normalize_tmp_bytes()
27}
28
29impl Parameters {
30    pub fn decrypt_rlwe_thread_safe_tmp_byte(&self, log_q: usize) -> usize {
31        decrypt_rlwe_thread_safe_tmp_byte(
32            self.module(),
33            (log_q + self.log_base2k() - 1) / self.log_base2k(),
34        )
35    }
36
37    pub fn decrypt_rlwe_thread_safe<T>(
38        &self,
39        res: &mut Plaintext<T>,
40        ct: &Ciphertext<T>,
41        sk: &SvpPPol,
42        tmp_bytes: &mut [u8],
43    ) where
44        T: VecZnxCommon,
45        Elem<T>: Infos + ElemVecZnx<T>,
46    {
47        decrypt_rlwe_thread_safe(self.module(), &mut res.0, &ct.0, sk, tmp_bytes)
48    }
49}
50
51pub fn decrypt_rlwe_thread_safe<T>(
52    module: &Module,
53    res: &mut Elem<T>,
54    a: &Elem<T>,
55    sk: &SvpPPol,
56    tmp_bytes: &mut [u8],
57) where
58    T: VecZnxCommon,
59    Elem<T>: Infos + ElemVecZnx<T>,
60{
61    let cols: usize = a.cols();
62
63    assert!(
64        tmp_bytes.len() >= decrypt_rlwe_thread_safe_tmp_byte(module, cols),
65        "invalid tmp_bytes: tmp_bytes.len()={} < decrypt_rlwe_thread_safe_tmp_byte={}",
66        tmp_bytes.len(),
67        decrypt_rlwe_thread_safe_tmp_byte(module, cols)
68    );
69
70    let res_dft_bytes: usize = module.bytes_of_vec_znx_dft(cols);
71
72    let mut res_dft: VecZnxDft = VecZnxDft::from_bytes(a.cols(), tmp_bytes);
73    let mut res_big: base2k::VecZnxBig = res_dft.as_vec_znx_big();
74
75    // res_dft <- DFT(ct[1]) * DFT(sk)
76    module.svp_apply_dft(&mut res_dft, sk, a.at(1), cols);
77    // res_big <- ct[1] x sk
78    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
79    // res_big <- ct[1] x sk + ct[0]
80    module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
81    // res <- normalize(ct[1] x sk + ct[0])
82    module.vec_znx_big_normalize(
83        a.log_base2k(),
84        res.at_mut(0),
85        &res_big,
86        &mut tmp_bytes[res_dft_bytes..],
87    );
88
89    res.log_base2k = a.log_base2k();
90    res.log_q = min(res.log_q(), a.log_q());
91    res.log_scale = a.log_scale();
92}