1use crate::{
 2    ciphertext::Ciphertext,
 3    elem::{Elem, ElemBasics},
 4    keys::SecretKey,
 5    parameters::Parameters,
 6    plaintext::Plaintext,
 7};
 8use base2k::{
 9    Infos, VecZnx, Module, SvpPPol, SvpPPolOps, 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(
38        &self,
39        res: &mut Plaintext<VecZnx>,
40        ct: &Ciphertext,
41        sk: &SvpPPol,
42        tmp_bytes: &mut [u8],
43    ) {
44        decrypt_rlwe_thread_safe(self.module(), &mut res.0, &ct.0, sk, tmp_bytes)
45    }
46}
47
48pub fn decrypt_rlwe_thread_safe<T>(
49    module: &Module,
50    res: &mut Elem<T>,
51    a: &Elem<T>,
52    sk: &SvpPPol,
53    tmp_bytes: &mut [u8],
54) where
55    T: VecZnxApi + Infos,
56{
57    assert!(
58        tmp_bytes.len() >= decrypt_rlwe_thread_safe_tmp_byte(module, a.limbs()),
59        "invalid tmp_bytes: tmp_bytes.len()={} < decrypt_rlwe_thread_safe_tmp_byte={}",
60        tmp_bytes.len(),
61        decrypt_rlwe_thread_safe_tmp_byte(module, a.limbs())
62    );
63
64    let res_dft_bytes: usize = module.bytes_of_vec_znx_dft(a.limbs());
65
66    let mut res_dft: VecZnxDft = VecZnxDft::from_bytes(a.limbs(), tmp_bytes);
67    let mut res_big: base2k::VecZnxBig = res_dft.as_vec_znx_big();
68
69    // res_dft <- DFT(ct[1]) * DFT(sk)
70    module.svp_apply_dft(&mut res_dft, sk, &a.value[1], a.limbs());
71    // res_big <- ct[1] x sk
72    module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, a.limbs());
73    // res_big <- ct[1] x sk + ct[0]
74    module.vec_znx_big_add_small_inplace(&mut res_big, &a.value[0]);
75    // res <- normalize(ct[1] x sk + ct[0])
76    module.vec_znx_big_normalize(
77        a.log_base2k(),
78        res.at_mut(0),
79        &res_big,
80        &mut tmp_bytes[res_dft_bytes..],
81    );
82
83    res.log_base2k = a.log_base2k();
84    res.log_q = min(res.log_q(), a.log_q());
85    res.log_scale = a.log_scale();
86}