refactor of key-switching & external product

base2k/src/module.rs

@@ -1,107 +1,107 @@
-use crate::GALOISGENERATOR;
-use crate::ffi::module::{MODULE, delete_module_info, module_info_t, new_module_info};
-use std::marker::PhantomData;
-
-#[derive(Copy, Clone)]
-#[repr(u8)]
-pub enum BACKEND {
-    FFT64,
-    NTT120,
-}
-
-pub trait Backend {
-    const KIND: BACKEND;
-    fn module_type() -> u32;
-}
-
-pub struct FFT64;
-pub struct NTT120;
-
-impl Backend for FFT64 {
-    const KIND: BACKEND = BACKEND::FFT64;
-    fn module_type() -> u32 {
-        0
-    }
-}
-
-impl Backend for NTT120 {
-    const KIND: BACKEND = BACKEND::NTT120;
-    fn module_type() -> u32 {
-        1
-    }
-}
-
-pub struct Module<B: Backend> {
-    pub ptr: *mut MODULE,
-    n: usize,
-    _marker: PhantomData<B>,
-}
-
-impl<B: Backend> Module<B> {
-    // Instantiates a new module.
-    pub fn new(n: usize) -> Self {
-        unsafe {
-            let m: *mut module_info_t = new_module_info(n as u64, B::module_type());
-            if m.is_null() {
-                panic!("Failed to create module.");
-            }
-            Self {
-                ptr: m,
-                n: n,
-                _marker: PhantomData,
-            }
-        }
-    }
-
-    pub fn n(&self) -> usize {
-        self.n
-    }
-
-    pub fn log_n(&self) -> usize {
-        (usize::BITS - (self.n() - 1).leading_zeros()) as _
-    }
-
-    pub fn cyclotomic_order(&self) -> u64 {
-        (self.n() << 1) as _
-    }
-
-    // Returns GALOISGENERATOR^|generator| * sign(generator)
-    pub fn galois_element(&self, generator: i64) -> i64 {
-        if generator == 0 {
-            return 1;
-        }
-        ((mod_exp_u64(GALOISGENERATOR, generator.abs() as usize) & (self.cyclotomic_order() - 1)) as i64) * generator.signum()
-    }
-
-    // Returns gen^-1
-    pub fn galois_element_inv(&self, generator: i64) -> i64 {
-        if generator == 0 {
-            panic!("cannot invert 0")
-        }
-        ((mod_exp_u64(
-            generator.abs() as u64,
-            (self.cyclotomic_order() - 1) as usize,
-        ) & (self.cyclotomic_order() - 1)) as i64)
-            * generator.signum()
-    }
-}
-
-impl<B: Backend> Drop for Module<B> {
-    fn drop(&mut self) {
-         unsafe { delete_module_info(self.ptr) }
-    }
-}
-
-fn mod_exp_u64(x: u64, e: usize) -> u64 {
-    let mut y: u64 = 1;
-    let mut x_pow: u64 = x;
-    let mut exp = e;
-    while exp > 0 {
-        if exp & 1 == 1 {
-            y = y.wrapping_mul(x_pow);
-        }
-        x_pow = x_pow.wrapping_mul(x_pow);
-        exp >>= 1;
-    }
-    y
-}
+use crate::GALOISGENERATOR;
+use crate::ffi::module::{MODULE, delete_module_info, module_info_t, new_module_info};
+use std::marker::PhantomData;
+
+#[derive(Copy, Clone)]
+#[repr(u8)]
+pub enum BACKEND {
+    FFT64,
+    NTT120,
+}
+
+pub trait Backend {
+    const KIND: BACKEND;
+    fn module_type() -> u32;
+}
+
+pub struct FFT64;
+pub struct NTT120;
+
+impl Backend for FFT64 {
+    const KIND: BACKEND = BACKEND::FFT64;
+    fn module_type() -> u32 {
+        0
+    }
+}
+
+impl Backend for NTT120 {
+    const KIND: BACKEND = BACKEND::NTT120;
+    fn module_type() -> u32 {
+        1
+    }
+}
+
+pub struct Module<B: Backend> {
+    pub ptr: *mut MODULE,
+    n: usize,
+    _marker: PhantomData<B>,
+}
+
+impl<B: Backend> Module<B> {
+    // Instantiates a new module.
+    pub fn new(n: usize) -> Self {
+        unsafe {
+            let m: *mut module_info_t = new_module_info(n as u64, B::module_type());
+            if m.is_null() {
+                panic!("Failed to create module.");
+            }
+            Self {
+                ptr: m,
+                n: n,
+                _marker: PhantomData,
+            }
+        }
+    }
+
+    pub fn n(&self) -> usize {
+        self.n
+    }
+
+    pub fn log_n(&self) -> usize {
+        (usize::BITS - (self.n() - 1).leading_zeros()) as _
+    }
+
+    pub fn cyclotomic_order(&self) -> u64 {
+        (self.n() << 1) as _
+    }
+
+    // Returns GALOISGENERATOR^|generator| * sign(generator)
+    pub fn galois_element(&self, generator: i64) -> i64 {
+        if generator == 0 {
+            return 1;
+        }
+        ((mod_exp_u64(GALOISGENERATOR, generator.abs() as usize) & (self.cyclotomic_order() - 1)) as i64) * generator.signum()
+    }
+
+    // Returns gen^-1
+    pub fn galois_element_inv(&self, generator: i64) -> i64 {
+        if generator == 0 {
+            panic!("cannot invert 0")
+        }
+        ((mod_exp_u64(
+            generator.abs() as u64,
+            (self.cyclotomic_order() - 1) as usize,
+        ) & (self.cyclotomic_order() - 1)) as i64)
+            * generator.signum()
+    }
+}
+
+impl<B: Backend> Drop for Module<B> {
+    fn drop(&mut self) {
+        unsafe { delete_module_info(self.ptr) }
+    }
+}
+
+fn mod_exp_u64(x: u64, e: usize) -> u64 {
+    let mut y: u64 = 1;
+    let mut x_pow: u64 = x;
+    let mut exp = e;
+    while exp > 0 {
+        if exp & 1 == 1 {
+            y = y.wrapping_mul(x_pow);
+        }
+        x_pow = x_pow.wrapping_mul(x_pow);
+        exp >>= 1;
+    }
+    y
+}

