1use crate::ciphertext::{Ciphertext, new_gadget_ciphertext};
 2use crate::elem::{Elem, ElemCommon};
 3use crate::encryptor::{encrypt_rlwe_sk, encrypt_rlwe_sk_tmp_bytes};
 4use base2k::{Module, Scalar, ScalarZnxDft, ScalarZnxDftOps, VecZnx, MatZnxDft};
 5use sampling::source::Source;
 6
 7pub struct SecretKey(pub Scalar);
 8
 9impl SecretKey {
10    pub fn new(module: &Module) -> Self {
11        SecretKey(Scalar::new(module.n()))
12    }
13
14    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
15        self.0.fill_ternary_prob(prob, source);
16    }
17
18    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
19        self.0.fill_ternary_hw(hw, source);
20    }
21
22    pub fn prepare(&self, module: &Module, sk_ppol: &mut ScalarZnxDft) {
23        module.svp_prepare(sk_ppol, &self.0)
24    }
25}
26
27pub struct PublicKey(pub Elem<VecZnx>);
28
29impl PublicKey {
30    pub fn new(module: &Module, log_base2k: usize, log_q: usize) -> PublicKey {
31        PublicKey(Elem::<VecZnx>::new(module, log_base2k, log_q, 2))
32    }
33
34    pub fn gen_thread_safe(
35        &mut self,
36        module: &Module,
37        sk: &ScalarZnxDft,
38        xe: f64,
39        xa_source: &mut Source,
40        xe_source: &mut Source,
41        tmp_bytes: &mut [u8],
42    ) {
43        encrypt_rlwe_sk(
44            module,
45            &mut self.0,
46            None,
47            sk,
48            xa_source,
49            xe_source,
50            xe,
51            tmp_bytes,
52        );
53    }
54
55    pub fn gen_thread_safe_tmp_bytes(module: &Module, log_base2k: usize, log_q: usize) -> usize {
56        encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q)
57    }
58}
59
60pub struct SwitchingKey(pub Ciphertext<MatZnxDft>);
61
62impl SwitchingKey {
63    pub fn new(module: &Module, log_base2k: usize, rows: usize, log_q: usize) -> SwitchingKey {
64        SwitchingKey(new_gadget_ciphertext(module, log_base2k, rows, log_q))
65    }
66
67    pub fn n(&self) -> usize {
68        self.0.n()
69    }
70
71    pub fn rows(&self) -> usize {
72        self.0.rows()
73    }
74
75    pub fn cols(&self) -> usize {
76        self.0.cols()
77    }
78
79    pub fn log_base2k(&self) -> usize {
80        self.0.log_base2k()
81    }
82}