From 0b8bf98b2a83ebe9054dde4c4bef665e3062cd64 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Fri, 31 Jan 2025 15:39:08 +0100
Subject: [PATCH] added vmp_prepare_dblptr

---
 .gitmodules                         |  5 +----
 base2k/README.md                    | 11 +++++++++++
 base2k/spqlios-arithmetic           |  1 +
 base2k/src/ffi/vmp.rs               |  6 +++---
 base2k/src/scalar.rs                |  2 +-
 base2k/src/vector_matrix_product.rs | 24 ++++++++++++++++++++++--
 6 files changed, 39 insertions(+), 10 deletions(-)
 create mode 100644 base2k/README.md
 create mode 160000 base2k/spqlios-arithmetic

diff --git a/.gitmodules b/.gitmodules
index 0cf28dc..7ba4798 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,6 +1,3 @@
-[submodule "spqlios/spqlios-arithmetic"]
-	path = spqlios/spqlios-arithmetic
-	url = https://github.com/tfhe/spqlios-arithmetic
 [submodule "base2k/spqlios-arithmetic"]
 	path = base2k/spqlios-arithmetic
-	url = https://github.com/tfhe/spqlios-arithmetic
+	url = https://github.com/tfhe/spqlios-arithmetic
\ No newline at end of file
diff --git a/base2k/README.md b/base2k/README.md
new file mode 100644
index 0000000..bbb334d
--- /dev/null
+++ b/base2k/README.md
@@ -0,0 +1,11 @@
+# DISCLAIMER: ONLY TESTED ON UBUNTU
+
+To use this crate you need to build spqlios-arithmetic, which is provided a as a git submodule:
+1) Initialize the sub-modile
+2) $ cd base2k/spqlios-arithmetic
+3) mdkir build
+4) cd build
+5) cmake ..
+6) make
+
+Steps 3 to 6 might change depending of your platform. See [spqlios-arithmetic/wiki/build](https://github.com/tfhe/spqlios-arithmetic/wiki/build) for additional information and build options.
\ No newline at end of file
diff --git a/base2k/spqlios-arithmetic b/base2k/spqlios-arithmetic
new file mode 160000
index 0000000..7ea875f
--- /dev/null
+++ b/base2k/spqlios-arithmetic
@@ -0,0 +1 @@
+Subproject commit 7ea875f61c51c67687ec9c15ae0eae04cda961f8
diff --git a/base2k/src/ffi/vmp.rs b/base2k/src/ffi/vmp.rs
index e202e4c..e51e31c 100644
--- a/base2k/src/ffi/vmp.rs
+++ b/base2k/src/ffi/vmp.rs
@@ -68,7 +68,7 @@ unsafe extern "C" {
 }
 
 unsafe extern "C" {
-    pub fn vmp_prepare_contiguous(
+    pub unsafe fn vmp_prepare_contiguous(
         module: *const MODULE,
         pmat: *mut VMP_PMAT,
         mat: *const i64,
@@ -78,10 +78,10 @@ unsafe extern "C" {
     );
 }
 unsafe extern "C" {
-    pub fn vmp_prepare_dblptr(
+    pub unsafe fn vmp_prepare_dblptr(
         module: *const MODULE,
         pmat: *mut VMP_PMAT,
-        mat: *mut *const i64,
+        mat: *const *const i64,
         nrows: u64,
         ncols: u64,
         tmp_space: *mut u8,
diff --git a/base2k/src/scalar.rs b/base2k/src/scalar.rs
index 59b85f7..50ff246 100644
--- a/base2k/src/scalar.rs
+++ b/base2k/src/scalar.rs
@@ -22,7 +22,7 @@ impl Scalar {
     }
 
     pub fn from_buffer(&mut self, n: usize, buf: &[i64]) {
-        let size = Self::buffer_size(n);
+        let size: usize = Self::buffer_size(n);
         assert!(
             buf.len() >= size,
             "invalid buffer: buf.len()={} < self.buffer_size(n={})={}",
diff --git a/base2k/src/vector_matrix_product.rs b/base2k/src/vector_matrix_product.rs
index b97d721..a546a91 100644
--- a/base2k/src/vector_matrix_product.rs
+++ b/base2k/src/vector_matrix_product.rs
@@ -1,7 +1,7 @@
 use crate::ffi::vmp::{
     delete_vmp_pmat, new_vmp_pmat, vmp_apply_dft, vmp_apply_dft_tmp_bytes, vmp_apply_dft_to_dft,
     vmp_apply_dft_to_dft_tmp_bytes, vmp_pmat_t, vmp_prepare_contiguous,
-    vmp_prepare_contiguous_tmp_bytes,
+    vmp_prepare_contiguous_tmp_bytes, vmp_prepare_dblptr,
 };
 use crate::{Module, VecZnx, VecZnxDft};
 
@@ -110,6 +110,21 @@ impl Module {
         }
     }
 
+    pub fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<&Vec<i64>>, buf: &mut [u8]) {
+        let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
+
+        unsafe {
+            vmp_prepare_dblptr(
+                self.0,
+                b.data(),
+                ptrs.as_ptr(),
+                b.rows() as u64,
+                b.cols() as u64,
+                buf.as_mut_ptr(),
+            );
+        }
+    }
+
     pub fn vmp_apply_dft_tmp_bytes(
         &self,
         c_limbs: usize,
@@ -209,7 +224,7 @@ pub struct Matrix3D<T> {
     pub n: usize,
 }
 
-impl<T: Default + Clone> Matrix3D<T> {
+impl<T: Default + Clone + std::marker::Copy> Matrix3D<T> {
     pub fn new(rows: usize, cols: usize, n: usize) -> Self {
         let size = rows * cols * n;
         Self {
@@ -231,4 +246,9 @@ impl<T: Default + Clone> Matrix3D<T> {
         let idx: usize = col * (self.n * self.rows) + row * self.n;
         &mut self.data[idx..idx + self.n]
     }
+
+    pub fn set_col(&mut self, col: usize, a: &[T]) {
+        let idx: usize = col * (self.n * self.rows);
+        self.data[idx..idx + self.rows * self.n].copy_from_slice(a);
+    }
 }