Added support for arbitrary extended LUT

backend/src/vec_znx.rs

@@ -249,7 +249,6 @@ impl<D: AsRef<[u8]> + AsMut<[u8]>> VecZnx<D>{
     }
 }
 
-
 fn normalize<D: AsMut<[u8]> + AsRef<[u8]>>(basek: usize, a: &mut VecZnx<D>, a_col: usize, tmp_bytes: &mut [u8]) {
     let n: usize = a.n();
 

core/src/blind_rotation/ccgi.rs

@@ -1,10 +1,15 @@
 use std::time::Instant;
 
-use backend::{MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxDftOps, VecZnxOps, ZnxView, ZnxViewMut, ZnxZero, FFT64};
+use backend::{
+    FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxDftOps, VecZnxOps, ZnxView, ZnxViewMut,
+    ZnxZero,
+};
 use itertools::izip;
 
 use crate::{
-    blind_rotation::key::BlindRotationKeyCGGI, lwe::ciphertext::LWECiphertextToRef, FourierGLWESecret, GGSWCiphertext, GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret, ScratchCore
+    GGSWCiphertext, GLWECiphertext, GLWECiphertextToMut, Infos, LWECiphertext, ScratchCore,
+    blind_rotation::{key::BlindRotationKeyCGGI, lut::LookUpTable},
+    lwe::ciphertext::LWECiphertextToRef,
 };
 
 pub fn cggi_blind_rotate_scratch_space(
@@ -21,26 +26,24 @@ pub fn cggi_blind_rotate_scratch_space(
             | GLWECiphertext::external_product_inplace_scratch_space(module, basek, k_lut, k_brk, 1, rank))
 }
 
-pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
+pub fn cggi_blind_rotate<DataRes, DataIn>(
     module: &Module<FFT64>,
     res: &mut GLWECiphertext<DataRes>,
     lwe: &LWECiphertext<DataIn>,
-    lut: &GLWEPlaintext<DataLUT>,
+    lut: &LookUpTable,
     brk: &BlindRotationKeyCGGI<FFT64>,
     scratch: &mut Scratch,
 ) where
     DataRes: AsRef<[u8]> + AsMut<[u8]>,
     DataIn: AsRef<[u8]>,
-    DataLUT: AsRef<[u8]>,
 {
     let basek = res.basek();
 
     let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
     let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
     let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
-    let lut_ref: GLWECiphertext<&[u8]> = lut.to_ref();
 
-    let cols = out_mut.rank()+1;
+    let cols: usize = out_mut.rank() + 1;
 
     mod_switch_2n(module, &mut lwe_2n, &lwe_ref);
 
@@ -50,7 +53,7 @@ pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
     out_mut.data.zero();
 
     // Initialize out to X^{b} * LUT(X)
-    module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut_ref.data, 0);
+    module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
 
     let block_size: usize = brk.block_size();
 
@@ -68,14 +71,10 @@ pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
         brk.data.chunks_exact(block_size)
     )
     .for_each(|(ai, ski)| {
-
         out_mut.dft(module, &mut acc_dft);
         acc_add_dft.data.zero();
 
-        izip!(ai.iter(), ski.iter())
+        izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
-            .enumerate()
-            .for_each(|(i, (aii, skii))| {
-
             // vmp_res = DFT(acc) * BRK[i]
             module.vmp_apply(&mut vmp_res.data, &acc_dft.data, &skii.data, scratch5);
 
@@ -98,7 +97,6 @@ pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
         });
 
         acc_dft.idft(module, &mut out_mut, scratch5);
-
     });
     let duration: std::time::Duration = start.elapsed();
     println!("external products: {} us", duration.as_micros());

core/src/blind_rotation/key.rs

@@ -69,14 +69,17 @@ impl BlindRotationKeyCGGI<FFT64> {
         }
     }
 
+    #[allow(dead_code)]
     pub(crate) fn rows(&self) -> usize {
         self.data[0].rows()
     }
 
+    #[allow(dead_code)]
     pub(crate) fn k(&self) -> usize {
         self.data[0].k()
     }
 
