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