added grlwe sk encryption

base2k/src/vec_znx_big.rs

@@ -1,6 +1,6 @@
 use crate::ffi::vec_znx_big;
 use crate::znx_base::{ZnxInfos, ZnxView};
-use crate::{Backend, DataView, DataViewMut, FFT64, Module, VecZnxDft, ZnxSliceSize, alloc_aligned};
+use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxSliceSize, alloc_aligned};
 use std::fmt;
 use std::marker::PhantomData;
 

rlwe/src/elem.rs

@@ -1,7 +1,6 @@
-use base2k::{
-    Backend, MatZnxDft, MatZnxDftAlloc, MatZnxDftToMut, MatZnxDftToRef, Module, VecZnx, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc,
-    VecZnxDftToMut, VecZnxDftToRef, VecZnxToMut, VecZnxToRef, ZnxInfos,
-};
+use base2k::ZnxInfos;
+
+use crate::utils::derive_size;
 
 pub trait Infos {
     type Inner: ZnxInfos;
@@ -46,257 +45,3 @@ pub trait Infos {
     /// Returns the bit precision of the ciphertext.
     fn log_k(&self) -> usize;
 }
-
-pub struct RLWECt<C> {
-    pub data: VecZnx<C>,
-    pub log_base2k: usize,
-    pub log_k: usize,
-}
-
-impl<T> Infos for RLWECt<T> {
-    type Inner = VecZnx<T>;
-
-    fn inner(&self) -> &Self::Inner {
-        &self.data
-    }
-
-    fn log_base2k(&self) -> usize {
-        self.log_base2k
-    }
-
-    fn log_k(&self) -> usize {
-        self.log_k
-    }
-}
-
-impl<C> VecZnxToMut for RLWECt<C>
-where
-    VecZnx<C>: VecZnxToMut,
-{
-    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
-        self.data.to_mut()
-    }
-}
-
-impl<C> VecZnxToRef for RLWECt<C>
-where
-    VecZnx<C>: VecZnxToRef,
-{
-    fn to_ref(&self) -> VecZnx<&[u8]> {
-        self.data.to_ref()
-    }
-}
-
-pub struct RLWEPt<C> {
-    pub data: VecZnx<C>,
-    pub log_base2k: usize,
-    pub log_k: usize,
-}
-
-impl<T> Infos for RLWEPt<T> {
-    type Inner = VecZnx<T>;
-
-    fn inner(&self) -> &Self::Inner {
-        &self.data
-    }
-
-    fn log_base2k(&self) -> usize {
-        self.log_base2k
-    }
-
-    fn log_k(&self) -> usize {
-        self.log_k
-    }
-}
-
-impl<C> VecZnxToMut for RLWEPt<C>
-where
-    VecZnx<C>: VecZnxToMut,
-{
-    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
-        self.data.to_mut()
-    }
-}
-
-impl<C> VecZnxToRef for RLWEPt<C>
-where
-    VecZnx<C>: VecZnxToRef,
-{
-    fn to_ref(&self) -> VecZnx<&[u8]> {
-        self.data.to_ref()
-    }
-}
-
-impl RLWECt<Vec<u8>> {
-    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
-        Self {
-            data: module.new_vec_znx(cols, derive_size(log_base2k, log_k)),
-            log_base2k: log_base2k,
-            log_k: log_k,
-        }
-    }
-}
-
-impl RLWEPt<Vec<u8>> {
-    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
-        Self {
-            data: module.new_vec_znx(1, derive_size(log_base2k, log_k)),
-            log_base2k: log_base2k,
-            log_k: log_k,
-        }
-    }
-}
-
-pub struct RLWECtDft<C, B: Backend> {
-    pub data: VecZnxDft<C, B>,
-    pub log_base2k: usize,
-    pub log_k: usize,
-}
-
-impl<B: Backend> RLWECtDft<Vec<u8>, B> {
-    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
-        Self {
-            data: module.new_vec_znx_dft(cols, derive_size(log_base2k, log_k)),
-            log_base2k: log_base2k,
-            log_k: log_k,
-        }
-    }
-}
-
-impl<T, B: Backend> Infos for RLWECtDft<T, B> {
-    type Inner = VecZnxDft<T, B>;
-
-    fn inner(&self) -> &Self::Inner {
-        &self.data
-    }
-
-    fn log_base2k(&self) -> usize {
-        self.log_base2k
-    }
-
-    fn log_k(&self) -> usize {
-        self.log_k
-    }
-}
-
-impl<C, B: Backend> VecZnxDftToMut<B> for RLWECtDft<C, B>
-where
-    VecZnxDft<C, B>: VecZnxDftToMut<B>,
-{
-    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
-        self.data.to_mut()
-    }
-}
-
-impl<C, B: Backend> VecZnxDftToRef<B> for RLWECtDft<C, B>
-where
-    VecZnxDft<C, B>: VecZnxDftToRef<B>,
-{
-    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
-        self.data.to_ref()
-    }
-}
-
-pub struct GRLWECt<C, B: Backend> {
-    pub data: MatZnxDft<C, B>,
-    pub log_base2k: usize,
-    pub log_k: usize,
-}
-
-impl<B: Backend> GRLWECt<Vec<u8>, B> {
-    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
-        Self {
-            data: module.new_mat_znx_dft(rows, 1, 2, derive_size(log_base2k, log_k)),
-            log_base2k: log_base2k,
-            log_k: log_k,
-        }
-    }
-}
-
-impl<T, B: Backend> Infos for GRLWECt<T, B> {
-    type Inner = MatZnxDft<T, B>;
-
-    fn inner(&self) -> &Self::Inner {
-        &self.data
-    }
-
-    fn log_base2k(&self) -> usize {
-        self.log_base2k
-    }
-
-    fn log_k(&self) -> usize {
-        self.log_k
-    }
-}
-
-impl<C, B: Backend> MatZnxDftToMut<B> for GRLWECt<C, B>
-where
-    MatZnxDft<C, B>: MatZnxDftToMut<B>,
-{
-    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
-        self.data.to_mut()
-    }
-}
-
-impl<C, B: Backend> MatZnxDftToRef<B> for GRLWECt<C, B>
-where
-    MatZnxDft<C, B>: MatZnxDftToRef<B>,
-{
-    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
-        self.data.to_ref()
-    }
-}
-
-pub struct RGSWCt<C, B: Backend> {
-    pub data: MatZnxDft<C, B>,
-    pub log_base2k: usize,
-    pub log_k: usize,
-}
-
-impl<B: Backend> RGSWCt<Vec<u8>, B> {
-    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
-        Self {
-            data: module.new_mat_znx_dft(rows, 2, 2, derive_size(log_base2k, log_k)),
-            log_base2k: log_base2k,
-            log_k: log_k,
-        }
-    }
-}
-
-impl<T, B: Backend> Infos for RGSWCt<T, B> {
-    type Inner = MatZnxDft<T, B>;
-
-    fn inner(&self) -> &Self::Inner {
-        &self.data
-    }
-
-    fn log_base2k(&self) -> usize {
-        self.log_base2k
-    }
-
-    fn log_k(&self) -> usize {
-        self.log_k
-    }
-}
-
-impl<C, B: Backend> MatZnxDftToMut<B> for RGSWCt<C, B>
-where
-    MatZnxDft<C, B>: MatZnxDftToMut<B>,
-{
-    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
-        self.data.to_mut()
-    }
-}
-
-impl<C, B: Backend> MatZnxDftToRef<B> for RGSWCt<C, B>
-where
-    MatZnxDft<C, B>: MatZnxDftToRef<B>,
-{
-    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
-        self.data.to_ref()
-    }
-}
-
-pub(crate) fn derive_size(log_base2k: usize, log_k: usize) -> usize {
-    (log_k + log_base2k - 1) / log_base2k
-}

