arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/rlwe/src/encryptor.rs
rlwe/src/encryptor.rsRaw
  1use crate::ciphertext::Ciphertext;
  2use crate::elem::{Elem, ElemCommon, ElemVecZnx};
  3use crate::keys::SecretKey;
  4use crate::parameters::Parameters;
  5use crate::plaintext::Plaintext;
  6use base2k::sampling::Sampling;
  7use base2k::{
  8    Infos, Module, Scalar, SvpPPol, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat,
  9    VmpPMatOps,
 10};
 11
 12use sampling::source::{Source, new_seed};
 13
 14impl Parameters {
 15    pub fn encrypt_rlwe_sk_tmp_bytes(&self, log_q: usize) -> usize {
 16        encrypt_rlwe_sk_tmp_bytes(self.module(), self.log_base2k(), log_q)
 17    }
 18    pub fn encrypt_rlwe_sk(
 19        &self,
 20        ct: &mut Ciphertext<VecZnx>,
 21        pt: Option<&Plaintext>,
 22        sk: &SvpPPol,
 23        source_xa: &mut Source,
 24        source_xe: &mut Source,
 25        tmp_bytes: &mut [u8],
 26    ) {
 27        encrypt_rlwe_sk(
 28            self.module(),
 29            &mut ct.0,
 30            pt.map(|pt| pt.at(0)),
 31            sk,
 32            source_xa,
 33            source_xe,
 34            self.xe(),
 35            tmp_bytes,
 36        )
 37    }
 38}
 39
 40pub struct EncryptorSk {
 41    sk: SvpPPol,
 42    source_xa: Source,
 43    source_xe: Source,
 44    initialized: bool,
 45    tmp_bytes: Vec<u8>,
 46}
 47
 48impl EncryptorSk {
 49    pub fn new(params: &Parameters, sk: Option<&SecretKey>) -> Self {
 50        let mut sk_svp_ppol: SvpPPol = params.module().new_svp_ppol();
 51        let mut initialized: bool = false;
 52        if let Some(sk) = sk {
 53            sk.prepare(params.module(), &mut sk_svp_ppol);
 54            initialized = true;
 55        }
 56        Self {
 57            sk: sk_svp_ppol,
 58            initialized,
 59            source_xa: Source::new(new_seed()),
 60            source_xe: Source::new(new_seed()),
 61            tmp_bytes: vec![0u8; params.encrypt_rlwe_sk_tmp_bytes(params.cols_qp())],
 62        }
 63    }
 64
 65    pub fn set_sk(&mut self, module: &Module, sk: &SecretKey) {
 66        sk.prepare(module, &mut self.sk);
 67        self.initialized = true;
 68    }
 69
 70    pub fn seed_source_xa(&mut self, seed: [u8; 32]) {
 71        self.source_xa = Source::new(seed)
 72    }
 73
 74    pub fn seed_source_xe(&mut self, seed: [u8; 32]) {
 75        self.source_xe = Source::new(seed)
 76    }
 77
 78    pub fn encrypt_rlwe_sk(&mut self, params: &Parameters, ct: &mut Ciphertext<VecZnx>, pt: Option<&Plaintext>) {
 79        assert!(
 80            self.initialized == true,
 81            "invalid call to [EncryptorSk.encrypt_rlwe_sk]: [EncryptorSk] has not been initialized with a [SecretKey]"
 82        );
 83        params.encrypt_rlwe_sk(
 84            ct,
 85            pt,
 86            &self.sk,
 87            &mut self.source_xa,
 88            &mut self.source_xe,
 89            &mut self.tmp_bytes,
 90        );
 91    }
 92
 93    pub fn encrypt_rlwe_sk_core(
 94        &self,
 95        params: &Parameters,
 96        ct: &mut Ciphertext<VecZnx>,
 97        pt: Option<&Plaintext>,
 98        source_xa: &mut Source,
 99        source_xe: &mut Source,
100        tmp_bytes: &mut [u8],
101    ) {
102        assert!(
103            self.initialized == true,
104            "invalid call to [EncryptorSk.encrypt_rlwe_sk]: [EncryptorSk] has not been initialized with a [SecretKey]"
105        );
106        params.encrypt_rlwe_sk(ct, pt, &self.sk, source_xa, source_xe, tmp_bytes);
107    }
108}
109
110pub fn encrypt_rlwe_sk_tmp_bytes(module: &Module, log_base2k: usize, log_q: usize) -> usize {
111    module.bytes_of_vec_znx_dft(1, (log_q + log_base2k - 1) / log_base2k) + module.vec_znx_big_normalize_tmp_bytes()
112}
113pub fn encrypt_rlwe_sk(
114    module: &Module,
115    ct: &mut Elem<VecZnx>,
116    pt: Option<&VecZnx>,
117    sk: &SvpPPol,
118    source_xa: &mut Source,
119    source_xe: &mut Source,
120    sigma: f64,
121    tmp_bytes: &mut [u8],
122) {
123    encrypt_rlwe_sk_core::<0>(module, ct, pt, sk, source_xa, source_xe, sigma, tmp_bytes)
124}
125
126fn encrypt_rlwe_sk_core<const PT_POS: u8>(
127    module: &Module,
128    ct: &mut Elem<VecZnx>,
129    pt: Option<&VecZnx>,
130    sk: &SvpPPol,
131    source_xa: &mut Source,
132    source_xe: &mut Source,
133    sigma: f64,
134    tmp_bytes: &mut [u8],
135) {
136    let cols: usize = ct.cols();
137    let log_base2k: usize = ct.log_base2k();
138    let log_q: usize = ct.log_q();
139
140    assert!(
141        tmp_bytes.len() >= encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q),
142        "invalid tmp_bytes: tmp_bytes={} < encrypt_rlwe_sk_tmp_bytes={}",
143        tmp_bytes.len(),
144        encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q)
145    );
146
147    let log_q: usize = ct.log_q();
148    let log_base2k: usize = ct.log_base2k();
149    let c1: &mut VecZnx = ct.at_mut(1);
150
151    // c1 <- Z_{2^prec}[X]/(X^{N}+1)
152    module.fill_uniform(log_base2k, c1, cols, source_xa);
153
154    let (tmp_bytes_vec_znx_dft, tmp_bytes_normalize) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(1, cols));
155
156    // Scratch space for DFT values
157    let mut buf_dft: VecZnxDft = VecZnxDft::from_bytes_borrow(module, 1, cols, tmp_bytes_vec_znx_dft);
158
159    // Applies buf_dft <- DFT(s) * DFT(c1)
160    module.svp_apply_dft(&mut buf_dft, sk, c1);
161
162    // Alias scratch space
163    let mut buf_big: VecZnxBig = buf_dft.as_vec_znx_big();
164
165    // buf_big = s x c1
166    module.vec_znx_idft_tmp_a(&mut buf_big, &mut buf_dft);
167
168    match PT_POS {
169        // c0 <- -s x c1 + m
170        0 => {
171            let c0: &mut VecZnx = ct.at_mut(0);
172            if let Some(pt) = pt {
173                module.vec_znx_big_sub_small_a_inplace(&mut buf_big, pt);
174                module.vec_znx_big_normalize(log_base2k, c0, &buf_big, tmp_bytes_normalize);
175            } else {
176                module.vec_znx_big_normalize(log_base2k, c0, &buf_big, tmp_bytes_normalize);
177                module.vec_znx_negate_inplace(c0);
178            }
179        }
180        // c1 <- c1 + m
181        1 => {
182            if let Some(pt) = pt {
183                module.vec_znx_add_inplace(c1, pt);
184                c1.normalize(log_base2k, tmp_bytes_normalize);
185            }
186            let c0: &mut VecZnx = ct.at_mut(0);
187            module.vec_znx_big_normalize(log_base2k, c0, &buf_big, tmp_bytes_normalize);
188            module.vec_znx_negate_inplace(c0);
189        }
190        _ => panic!("PT_POS must be 1 or 2"),
191    }
192
193    // c0 <- -s x c1 + m + e
194    module.add_normal(
195        log_base2k,
196        ct.at_mut(0),
197        log_q,
198        source_xe,
199        sigma,
200        (sigma * 6.0).ceil(),
201    );
202}
203
204impl Parameters {
205    pub fn encrypt_grlwe_sk_tmp_bytes(&self, rows: usize, log_q: usize) -> usize {
206        encrypt_grlwe_sk_tmp_bytes(self.module(), self.log_base2k(), rows, log_q)
207    }
208}
209
210pub fn encrypt_grlwe_sk_tmp_bytes(module: &Module, log_base2k: usize, rows: usize, log_q: usize) -> usize {
211    let cols = (log_q + log_base2k - 1) / log_base2k;
212    Elem::<VecZnx>::bytes_of(module, log_base2k, log_q, 2)
213        + Plaintext::bytes_of(module, log_base2k, log_q)
214        + encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q)
215        + module.vmp_prepare_tmp_bytes(rows, cols)
216}
217
218pub fn encrypt_grlwe_sk(
219    module: &Module,
220    ct: &mut Ciphertext<VmpPMat>,
221    m: &Scalar,
222    sk: &SvpPPol,
223    source_xa: &mut Source,
224    source_xe: &mut Source,
225    sigma: f64,
226    tmp_bytes: &mut [u8],
227) {
228    let log_q: usize = ct.log_q();
229    let log_base2k: usize = ct.log_base2k();
230    let (left, right) = ct.0.value.split_at_mut(1);
231    encrypt_grlwe_sk_core::<0>(
232        module,
233        log_base2k,
234        [&mut left[0], &mut right[0]],
235        log_q,
236        m,
237        sk,
238        source_xa,
239        source_xe,
240        sigma,
241        tmp_bytes,
242    )
243}
244
245impl Parameters {
246    pub fn encrypt_rgsw_sk_tmp_bytes(&self, rows: usize, log_q: usize) -> usize {
247        encrypt_rgsw_sk_tmp_bytes(self.module(), self.log_base2k(), rows, log_q)
248    }
249}
250
251pub fn encrypt_rgsw_sk_tmp_bytes(module: &Module, log_base2k: usize, rows: usize, log_q: usize) -> usize {
252    let cols = (log_q + log_base2k - 1) / log_base2k;
253    Elem::<VecZnx>::bytes_of(module, log_base2k, log_q, 2)
254        + Plaintext::bytes_of(module, log_base2k, log_q)
255        + encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q)
256        + module.vmp_prepare_tmp_bytes(rows, cols)
257}
258
259pub fn encrypt_rgsw_sk(
260    module: &Module,
261    ct: &mut Ciphertext<VmpPMat>,
262    m: &Scalar,
263    sk: &SvpPPol,
264    source_xa: &mut Source,
265    source_xe: &mut Source,
266    sigma: f64,
267    tmp_bytes: &mut [u8],
268) {
269    let log_q: usize = ct.log_q();
270    let log_base2k: usize = ct.log_base2k();
271
272    let (left, right) = ct.0.value.split_at_mut(2);
273    let (ll, lr) = left.split_at_mut(1);
274    let (rl, rr) = right.split_at_mut(1);
275
276    encrypt_grlwe_sk_core::<0>(
277        module,
278        log_base2k,
279        [&mut ll[0], &mut lr[0]],
280        log_q,
281        m,
282        sk,
283        source_xa,
284        source_xe,
285        sigma,
286        tmp_bytes,
287    );
288    encrypt_grlwe_sk_core::<1>(
289        module,
290        log_base2k,
291        [&mut rl[0], &mut rr[0]],
292        log_q,
293        m,
294        sk,
295        source_xa,
296        source_xe,
297        sigma,
298        tmp_bytes,
299    );
300}
301
302fn encrypt_grlwe_sk_core<const PT_POS: u8>(
303    module: &Module,
304    log_base2k: usize,
305    mut ct: [&mut VmpPMat; 2],
306    log_q: usize,
307    m: &Scalar,
308    sk: &SvpPPol,
309    source_xa: &mut Source,
310    source_xe: &mut Source,
311    sigma: f64,
312    tmp_bytes: &mut [u8],
313) {
314    let rows: usize = ct[0].rows();
315
316    let min_tmp_bytes_len = encrypt_grlwe_sk_tmp_bytes(module, log_base2k, rows, log_q);
317
318    assert!(
319        tmp_bytes.len() >= min_tmp_bytes_len,
320        "invalid tmp_bytes: tmp_bytes.len()={} < encrypt_grlwe_sk_tmp_bytes={}",
321        tmp_bytes.len(),
322        min_tmp_bytes_len
323    );
324
325    let bytes_of_elem: usize = Elem::<VecZnx>::bytes_of(module, log_base2k, log_q, 2);
326    let bytes_of_pt: usize = Plaintext::bytes_of(module, log_base2k, log_q);
327    let bytes_of_enc_sk: usize = encrypt_rlwe_sk_tmp_bytes(module, log_base2k, log_q);
328
329    let (tmp_bytes_pt, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_pt);
330    let (tmp_bytes_enc_sk, tmp_bytes) = tmp_bytes.split_at_mut(bytes_of_enc_sk);
331    let (tmp_bytes_elem, tmp_bytes_vmp_prepare_row) = tmp_bytes.split_at_mut(bytes_of_elem);
332
333    let mut tmp_elem: Elem<VecZnx> = Elem::<VecZnx>::from_bytes_borrow(module, log_base2k, log_q, 2, tmp_bytes_elem);
334    let mut tmp_pt: Plaintext = Plaintext::from_bytes_borrow(module, log_base2k, log_q, tmp_bytes_pt);
335
336    (0..rows).for_each(|row_i| {
337        // Sets the i-th row of the RLWE sample to m (i.e. m * 2^{-log_base2k*i})
338        tmp_pt.at_mut(0).at_mut(row_i).copy_from_slice(&m.raw());
339
340        // Encrypts RLWE(m * 2^{-log_base2k*i})
341        encrypt_rlwe_sk_core::<PT_POS>(
342            module,
343            &mut tmp_elem,
344            Some(tmp_pt.at(0)),
345            sk,
346            source_xa,
347            source_xe,
348            sigma,
349            tmp_bytes_enc_sk,
350        );
351
352        // Zeroes the ith-row of tmp_pt
353        tmp_pt.at_mut(0).at_mut(row_i).fill(0);
354
355        // GRLWE[row_i][0||1] = [-as + m * 2^{-i*log_base2k} + e*2^{-log_q} || a]
356        module.vmp_prepare_row(
357            ct[0],
358            tmp_elem.at(0).raw(),
359            row_i,
360            tmp_bytes_vmp_prepare_row,
361        );
362        module.vmp_prepare_row(
363            &mut ct[1],
364            tmp_elem.at(1).raw(),
365            row_i,
366            tmp_bytes_vmp_prepare_row,
367        );
368    });
369}