+    #[allow(dead_code)]
     pub(crate) fn rank(&self) -> usize {
         self.data[0].rank()
     }

core/src/blind_rotation/lut.rs

@@ -0,0 +1,84 @@
+use backend::{FFT64, Module, ScratchOwned, VecZnx, VecZnxAlloc, VecZnxOps, ZnxInfos, ZnxViewMut, alloc_aligned};
+
+pub struct LookUpTable {
+    pub(crate) data: Vec<VecZnx<Vec<u8>>>,
+    pub(crate) basek: usize,
+    pub(crate) k: usize,
+}
+
+impl LookUpTable {
+    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, extend_factor: usize) -> Self {
+        let size: usize = k.div_ceil(basek);
+        let mut data: Vec<VecZnx<Vec<u8>>> = Vec::with_capacity(extend_factor);
+        (0..extend_factor).for_each(|_| {
+            data.push(module.new_vec_znx(1, size));
+        });
+        Self { data, basek, k }
+    }
+
+    pub fn set(&mut self, module: &Module<FFT64>, f: fn(i64) -> i64, message_modulus: usize) {
+        let basek: usize = self.basek;
+
+        // Get the number minimum limb to store the message modulus
+        let limbs: usize = message_modulus.div_ceil(1 << basek);
+
+        // Scaling factor
+        let scale: i64 = (1 << (basek * limbs - 1)).div_round(message_modulus) as i64;
+
+        // Updates function
+        let f_scaled = |x: i64| (f(x) % message_modulus as i64) * scale;
+
+        // If LUT size > module.n()
+        let domain_size: usize = self.data[0].n() * self.data.len();
+
+        let size: usize = self.k.div_ceil(self.basek);
+
+        // Equivalent to AUTO([f(0), f(1), ..., f(n-1)], -1)
+        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::new::<i64>(domain_size, 1, size);
+        {
+            let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
+
+            let start: usize = 0;
+            let end: usize = (domain_size).div_round(message_modulus);
+
+            let y: i64 = f_scaled(0);
+            (start..end).for_each(|i| {
+                lut_at[i] = y;
+            });
+
+            (1..message_modulus).for_each(|x| {
+                let start: usize = (x * domain_size).div_round(message_modulus);
+                let end: usize = ((x + 1) * domain_size).div_round(message_modulus);
+                let y: i64 = f_scaled(x as i64);
+                (start..end).for_each(|i| {
+                    lut_at[domain_size - i] = -y;
+                })
+            });
+        }
+
+        // Rotates half the step to the left
+        let half_step: usize = domain_size.div_round(message_modulus << 1);
+        module.vec_znx_rotate_inplace(-(half_step as i64), &mut lut_full, 0);
+
+        let mut tmp_bytes: Vec<u8> = alloc_aligned(lut_full.n() * size_of::<i64>());
+        lut_full.normalize(self.basek, 0, &mut tmp_bytes);
+
+        if self.data.len() > 1 {
+            let mut scratch: ScratchOwned = ScratchOwned::new(module.bytes_of_vec_znx(1, size));
+            module.vec_znx_split(&mut self.data, 0, &lut_full, 0, scratch.borrow());
+        } else {
+            module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
+        }
+    }
+}
+
+pub trait DivRound {
+    fn div_round(self, rhs: Self) -> Self;
+}
+
+impl DivRound for usize {
+    #[inline]
+    fn div_round(self, rhs: Self) -> Self {
+        (self + rhs / 2) / rhs
+    }
+}

core/src/blind_rotation/mod.rs

@@ -1,6 +1,7 @@
 // pub mod cggi;
 pub mod ccgi;
 pub mod key;
+pub mod lut;
 
 #[cfg(test)]
 pub mod test_fft64;

core/src/blind_rotation/test_fft64/cggi.rs

@@ -1,10 +1,16 @@
 use std::time::Instant;
 