rlwe/src/elem_grlwe.rs

@@ -0,0 +1,53 @@
+use base2k::{Backend, MatZnxDft, MatZnxDftAlloc, MatZnxDftToMut, MatZnxDftToRef, Module};
+
+use crate::{elem::Infos, utils::derive_size};
+
+pub struct GRLWECt<C, B: Backend> {
+    pub data: MatZnxDft<C, B>,
+    pub log_base2k: usize,
+    pub log_k: usize,
+}
+
+impl<B: Backend> GRLWECt<Vec<u8>, B> {
+    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
+        Self {
+            data: module.new_mat_znx_dft(rows, 1, 2, derive_size(log_base2k, log_k)),
+            log_base2k: log_base2k,
+            log_k: log_k,
+        }
+    }
+}
+
+impl<T, B: Backend> Infos for GRLWECt<T, B> {
+    type Inner = MatZnxDft<T, B>;
+
+    fn inner(&self) -> &Self::Inner {
+        &self.data
+    }
+
+    fn log_base2k(&self) -> usize {
+        self.log_base2k
+    }
+
+    fn log_k(&self) -> usize {
+        self.log_k
+    }
+}
+
+impl<C, B: Backend> MatZnxDftToMut<B> for GRLWECt<C, B>
+where
+    MatZnxDft<C, B>: MatZnxDftToMut<B>,
+{
+    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
+        self.data.to_mut()
+    }
+}
+
+impl<C, B: Backend> MatZnxDftToRef<B> for GRLWECt<C, B>
+where
+    MatZnxDft<C, B>: MatZnxDftToRef<B>,
+{
+    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
+        self.data.to_ref()
+    }
+}

