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