1use crate::ciphertext::{Ciphertext, GadgetCiphertext};
  2use crate::elem::Elem;
  3use crate::keys::SecretKey;
  4use crate::parameters::Parameters;
  5use crate::plaintext::Plaintext;
  6use base2k::ffi::znx::znx_zero_i64_ref;
  7use base2k::sampling::Sampling;
  8use base2k::{
  9    Module, Scalar, SvpPPol, SvpPPolOps, VecZnx, VecZnxBig, VecZnxDft, VecZnxOps, VmpPMatOps,
 10};
 11use sampling::source::{Source, new_seed};
 12
 13pub struct EncryptorSk {
 14    sk: SvpPPol,
 15    source_xa: Source,
 16    source_xe: Source,
 17    initialized: bool,
 18    tmp_bytes: Vec<u8>,
 19}
 20
 21impl EncryptorSk {
 22    pub fn new(params: &Parameters, sk: Option<&SecretKey>) -> Self {
 23        let mut sk_svp_ppol: SvpPPol = params.module().svp_new_ppol();
 24        let mut initialized: bool = false;
 25        if let Some(sk) = sk {
 26            sk.prepare(params.module(), &mut sk_svp_ppol);
 27            initialized = true;
 28        }
 29        Self {
 30            sk: sk_svp_ppol,
 31            initialized,
 32            source_xa: Source::new(new_seed()),
 33            source_xe: Source::new(new_seed()),
 34            tmp_bytes: vec![0u8; params.encrypt_rlwe_sk_tmp_bytes(params.limbs_qp())],
 35        }
 36    }
 37
 38    pub fn set_sk(&mut self, module: &Module, sk: &SecretKey) {
 39        sk.prepare(module, &mut self.sk);
 40        self.initialized = true;
 41    }
 42
 43    pub fn seed_source_xa(&mut self, seed: [u8; 32]) {
 44        self.source_xa = Source::new(seed)
 45    }
 46
 47    pub fn seed_source_xe(&mut self, seed: [u8; 32]) {
 48        self.source_xe = Source::new(seed)
 49    }
 50
 51    pub fn encrypt_rlwe_sk(
 52        &mut self,
 53        params: &Parameters,
 54        ct: &mut Ciphertext,
 55        pt: Option<&Plaintext>,
 56    ) {
 57        assert!(
 58            self.initialized == true,
 59            "invalid call to [EncryptorSk.encrypt_rlwe_sk]: [EncryptorSk] has not been initialized with a [SecretKey]"
 60        );
 61        params.encrypt_rlwe_sk_thread_safe(
 62            ct,
 63            pt,
 64            &self.sk,
 65            &mut self.source_xa,
 66            &mut self.source_xe,
 67            &mut self.tmp_bytes,
 68        );
 69    }
 70
 71    pub fn encrypt_rlwe_sk_thread_safe(
 72        &self,
 73        params: &Parameters,
 74        ct: &mut Ciphertext,
 75        pt: Option<&Plaintext>,
 76        source_xa: &mut Source,
 77        source_xe: &mut Source,
 78        tmp_bytes: &mut [u8],
 79    ) {
 80        assert!(
 81            self.initialized == true,
 82            "invalid call to [EncryptorSk.encrypt_rlwe_sk_thread_safe]: [EncryptorSk] has not been initialized with a [SecretKey]"
 83        );
 84        params.encrypt_rlwe_sk_thread_safe(ct, pt, &self.sk, source_xa, source_xe, tmp_bytes);
 85    }
 86}
 87
 88impl Parameters {
 89    pub fn encrypt_rlwe_sk_tmp_bytes(&self, log_q: usize) -> usize {
 90        encrypt_rlwe_sk_tmp_bytes(self.module(), self.log_base2k(), log_q)
 91    }
 92
 93    pub fn encrypt_rlwe_sk_thread_safe(
 94        &self,
 95        ct: &mut Ciphertext,
 96        pt: Option<&Plaintext>,
 97        sk: &SvpPPol,
 98        source_xa: &mut Source,
 99        source_xe: &mut Source,
100        tmp_bytes: &mut [u8],
101    ) {
102        encrypt_rlwe_sk_thread_safe(
103            self.module(),
104            &mut ct.0,
105            pt.map(|pt| &pt.0),
106            sk,
107            source_xa,
108            source_xe,
109            self.xe(),
110            tmp_bytes,
111        )
112    }
113}
114
115pub fn encrypt_rlwe_sk_tmp_bytes(module: &Module, log_base2k: usize, log_q: usize) -> usize {
116    module.bytes_of_vec_znx_dft((log_q + log_base2k - 1) / log_base2k)
117        + module.vec_znx_big_normalize_tmp_bytes()
118}
119
120pub fn encrypt_rlwe_sk_thread_safe(
121    module: &Module,
122    ct: &mut Elem,
123    pt: Option<&Elem>,
124    sk: &SvpPPol,
125    source_xa: &mut Source,
126    source_xe: &mut Source,
127    sigma: f64,
128    tmp_bytes: &mut [u8],
129) {
130    let limbs: usize = ct.limbs();
131    let log_base2k: usize = ct.log_base2k();
132    let log_q: usize = ct.log_q();
133
134    assert!(
135        tmp_bytes.len() >= encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q),
136        "invalid tmp_bytes: tmp_bytes={} < encrypt_rlwe_sk_tmp_bytes={}",
137        tmp_bytes.len(),
138        encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q)
139    );
140
141    let log_q: usize = ct.log_q();
142    let log_base2k: usize = ct.log_base2k();
143    let c1: &mut VecZnx = ct.at_mut(1);
144
145    // c1 <- Z_{2^prec}[X]/(X^{N}+1)
146    c1.fill_uniform(log_base2k, limbs, source_xa);
147
148    let bytes_of_vec_znx_dft = module.bytes_of_vec_znx_dft(limbs);
149
150    // Scratch space for DFT values
151    let mut buf_dft: VecZnxDft =
152        VecZnxDft::from_bytes(limbs, &mut tmp_bytes[..bytes_of_vec_znx_dft]);
153
154    // Applies buf_dft <- s * c1
155    module.svp_apply_dft(&mut buf_dft, sk, c1, limbs);
156
157    // Alias scratch space
158    let mut buf_big: VecZnxBig = buf_dft.as_vec_znx_big();
159
160    if let Some(pt) = pt {
161        // buf_big <- m - buf_big
162        module.vec_znx_big_sub_small_a_inplace(&mut buf_big, pt.at(0));
163    };
164
165    let carry: &mut [u8] = &mut tmp_bytes[bytes_of_vec_znx_dft..];
166
167    // c0 <- normalize(buf_big) + e
168    let c0: &mut VecZnx = ct.at_mut(0);
169    module.vec_znx_big_normalize(log_base2k, c0, &buf_big, carry);
170    c0.add_normal(log_base2k, log_q, source_xe, sigma, (sigma * 6.0).ceil());
171}
172
173pub fn encrypt_grlwe_sk_tmp_bytes(
174    module: &Module,
175    log_base2k: usize,
176    rows: usize,
177    log_q: usize,
178) -> usize {
179    let cols = (log_q + log_base2k - 1) / log_base2k;
180    Elem::bytes_of(module, log_base2k, log_q, 1)
181        + Plaintext::bytes_of(module, log_base2k, log_q)
182        + encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q)
183        + module.vmp_prepare_tmp_bytes(rows, cols)
184}
185
186pub fn encrypt_grlwe_sk_thread_safe(
187    module: &Module,
188    ct: &mut GadgetCiphertext,
189    m: &Scalar,
190    sk: &SvpPPol,
191    source_xa: &mut Source,
192    source_xe: &mut Source,
193    sigma: f64,
194    tmp_bytes: &mut [u8],
195) {
196    let rows: usize = ct.rows();
197    let log_q: usize = ct.log_q();
198    let log_base2k: usize = ct.log_base2k();
199
200    let min_tmp_bytes_len = encrypt_grlwe_sk_tmp_bytes(module, log_base2k, rows, log_q);
201
202    assert!(
203        tmp_bytes.len() >= min_tmp_bytes_len,
204        "invalid tmp_bytes: tmp_bytes.len()={} < encrypt_grlwe_sk_tmp_bytes={}",
205        tmp_bytes.len(),
206        min_tmp_bytes_len
207    );
208
209    let mut ptr: usize = 0;
210    let mut tmp_elem: Elem = Elem::from_bytes(module, log_base2k, ct.log_q(), 1, tmp_bytes);
211    let bytes_of_elem: usize = Elem::bytes_of(module, log_base2k, log_q, 1);
212    ptr += bytes_of_elem;
213
214    let mut tmp_pt: Plaintext =
215        Plaintext::from_bytes(module, log_base2k, log_q, &mut tmp_bytes[ptr..]);
216    ptr += Plaintext::bytes_of(module, log_base2k, log_q);
217
218    let (tmp_bytes_encrypt_sk, tmp_bytes_vmp_prepare_row) =
219        tmp_bytes[ptr..].split_at_mut(encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q));
220
221    (0..rows).for_each(|row_i| {
222        // Sets the i-th row of the RLWE sample to m (i.e. m * 2^{-log_base2k*i})
223        tmp_pt.0.value[0].at_mut(row_i).copy_from_slice(&m.0);
224
225        // Encrypts RLWE(m * 2^{-log_base2k*i})
226        encrypt_rlwe_sk_thread_safe(
227            module,
228            &mut tmp_elem,
229            Some(&tmp_pt.0),
230            sk,
231            source_xa,
232            source_xe,
233            sigma,
234            tmp_bytes_encrypt_sk,
235        );
236
237        // Zeroes the ith-row of tmp_pt
238        tmp_pt.0.value[0].at_mut(row_i).fill(0);
239
240        // GRLWE[row_i][0] = -as + m * 2^{-i*log_base2k} + e*2^{-log_q}
241        module.vmp_prepare_row(
242            &mut ct.value[0],
243            tmp_elem.at(0),
244            row_i,
245            tmp_bytes_vmp_prepare_row,
246        );
247
248        // GRLWE[row_i][1] = a
249        module.vmp_prepare_row(
250            &mut ct.value[1],
251            tmp_elem.at(1),
252            row_i,
253            tmp_bytes_vmp_prepare_row,
254        );
255    })
256}