rlwe/src/elem_rgsw.rs

@@ -0,0 +1,140 @@
+use base2k::{
+    Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
+    ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxDftAlloc, VecZnxDftOps, ZnxView, ZnxViewMut,
+};
+use sampling::source::Source;
+
+use crate::{
+    elem::Infos,
+    elem_grlwe::GRLWECt,
+    elem_rlwe::{RLWECt, RLWECtDft, RLWEPt},
+    keys::SecretKeyDft,
+    utils::derive_size,
+};
+
+pub struct RGSWCt<C, B: Backend> {
+    pub data: MatZnxDft<C, B>,
+    pub log_base2k: usize,
+    pub log_k: usize,
+}
+
+impl<B: Backend> RGSWCt<Vec<u8>, B> {
+    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
+        Self {
+            data: module.new_mat_znx_dft(rows, 2, 2, derive_size(log_base2k, log_k)),
+            log_base2k: log_base2k,
+            log_k: log_k,
+        }
+    }
+}
+
+impl<T, B: Backend> Infos for RGSWCt<T, B> {
+    type Inner = MatZnxDft<T, B>;
+
+    fn inner(&self) -> &Self::Inner {
+        &self.data
+    }
+
+    fn log_base2k(&self) -> usize {
+        self.log_base2k
+    }
+
+    fn log_k(&self) -> usize {
+        self.log_k
+    }
+}
+
+impl<C, B: Backend> MatZnxDftToMut<B> for RGSWCt<C, B>
+where
+    MatZnxDft<C, B>: MatZnxDftToMut<B>,
+{
+    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
+        self.data.to_mut()
+    }
+}
+
+impl<C, B: Backend> MatZnxDftToRef<B> for RGSWCt<C, B>
+where
+    MatZnxDft<C, B>: MatZnxDftToRef<B>,
+{
+    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
+        self.data.to_ref()
+    }
+}
+
+impl GRLWECt<Vec<u8>, FFT64> {
+    pub fn encrypt_sk_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
+        RLWECt::encrypt_sk_scratch_bytes(module, size)
+            + module.bytes_of_vec_znx(2, size)
+            + module.bytes_of_vec_znx(1, size)
+            + module.bytes_of_vec_znx_dft(2, size)
+    }
+
+    pub fn encrypt_pk_scratch_bytes(module: &Module<FFT64>, pk_size: usize) -> usize {
+        RLWECt::encrypt_pk_scratch_bytes(module, pk_size)
+    }
+
+    pub fn decrypt_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
+        RLWECtDft::decrypt_scratch_bytes(module, size)
+    }
+}
+
+pub fn encrypt_grlwe_sk<C, P, S>(
+    module: &Module<FFT64>,
+    ct: &mut GRLWECt<C, FFT64>,
+    pt: &ScalarZnx<P>,
+    sk: &SecretKeyDft<S, FFT64>,
+    source_xa: &mut Source,
+    source_xe: &mut Source,
+    sigma: f64,
+    bound: f64,
+    scratch: &mut Scratch,
+) where
+    MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
+    ScalarZnx<P>: ScalarZnxToRef,
+    ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
+{
+    let rows: usize = ct.rows();
+    let size: usize = ct.size();
+
+    let (tmp_znx_pt, scrach_1) = scratch.tmp_vec_znx(module, 1, size);
+    let (tmp_znx_ct, scrach_2) = scrach_1.tmp_vec_znx(module, 2, size);
+    let (mut tmp_dft, scratch_3) = scrach_2.tmp_vec_znx_dft(module, 2, size);
+
+    let mut tmp_pt: RLWEPt<&mut [u8]> = RLWEPt {
+        data: tmp_znx_pt,
+        log_base2k: ct.log_base2k(),
+        log_k: ct.log_k(),
+    };
+
+    let mut tmp_ct: RLWECt<&mut [u8]> = RLWECt {
+        data: tmp_znx_ct,
+        log_base2k: ct.log_base2k(),
+        log_k: ct.log_k(),
+    };
+
+    (0..rows).for_each(|row_i| {
+        tmp_pt
+            .data
+            .at_mut(0, row_i)
+            .copy_from_slice(&pt.to_ref().raw());
+
+        tmp_ct.encrypt_sk(
+            module,
+            Some(&tmp_pt),
+            sk,
+            source_xa,
+            source_xe,
+            sigma,
+            bound,
+            scratch_3,
+        );
+
+        tmp_pt.data.at_mut(0, row_i).fill(0);
+
+        module.vec_znx_dft(&mut tmp_dft, 0, &tmp_ct, 0);
+        module.vec_znx_dft(&mut tmp_dft, 1, &tmp_ct, 1);
+
+        module.vmp_prepare_row(ct, row_i, 0, &tmp_dft);
+    });
+}

