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