Added automorphism for glwe

base2k/src/scalar_znx.rs

@@ -1,3 +1,4 @@
+use crate::ffi::vec_znx;
 use crate::znx_base::ZnxInfos;
 use crate::{
     Backend, DataView, DataViewMut, Module, VecZnx, VecZnxToMut, VecZnxToRef, ZnxSliceSize, ZnxView, ZnxViewMut, alloc_aligned,
@@ -122,6 +123,69 @@ impl<B: Backend> ScalarZnxAlloc for Module<B> {
     }
 }
 
+pub trait ScalarZnxOps {
+    fn scalar_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
+    where
+        R: ScalarZnxToMut,
+        A: ScalarZnxToRef;
+
+    /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
+    fn vec_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
+    where
+        A: ScalarZnxToMut;
+}
+
+impl<B: Backend> ScalarZnxOps for Module<B> {
+    fn scalar_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
+    where
+        R: ScalarZnxToMut,
+        A: ScalarZnxToRef,
+    {
+        let a: ScalarZnx<&[u8]> = a.to_ref();
+        let mut res: ScalarZnx<&mut [u8]> = res.to_mut();
+        #[cfg(debug_assertions)]
+        {
+            assert_eq!(a.n(), self.n());
+            assert_eq!(res.n(), self.n());
+        }
+        unsafe {
+            vec_znx::vec_znx_automorphism(
+                self.ptr,
+                k,
+                res.at_mut_ptr(res_col, 0),
+                res.size() as u64,
+                res.sl() as u64,
+                a.at_ptr(a_col, 0),
+                a.size() as u64,
+                a.sl() as u64,
+            )
+        }
+    }
+
+    fn vec_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
+    where
+        A: ScalarZnxToMut,
+    {
+        let mut a: ScalarZnx<&mut [u8]> = a.to_mut();
+        #[cfg(debug_assertions)]
+        {
+            assert_eq!(a.n(), self.n());
+        }
+        unsafe {
+            vec_znx::vec_znx_automorphism(
+                self.ptr,
+                k,
+                a.at_mut_ptr(a_col, 0),
+                a.size() as u64,
+                a.sl() as u64,
+                a.at_ptr(a_col, 0),
+                a.size() as u64,
+                a.sl() as u64,
+            )
+        }
+    }
+}
+
 impl<D> ScalarZnx<D> {
     pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
         Self { data, n, cols }