From cac4b3549d0905fb79fb5012d338aad0c106ae7b Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Mon, 24 Feb 2025 15:29:52 +0100
Subject: [PATCH] added generic copy_from for VecZnxApi

---
 base2k/examples/vector_matrix_product.rs |  4 +-
 base2k/src/stats.rs                      |  8 +--
 base2k/src/vec_znx.rs                    | 76 +++++++++++++++++++++---
 3 files changed, 73 insertions(+), 15 deletions(-)

diff --git a/base2k/examples/vector_matrix_product.rs b/base2k/examples/vector_matrix_product.rs
index 11a5783..0eae265 100644
--- a/base2k/examples/vector_matrix_product.rs
+++ b/base2k/examples/vector_matrix_product.rs
@@ -1,6 +1,6 @@
 use base2k::{
     Encoding, Free, Infos, Module, VecZnx, VecZnxApi, VecZnxBig, VecZnxBigOps, VecZnxDft,
-    VecZnxDftOps, VecZnxOps, VmpPMat, VmpPMatOps, FFT64,
+    VecZnxDftOps, VecZnxOps, VecZnxVec, VmpPMat, VmpPMatOps, FFT64,
 };
 
 fn main() {
@@ -40,7 +40,7 @@ fn main() {
         vecznx[i].data[i * n + 1] = 1 as i64;
     });
 
-    let slices: Vec<&[i64]> = vecznx.iter().map(|v| v.data.as_slice()).collect();
+    let slices: Vec<&[i64]> = vecznx.dblptr();
 
     let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
     module.vmp_prepare_dblptr(&mut vmp_pmat, &slices, &mut buf);
diff --git a/base2k/src/stats.rs b/base2k/src/stats.rs
index 60568b3..20ef60a 100644
--- a/base2k/src/stats.rs
+++ b/base2k/src/stats.rs
@@ -1,6 +1,6 @@
-use crate::{Infos, VecZnx, Encoding};
+use crate::{Encoding, Infos, VecZnx};
 use rug::float::Round;
-use rug::ops::{AddAssignRound, SubAssignRound, DivAssignRound};
+use rug::ops::{AddAssignRound, DivAssignRound, SubAssignRound};
 use rug::Float;
 
 impl VecZnx {
@@ -18,9 +18,7 @@ impl VecZnx {
             x.sub_assign_round(&avg, Round::Nearest);
         });
         let mut std: Float = Float::with_val(prec, 0);
-        data.iter().for_each(|x| {
-            std += x*x
-        });
+        data.iter().for_each(|x| std += x * x);
         std.div_assign_round(Float::with_val(prec, data.len()), Round::Nearest);
         std = std.sqrt();
         std.to_f64()
diff --git a/base2k/src/vec_znx.rs b/base2k/src/vec_znx.rs
index 01118be..13c0bca 100644
--- a/base2k/src/vec_znx.rs
+++ b/base2k/src/vec_znx.rs
@@ -7,7 +7,22 @@ use crate::{Infos, Module};
 use itertools::izip;
 use std::cmp::min;
 
-pub trait VecZnxApi {
+pub trait VecZnxVec {
+    fn dblptr(&self) -> Vec<&[i64]>;
+    fn dblptr_mut(&mut self) -> Vec<&mut [i64]>;
+}
+
+impl<T: VecZnxApi + Infos> VecZnxVec for Vec<T> {
+    fn dblptr(&self) -> Vec<&[i64]> {
+        self.iter().map(|v| v.raw()).collect()
+    }
+
+    fn dblptr_mut(&mut self) -> Vec<&mut [i64]> {
+        self.iter_mut().map(|v| v.raw_mut()).collect()
+    }
+}
+
+pub trait VecZnxApi: AsRef<Self> + AsMut<Self> {
     type Owned: VecZnxApi + Infos;
 
     fn from_bytes(n: usize, cols: usize, bytes: &mut [u8]) -> Self::Owned;
@@ -16,6 +31,11 @@ pub trait VecZnxApi {
     /// new backend array.
     fn bytes_of(n: usize, cols: usize) -> usize;
 
+    /// Copy the data of a onto self.
+    fn copy_from<T: VecZnxApi + Infos>(&mut self, a: &T)
+    where
+        Self: AsMut<T>;
+
     /// Returns the backing array.
     fn raw(&self) -> &[i64];
 
@@ -83,6 +103,18 @@ pub struct VecZnxBorrow {
     pub data: *mut i64,
 }
 
+impl AsMut<VecZnxBorrow> for VecZnxBorrow {
+    fn as_mut(&mut self) -> &mut VecZnxBorrow {
+        self
+    }
+}
+
+impl AsRef<VecZnxBorrow> for VecZnxBorrow {
+    fn as_ref(&self) -> &VecZnxBorrow {
+        self
+    }
+}
+
 impl VecZnxApi for VecZnxBorrow {
     type Owned = VecZnxBorrow;
 
@@ -113,6 +145,13 @@ impl VecZnxApi for VecZnxBorrow {
         bytes_of_vec_znx(n, cols)
     }
 
+    fn copy_from<T: VecZnxApi + Infos>(&mut self, a: &T)
+    where
+        Self: AsMut<T>,
+    {
+        copy_vec_znx_from::<T>(self.as_mut(), a);
+    }
+
     fn as_ptr(&self) -> *const i64 {
         self.data
     }
@@ -203,6 +242,13 @@ impl VecZnxApi for VecZnx {
         bytes_of_vec_znx(n, cols)
     }
 
+    fn copy_from<T: VecZnxApi + Infos>(&mut self, a: &T)
+    where
+        Self: AsMut<T>,
+    {
+        copy_vec_znx_from(self.as_mut(), a);
+    }
+
     fn raw(&self) -> &[i64] {
         &self.data
     }
@@ -272,6 +318,27 @@ pub struct VecZnx {
     pub data: Vec<i64>,
 }
 
+impl AsMut<VecZnx> for VecZnx {
+    fn as_mut(&mut self) -> &mut VecZnx {
+        self
+    }
+}
+
+impl AsRef<VecZnx> for VecZnx {
+    fn as_ref(&self) -> &VecZnx {
+        self
+    }
+}
+
+/// Copies the coefficients of `a` on the receiver.
+/// Copy is done with the minimum size matching both backing arrays.
+pub fn copy_vec_znx_from<T: VecZnxApi + Infos>(b: &mut T, a: &T) {
+    let data_a: &[i64] = a.raw();
+    let data_b: &mut [i64] = b.raw_mut();
+    let size = min(data_b.len(), data_a.len());
+    data_b[..size].copy_from_slice(&data_a[..size])
+}
+
 impl VecZnx {
     /// Allocates a new [VecZnx] composed of #cols polynomials of Z\[X\].
     pub fn new(n: usize, cols: usize) -> Self {
@@ -281,13 +348,6 @@ impl VecZnx {
         }
     }
 
-    /// Copies the coefficients of `a` on the receiver.
-    /// Copy is done with the minimum size matching both backing arrays.
-    pub fn copy_from(&mut self, a: &VecZnx) {
-        let size = min(self.data.len(), a.data.len());
-        self.data[..size].copy_from_slice(&a.data[..size])
-    }
-
     /// Truncates the precision of the [VecZnx] by k bits.
     ///
     /// # Arguments