rlwe/src/elem_rlwe.rs

@@ -1,22 +1,182 @@
-use std::cmp::min;
-
 use base2k::{
     AddNormal, Backend, FFT64, FillUniform, Module, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps,
     ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc,
     VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, VecZnxToMut, VecZnxToRef,
 };
-
 use sampling::source::Source;
 
 use crate::{
-    elem::{Infos, RLWECt, RLWECtDft, RLWEPt},
+    elem::Infos,
     keys::{PublicKey, SecretDistribution, SecretKeyDft},
+    utils::derive_size,
 };
 
-pub fn encrypt_rlwe_sk_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
+pub struct RLWECt<C> {
+    pub data: VecZnx<C>,
+    pub log_base2k: usize,
+    pub log_k: usize,
+}
+
+impl RLWECt<Vec<u8>> {
+    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
+        Self {
+            data: module.new_vec_znx(cols, derive_size(log_base2k, log_k)),
+            log_base2k: log_base2k,
+            log_k: log_k,
+        }
+    }
+}
+
+impl<T> Infos for RLWECt<T> {
+    type Inner = VecZnx<T>;
+
+    fn inner(&self) -> &Self::Inner {
+        &self.data
+    }
+
+    fn log_base2k(&self) -> usize {
+        self.log_base2k
+    }
+
+    fn log_k(&self) -> usize {
+        self.log_k
+    }
+}
+
+impl<C> VecZnxToMut for RLWECt<C>
+where
+    VecZnx<C>: VecZnxToMut,
+{
+    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
+        self.data.to_mut()
+    }
+}
+
+impl<C> VecZnxToRef for RLWECt<C>
+where
+    VecZnx<C>: VecZnxToRef,
+{
+    fn to_ref(&self) -> VecZnx<&[u8]> {
+        self.data.to_ref()
+    }
+}
+
+pub struct RLWEPt<C> {
+    pub data: VecZnx<C>,
+    pub log_base2k: usize,
+    pub log_k: usize,
+}
+
+impl<T> Infos for RLWEPt<T> {
+    type Inner = VecZnx<T>;
+
+    fn inner(&self) -> &Self::Inner {
+        &self.data
+    }
+
+    fn log_base2k(&self) -> usize {
+        self.log_base2k
+    }
+
+    fn log_k(&self) -> usize {
+        self.log_k
+    }
+}
+
+impl<C> VecZnxToMut for RLWEPt<C>
+where
+    VecZnx<C>: VecZnxToMut,
+{
+    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
+        self.data.to_mut()
+    }
+}
+
+impl<C> VecZnxToRef for RLWEPt<C>
+where
+    VecZnx<C>: VecZnxToRef,
+{
+    fn to_ref(&self) -> VecZnx<&[u8]> {
+        self.data.to_ref()
+    }
+}
+
+impl RLWEPt<Vec<u8>> {
+    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
+        Self {
+            data: module.new_vec_znx(1, derive_size(log_base2k, log_k)),
+            log_base2k: log_base2k,
+            log_k: log_k,
+        }
+    }
+}
+
+pub struct RLWECtDft<C, B: Backend> {
+    pub data: VecZnxDft<C, B>,
+    pub log_base2k: usize,
+    pub log_k: usize,
+}
+
+impl<B: Backend> RLWECtDft<Vec<u8>, B> {
+    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
+        Self {
+            data: module.new_vec_znx_dft(cols, derive_size(log_base2k, log_k)),
+            log_base2k: log_base2k,
+            log_k: log_k,
+        }
+    }
+}
+
+impl<T, B: Backend> Infos for RLWECtDft<T, B> {
+    type Inner = VecZnxDft<T, B>;
+
+    fn inner(&self) -> &Self::Inner {
+        &self.data
+    }
+
+    fn log_base2k(&self) -> usize {
+        self.log_base2k
+    }
+
+    fn log_k(&self) -> usize {
+        self.log_k
+    }
+}
+
+impl<C, B: Backend> VecZnxDftToMut<B> for RLWECtDft<C, B>
+where
+    VecZnxDft<C, B>: VecZnxDftToMut<B>,
+{
+    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
+        self.data.to_mut()
+    }
+}
+
+impl<C, B: Backend> VecZnxDftToRef<B> for RLWECtDft<C, B>
+where
+    VecZnxDft<C, B>: VecZnxDftToRef<B>,
+{
+    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
+        self.data.to_ref()
+    }
+}
+
+impl RLWECt<Vec<u8>> {
+    pub fn encrypt_sk_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
         (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
     }
 
+    pub fn encrypt_pk_scratch_bytes<B: Backend>(module: &Module<B>, pk_size: usize) -> usize {
+        ((module.bytes_of_vec_znx_dft(1, pk_size) + module.bytes_of_vec_znx_big(1, pk_size)) | module.bytes_of_scalar_znx(1))
+            + module.bytes_of_scalar_znx_dft(1)
+            + module.vec_znx_big_normalize_tmp_bytes()
+    }
+
+    pub fn decrypt_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
+        (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
+    }
+}
+
 pub fn encrypt_rlwe_sk<C, P, S>(
     module: &Module<FFT64>,
     ct: &mut RLWECt<C>,
@@ -94,11 +254,7 @@ pub fn decrypt_rlwe<P, C, S>(
     module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
 
     pt.log_base2k = ct.log_base2k();
-    pt.log_k = min(pt.log_k(), ct.log_k());
-}
-
-pub fn decrypt_rlwe_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
-    (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
+    pt.log_k = pt.log_k().min(ct.log_k());
 }
 
 impl<C> RLWECt<C> {
@@ -207,13 +363,22 @@ pub(crate) fn encrypt_zero_rlwe_dft_sk<C, S>(
     module.vec_znx_dft(ct, 0, &tmp_znx, 0);
 }
 
-pub(crate) fn encrypt_zero_rlwe_dft_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
+impl RLWECtDft<Vec<u8>, FFT64> {
+    pub fn encrypt_zero_sk_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
         (module.bytes_of_vec_znx(1, size) | module.bytes_of_vec_znx_dft(1, size))
             + module.bytes_of_vec_znx_big(1, size)
             + module.bytes_of_vec_znx(1, size)
             + module.vec_znx_big_normalize_tmp_bytes()
     }
 
+    pub fn decrypt_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
+        (module.vec_znx_big_normalize_tmp_bytes()
+            | module.bytes_of_vec_znx_dft(1, size)
+            | (module.bytes_of_vec_znx_big(1, size) + module.vec_znx_idft_tmp_bytes()))
+            + module.bytes_of_vec_znx_big(1, size)
+    }
+}
+
 pub fn decrypt_rlwe_dft<P, C, S>(
     module: &Module<FFT64>,
     pt: &mut RLWEPt<P>,
@@ -246,14 +411,7 @@ pub fn decrypt_rlwe_dft<P, C, S>(
     module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
 
     pt.log_base2k = ct.log_base2k();
-    pt.log_k = min(pt.log_k(), ct.log_k());
-}
-
-pub fn decrypt_rlwe_dft_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
-    (module.vec_znx_big_normalize_tmp_bytes()
-        | module.bytes_of_vec_znx_dft(1, size)
-        | (module.bytes_of_vec_znx_big(1, size) + module.vec_znx_idft_tmp_bytes()))
-        + module.bytes_of_vec_znx_big(1, size)
+    pt.log_k = pt.log_k().min(ct.log_k());
 }
 
 impl<C> RLWECtDft<C, FFT64> {
@@ -290,12 +448,6 @@ impl<C> RLWECtDft<C, FFT64> {
     }
 }
 