-use backend::{Encoding, Module, ScratchOwned, Stats, ZnxView, ZnxViewMut, FFT64};
+use backend::{Encoding, FFT64, Module, ScratchOwned, Stats, VecZnxOps, ZnxView};
 use sampling::source::Source;
 
 use crate::{
-    blind_rotation::{ccgi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, mod_switch_2n}, key::BlindRotationKeyCGGI}, lwe::{ciphertext::{LWECiphertextToMut, LWECiphertextToRef}, LWEPlaintext}, FourierGLWESecret, GLWECiphertext, GLWEOps, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret
+    FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret,
+    blind_rotation::{
+        ccgi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, mod_switch_2n},
+        key::BlindRotationKeyCGGI,
+        lut::LookUpTable,
+    },
+    lwe::{LWEPlaintext, ciphertext::LWECiphertextToRef},
 };
 
 #[test]
@@ -21,6 +27,8 @@ fn blind_rotation() {
     let rank: usize = 1;
     let block_size: usize = 7;
 
+    let message_modulus: usize = 64;
+
     let mut source_xs: Source = Source::new([0u8; 32]);
     let mut source_xe: Source = Source::new([0u8; 32]);
     let mut source_xa: Source = Source::new([0u8; 32]);
@@ -32,12 +40,14 @@ fn blind_rotation() {
     let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
     sk_lwe.fill_binary_block(block_size, &mut source_xs);
 
-    let mut scratch: ScratchOwned = ScratchOwned::new(BlindRotationKeyCGGI::generate_from_sk_scratch_space(
+    let mut scratch: ScratchOwned = ScratchOwned::new(
-        &module, basek, k_brk, rank,
+        BlindRotationKeyCGGI::generate_from_sk_scratch_space(&module, basek, k_brk, rank)
-    ) | cggi_blind_rotate_scratch_space(&module, basek, k_lut, k_brk, rows_brk, rank));
+            | cggi_blind_rotate_scratch_space(&module, basek, k_lut, k_brk, rows_brk, rank),
+    );
 
     let start: Instant = Instant::now();
     let mut brk: BlindRotationKeyCGGI<FFT64> = BlindRotationKeyCGGI::allocate(&module, n_lwe, basek, k_brk, rows_brk, rank);
+
     brk.generate_from_sk(
         &module,
         &sk_glwe_dft,
@@ -47,6 +57,7 @@ fn blind_rotation() {
         3.2,
         scratch.borrow(),
     );
+
     let duration: std::time::Duration = start.elapsed();
     println!("brk-gen: {} ms", duration.as_millis());
 
@@ -67,21 +78,17 @@ fn blind_rotation() {
 
     println!("{}", pt_lwe.data);
 
-    let mut lut: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_lut);
+    fn lut_fn(x: i64) -> i64 {
-
+        2 * x + 1
-    
+    }
-    lut.data.at_mut(0, 0)[0] = 0;
-    (1..module.n()).for_each(|i|{
-        lut.data.at_mut(0, 0)[i] = - ((module.n() as i64 - i as i64  - 1)<<(basek - module.log_n() - 1));
-    });
-
-    
 
+    let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, 1);
+    lut.set(&module, lut_fn, message_modulus);
 
     let mut res: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_lut, rank);
 
     let start: Instant = Instant::now();
-    (0..1).for_each(|i|{
+    (0..1).for_each(|_| {
         cggi_blind_rotate(&module, &mut res, &lwe, &lut, &brk, scratch.borrow());
     });
 
@@ -92,19 +99,27 @@ fn blind_rotation() {
 
     res.decrypt(&module, &mut pt_have, &sk_glwe_dft, scratch.borrow());
 
+    println!("pt_have: {}", pt_have.data);
+
     let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
 
     mod_switch_2n(&module, &mut lwe_2n, &lwe.to_ref());
 
-    let pt_want: i64 = (lwe_2n[0] + lwe_2n[1..].iter().zip(sk_lwe.data.at(0, 0)).map(|(x, y)| x * y).sum::<i64>()) % (module.n() as i64 * 2);
+    let pt_want: i64 = (lwe_2n[0]
+        + lwe_2n[1..]
+            .iter()
+            .zip(sk_lwe.data.at(0, 0))
+            .map(|(x, y)| x * y)
+            .sum::<i64>())
+        % (module.n() as i64 * 2);
 
-    lut.rotate_inplace(&module, pt_want);
+    module.vec_znx_rotate_inplace(pt_want, &mut lut.data[0], 0);
 
-    lut.sub_inplace_ab(&module, &pt_have);
+    println!("pt_want: {}", lut.data[0]);
 
-    let noise: f64 = lut.data.std(0, basek);
+    module.vec_znx_sub_ab_inplace(&mut lut.data[0], 0, &pt_have.data, 0);
+
+    let noise: f64 = lut.data[0].std(0, basek);
 
     println!("noise: {}", noise);
-
-    
 }

core/src/glwe/test_fft64/automorphism.rs

@@ -3,8 +3,7 @@ use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats, VecZnxOps};
 use sampling::source::Source;
 
 use crate::{
-    FourierGLWESecret, GLWEAutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
+    FourierGLWESecret, GLWEAutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, noise::log2_std_noise_gglwe_product,
-    noise::log2_std_noise_gglwe_product,
 };
 
 #[test]

core/src/glwe/test_fft64/external_product.rs

@@ -1,9 +1,7 @@
 use backend::{FFT64, FillUniform, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, ZnxViewMut};
 use sampling::source::Source;
 
-use crate::{
+use crate::{FourierGLWESecret, GGSWCiphertext, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, noise::noise_ggsw_product};
-    FourierGLWESecret, GGSWCiphertext, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, noise::noise_ggsw_product,
-};
 
 #[test]
 fn apply() {

core/src/glwe/test_fft64/keyswitch.rs

@@ -2,8 +2,7 @@ use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats, VecZnxOps};
 use sampling::source::Source;
 
 use crate::{
-    FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, GLWESwitchingKey, Infos,
+    FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, GLWESwitchingKey, Infos, noise::log2_std_noise_gglwe_product,
-    noise::log2_std_noise_gglwe_product,
 };
 
 #[test]

core/src/glwe/test_fft64/trace.rs

@@ -4,8 +4,7 @@ use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats, VecZnxOps, ZnxVie
 use sampling::source::Source;
 
 use crate::{
-    FourierGLWESecret, GLWEAutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
+    FourierGLWESecret, GLWEAutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, noise::var_noise_gglwe_product,
-    noise::var_noise_gglwe_product,
 };
 
 #[test]

core/src/lwe/decryption.rs

@@ -1,15 +1,28 @@
-use backend::{alloc_aligned, ZnxView, ZnxViewMut};
+use backend::{ZnxView, ZnxViewMut, alloc_aligned};
 
-use crate::{lwe::{LWEPlaintext}, Infos, LWECiphertext, LWESecret, SetMetaData};
+use crate::{Infos, LWECiphertext, LWESecret, SetMetaData, lwe::LWEPlaintext};
 
-impl<DataSelf> LWECiphertext<DataSelf> where DataSelf: AsRef<[u8]>{
+impl<DataSelf> LWECiphertext<DataSelf>
-    pub fn decrypt<DataPt, DataSk>(&self, pt: &mut LWEPlaintext<DataPt>, sk: &LWESecret<DataSk>) where DataPt: AsRef<[u8]> + AsMut<[u8]>, DataSk: AsRef<[u8]>{
+where
-        #[cfg(debug_assertions)]{
+    DataSelf: AsRef<[u8]>,
+{
+    pub fn decrypt<DataPt, DataSk>(&self, pt: &mut LWEPlaintext<DataPt>, sk: &LWESecret<DataSk>)
+    where
+        DataPt: AsRef<[u8]> + AsMut<[u8]>,
+        DataSk: AsRef<[u8]>,
+    {
+        #[cfg(debug_assertions)]
+        {
             assert_eq!(self.n(), sk.n());
         }
 
         (0..pt.size().min(self.size())).for_each(|i| {
-            pt.data.at_mut(0, i)[0] = self.data.at(0, i)[0] + self.data.at(0, i)[1..].iter().zip(sk.data.at(0, 0)).map(|(x, y)| x * y).sum::<i64>();
+            pt.data.at_mut(0, i)[0] = self.data.at(0, i)[0]
+                + self.data.at(0, i)[1..]
+                    .iter()
+                    .zip(sk.data.at(0, 0))
+                    .map(|(x, y)| x * y)
+                    .sum::<i64>();
         });
 
         let mut tmp_bytes: Vec<u8> = alloc_aligned(size_of::<i64>());

core/src/lwe/encryption.rs

@@ -1,14 +1,25 @@
-use backend::{alloc_aligned, AddNormal, FillUniform, VecZnx, ZnxView, ZnxViewMut};
+use backend::{AddNormal, FillUniform, VecZnx, ZnxView, ZnxViewMut, alloc_aligned};
 use sampling::source::Source;
 
-use crate::{lwe::LWEPlaintext, Infos, LWECiphertext, LWESecret, SIX_SIGMA};
+use crate::{Infos, LWECiphertext, LWESecret, SIX_SIGMA, lwe::LWEPlaintext};
 
-
+impl<DataSelf> LWECiphertext<DataSelf>
-
+where
-impl<DataSelf> LWECiphertext<DataSelf> where DataSelf: AsMut<[u8]> + AsRef<[u8]>{
+    DataSelf: AsMut<[u8]> + AsRef<[u8]>,
-    pub fn encrypt_sk<DataPt, DataSk>(&mut self, pt: &LWEPlaintext<DataPt>, sk: &LWESecret<DataSk>, source_xa: &mut Source, source_xe: &mut Source, sigma: f64) where DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>{
+{
-
+    pub fn encrypt_sk<DataPt, DataSk>(
-        #[cfg(debug_assertions)]{
+        &mut self,
+        pt: &LWEPlaintext<DataPt>,
+        sk: &LWESecret<DataSk>,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        sigma: f64,
+    ) where
+        DataPt: AsRef<[u8]>,
+        DataSk: AsRef<[u8]>,
+    {
+        #[cfg(debug_assertions)]
+        {
             assert_eq!(self.n(), sk.n())
         }
 
@@ -18,7 +29,12 @@ impl<DataSelf> LWECiphertext<DataSelf> where DataSelf: AsMut<[u8]> + AsRef<[u8]>
         let mut tmp_znx: VecZnx<Vec<u8>> = VecZnx::<Vec<u8>>::new::<i64>(1, 1, self.size());
 
         (0..self.size()).for_each(|i| {
-            tmp_znx.at_mut(0, i)[0] = pt.data.at(0, i)[0] - self.data.at(0, i)[1..].iter().zip(sk.data.at(0, 0)).map(|(x, y)| x * y).sum::<i64>();
+            tmp_znx.at_mut(0, i)[0] = pt.data.at(0, i)[0]
+                - self.data.at(0, i)[1..]
+                    .iter()
+                    .zip(sk.data.at(0, 0))
+                    .map(|(x, y)| x * y)
+                    .sum::<i64>();
         });
 
         tmp_znx.add_normal(basek, 0, self.k(), source_xe, sigma, sigma * SIX_SIGMA);
@@ -30,6 +46,5 @@ impl<DataSelf> LWECiphertext<DataSelf> where DataSelf: AsMut<[u8]> + AsRef<[u8]>
         (0..self.size()).for_each(|i| {
             self.data.at_mut(0, i)[0] = tmp_znx.at(0, i)[0];
         });
-
     }
 }

core/src/lwe/mod.rs

@@ -1,9 +1,9 @@
 pub mod ciphertext;
-pub mod secret;
-pub mod encryption;
 pub mod decryption;
+pub mod encryption;
 pub mod plaintext;
+pub mod secret;
 
 pub use ciphertext::LWECiphertext;
-pub use secret::LWESecret;
 pub use plaintext::LWEPlaintext;
+pub use secret::LWESecret;