fixed potential buffer overflow & added vec_znx_normalize_tmp_bytes

base2k/src/vec_znx.rs

@@ -469,6 +469,7 @@ pub fn rsh<T: VecZnxCommon>(log_base2k: usize, a: &mut T, k: usize, carry: &mut
 pub trait VecZnxCommon: VecZnxApi + Infos {}
 
 pub trait VecZnxOps {
+
     /// Allocates a new [VecZnx].
     ///
     /// # Arguments
@@ -480,6 +481,8 @@ pub trait VecZnxOps {
     /// a new [VecZnx] through [VecZnx::from_bytes].
     fn bytes_of_vec_znx(&self, cols: usize) -> usize;
 
+    fn vec_znx_normalize_tmp_bytes(&self) -> usize;
+
     /// c <- a + b.
     fn vec_znx_add<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
         &self,
@@ -557,6 +560,12 @@ impl VecZnxOps for Module {
         self.n() * cols * 8
     }
 
+    fn vec_znx_normalize_tmp_bytes(&self) -> usize{
+        unsafe{
+            vec_znx::vec_znx_normalize_tmp_bytes(self.0) as usize
+        }   
+    }
+
     // c <- a + b
     fn vec_znx_add<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
         &self,

base2k/src/vmp.rs

@@ -439,6 +439,7 @@ impl VmpPMatOps for Module {
     }
 
     fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, buf: &mut [u8]) {
+        debug_assert!(a.len() == b.cols() * self.n());
         unsafe {
             vmp::vmp_prepare_row(
                 self.0,