-pub fn encrypt_rlwe_pk_scratch_bytes<B: Backend>(module: &Module<B>, pk_size: usize) -> usize {
-    ((module.bytes_of_vec_znx_dft(1, pk_size) + module.bytes_of_vec_znx_big(1, pk_size)) | module.bytes_of_scalar_znx(1))
-        + module.bytes_of_scalar_znx_dft(1)
-        + module.vec_znx_big_normalize_tmp_bytes()
-}
-
 pub(crate) fn encrypt_rlwe_pk<C, P, S>(
     module: &Module<FFT64>,
     ct: &mut RLWECt<C>,
@@ -369,13 +521,10 @@ mod tests {
     use sampling::source::Source;
 
     use crate::{
-        elem::{Infos, RLWECt, RLWECtDft, RLWEPt},
-        encryption::{decrypt_rlwe_dft_scratch_bytes, encrypt_zero_rlwe_dft_scratch_bytes},
+        elem_rlwe::{Infos, RLWECt, RLWECtDft, RLWEPt},
         keys::{PublicKey, SecretKey, SecretKeyDft},
     };
 
-    use super::{decrypt_rlwe_scratch_bytes, encrypt_rlwe_pk_scratch_bytes, encrypt_rlwe_sk_scratch_bytes};
-
     #[test]
     fn encrypt_sk_vec_znx_fft64() {
         let module: Module<FFT64> = Module::<FFT64>::new(32);
@@ -393,8 +542,9 @@ mod tests {
         let mut source_xe: Source = Source::new([0u8; 32]);
         let mut source_xa: Source = Source::new([0u8; 32]);
 
-        let mut scratch: ScratchOwned =
-            ScratchOwned::new(encrypt_rlwe_sk_scratch_bytes(&module, ct.size()) | decrypt_rlwe_scratch_bytes(&module, ct.size()));
+        let mut scratch: ScratchOwned = ScratchOwned::new(
+            RLWECt::encrypt_sk_scratch_bytes(&module, ct.size()) | RLWECt::decrypt_scratch_bytes(&module, ct.size()),
+        );
 
         let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
         sk.fill_ternary_prob(0.5, &mut source_xs);
@@ -469,9 +619,8 @@ mod tests {
         let mut ct_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_ct, 2);
 
         let mut scratch: ScratchOwned = ScratchOwned::new(
-            encrypt_rlwe_sk_scratch_bytes(&module, ct_dft.size())
-                | decrypt_rlwe_dft_scratch_bytes(&module, ct_dft.size())
-                | encrypt_zero_rlwe_dft_scratch_bytes(&module, ct_dft.size()),
+            RLWECtDft::decrypt_scratch_bytes(&module, ct_dft.size())
+                | RLWECtDft::encrypt_zero_sk_scratch_bytes(&module, ct_dft.size()),
         );
 
         ct_dft.encrypt_zero_sk(
@@ -523,9 +672,9 @@ mod tests {
         );
 
         let mut scratch: ScratchOwned = ScratchOwned::new(
-            encrypt_rlwe_sk_scratch_bytes(&module, ct.size())
-                | decrypt_rlwe_scratch_bytes(&module, ct.size())
-                | encrypt_rlwe_pk_scratch_bytes(&module, pk.size()),
+            RLWECt::encrypt_sk_scratch_bytes(&module, ct.size())
+                | RLWECt::decrypt_scratch_bytes(&module, ct.size())
+                | RLWECt::encrypt_pk_scratch_bytes(&module, pk.size()),
         );
 
         let mut data_want: Vec<i64> = vec![0i64; module.n()];

rlwe/src/keys.rs

@@ -4,10 +4,7 @@ use base2k::{
 };
 use sampling::source::Source;
 
-use crate::{
-    elem::{Infos, RLWECtDft},
-    encryption::encrypt_zero_rlwe_dft_scratch_bytes,
-};
+use crate::{elem::Infos, elem_rlwe::RLWECtDft};
 
 #[derive(Clone, Copy, Debug)]
 pub enum SecretDistribution {
@@ -182,7 +179,10 @@ impl<C> PublicKey<C, FFT64> {
         }
 
         // Its ok to allocate scratch space here since pk is usually generated only once.
-        let mut scratch: ScratchOwned = ScratchOwned::new(encrypt_zero_rlwe_dft_scratch_bytes(module, self.size()));
+        let mut scratch: ScratchOwned = ScratchOwned::new(RLWECtDft::encrypt_zero_sk_scratch_bytes(
+            module,
+            self.size(),
+        ));
         self.data.encrypt_zero_sk(
             module,
             sk_dft,

rlwe/src/lib.rs

@@ -1,3 +1,6 @@
 pub mod elem;
-pub mod encryption;
+pub mod elem_grlwe;
+pub mod elem_rgsw;
+pub mod elem_rlwe;
 pub mod keys;
+mod utils;

rlwe/src/utils.rs

@@ -0,0 +1,3 @@
+pub(crate) fn derive_size(log_base2k: usize, log_k: usize) -> usize {
+    (log_k + log_base2k - 1) / log_base2k
+}