base2k/src/vec_znx_dft_ops.rs

@@ -42,8 +42,13 @@ pub trait VecZnxDftOps<B: Backend> {
     /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
     fn vec_znx_idft_tmp_bytes(&self) -> usize;
 
+    fn vec_znx_dft_copy<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
+    where
+        R: VecZnxDftToMut<B>,
+        A: VecZnxDftToRef<B>;
+
     /// b <- IDFT(a), uses a as scratch space.
-    fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_cols: usize)
+    fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
     where
         R: VecZnxBigToMut<B>,
         A: VecZnxDftToMut<B>;
@@ -79,13 +84,33 @@ impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
 }
 
 impl VecZnxDftOps<FFT64> for Module<FFT64> {
+    fn vec_znx_dft_copy<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
+    where
+        R: VecZnxDftToMut<FFT64>,
+        A: VecZnxDftToRef<FFT64>,
+    {
+        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
+        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
+
+        let min_size: usize = min(res_mut.size(), a_ref.size());
+
+        (0..min_size).for_each(|j| {
+            res_mut
+                .at_mut(res_col, j)
+                .copy_from_slice(a_ref.at(a_col, j));
+        });
+        (min_size..res_mut.size()).for_each(|j| {
+            res_mut.zero_at(res_col, j);
+        })
+    }
+
     fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
     where
         R: VecZnxBigToMut<FFT64>,
         A: VecZnxDftToMut<FFT64>,
     {
-        let mut res_mut = res.to_mut();
-        let mut a_mut = a.to_mut();
+        let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
+        let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
 
         let min_size: usize = min(res_mut.size(), a_mut.size());
 
@@ -136,14 +161,14 @@ impl VecZnxDftOps<FFT64> for Module<FFT64> {
     /// b <- DFT(a)
     ///
     /// # Panics
-    /// If b.cols < a_cols
+    /// If b.cols < a_col
     fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
     where
         R: VecZnxDftToMut<FFT64>,
         A: VecZnxToRef,
     {
-        let mut res_mut = res.to_mut();
-        let a_ref = a.to_ref();
+        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
+        let a_ref: crate::VecZnx<&[u8]> = a.to_ref();
 
         let min_size: usize = min(res_mut.size(), a_ref.size());
 
@@ -170,8 +195,8 @@ impl VecZnxDftOps<FFT64> for Module<FFT64> {
         R: VecZnxBigToMut<FFT64>,
         A: VecZnxDftToRef<FFT64>,
     {
-        let mut res_mut = res.to_mut();
-        let a_ref = a.to_ref();
+        let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
+        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
 
         let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_idft_tmp_bytes());