Merge branch 'jay-ckks' of https://github.com/Janmajayamall/gauss int…

…o jay-ckks
gausslabs · Apr 19, 2024 · 3890616 · 3890616
2 parents 56898a8 + 4c21e48
commit 3890616
Show file tree

Hide file tree

Showing 8 changed files with 135 additions and 50 deletions.
diff --git a/src/ciphertext.rs b/src/ciphertext.rs
@@ -1,7 +1,7 @@
 use std::{cell::RefCell, fmt::Display, sync::Arc};
 
 use crate::core_crypto::{
-    matrix::{Matrix, MatrixMut, Row, RowMut},
+    matrix::{Matrix, MatrixEntity, MatrixMut, Row, RowMut},
     modulus::NativeModulusBackend,
     ntt::NativeNTTBackend,
     num::UnsignedInteger,
@@ -17,7 +17,8 @@ pub trait Ciphertext {
     type Scalar: UnsignedInteger;
     type Row: Row<Element = Self::Scalar> + RowMut<Element = Self::Scalar>;
     type Poly: Matrix<MatElement = Self::Scalar>
-        + MatrixMut<MatElement = Self::Scalar, R = Self::Row>;
+        + MatrixMut<MatElement = Self::Scalar, R = Self::Row>
+        + MatrixEntity;
 
     fn representation(&self) -> Representation;
     fn representation_mut(&mut self) -> &mut Representation;

diff --git a/src/core_crypto/matrix.rs b/src/core_crypto/matrix.rs
@@ -2,17 +2,16 @@ use std::fmt::Debug;
 
 use aligned_vec::{avec, AVec};
 use itertools::Itertools;
+use num_traits::Zero;
 
 use crate::core_crypto::num::UnsignedInteger;
 
-pub trait Matrix: AsRef<[Self::R]> + Debug {
+pub trait Matrix: AsRef<[Self::R]> {
     type MatElement;
     type R: Row<Element = Self::MatElement>;
 
     fn dimension(&self) -> (usize, usize);
 
-    fn zeros(row: usize, col: usize) -> Self;
-
     fn get_row(&self, row_idx: usize) -> &Self::R {
         &self.as_ref()[row_idx]
     }
@@ -43,6 +42,16 @@ where
         self.as_mut()[row_idx].as_mut()[column_idx] = val;
     }
 
+    /*  Ideally, you want the return type to be MatrixMut, However this is not doable because of object safety
+    However, when R is a vector or align vector, the result will always be a matrixmut because we explicitly implemented MatrixMut
+    for them.
+
+    Note the postion idx is included in the second half
+    */
+    fn split_at_row(&mut self, idx: usize) -> (&mut [<Self as Matrix>::R], &mut [<Self as Matrix>::R]) {
+        self.as_mut().split_at_mut(idx)
+    }
+
     fn get_col_iter_mut(
         &mut self,
         column_idx: usize,
@@ -67,6 +76,13 @@ where
     }
 }
 
+pub trait MatrixEntity: Matrix
+where
+    <Self as Matrix>::MatElement: Zero,
+{
+    fn zeros(row: usize, col: usize) -> Self;
+}
+
 pub trait Drop2Dimension: Matrix {
     /// Drop 2 dimensionality of matrix and reduces it self to a single row
     /// vector
@@ -128,40 +144,80 @@ impl<T> IntoRowOwned for aligned_vec::ABox<[T]> {
     }
 }
 
-impl<T> Matrix for AVec<AVec<T>>
-where
-    T: UnsignedInteger,
-{
+impl<T> Matrix for AVec<AVec<T>> {
     type MatElement = T;
     type R = AVec<T>;
 
-    fn zeros(row: usize, col: usize) -> Self {
-        avec![avec![T::zero(); col]; row]
-    }
-
     fn dimension(&self) -> (usize, usize) {
         (self.len(), self[0].len())
     }
 }
 
-impl<T> MatrixMut for AVec<AVec<T>> where T: UnsignedInteger {}
+impl<T> MatrixMut for AVec<AVec<T>> {}
 
-impl<T> Matrix for Vec<Vec<T>>
+impl<T> MatrixEntity for AVec<AVec<T>>
 where
-    T: UnsignedInteger,
+    T: Zero + Clone,
 {
+    fn zeros(row: usize, col: usize) -> Self {
+        avec![avec![T::zero();col];row]
+    }
+}
+
+impl<T> Matrix for Vec<Vec<T>> {
     type MatElement = T;
     type R = Vec<T>;
 
+    fn dimension(&self) -> (usize, usize) {
+        (self.len(), self[0].len())
+    }
+}
+impl<T> MatrixMut for Vec<Vec<T>> {}
+
+impl<T> MatrixEntity for Vec<Vec<T>>
+where
+    T: Zero + Clone,
+{
     fn zeros(row: usize, col: usize) -> Self {
         vec![vec![T::zero(); col]; row]
     }
+}
+impl<T> Matrix for &[Vec<T>] {
+    type MatElement = T;
+    type R = Vec<T>;
+
+    fn dimension(&self) -> (usize, usize) {
+        (self.len(), self[0].len())
+    }
+}
+
+impl<T> Matrix for &mut [Vec<T>] {
+    type MatElement = T;
+    type R = Vec<T>;
 
     fn dimension(&self) -> (usize, usize) {
         (self.len(), self[0].len())
     }
 }
-impl<T> MatrixMut for Vec<Vec<T>> where T: UnsignedInteger {}
+impl<T> MatrixMut for &mut [Vec<T>] {}
+
+impl<T> Matrix for &[AVec<T>] {
+    type MatElement = T;
+    type R = AVec<T>;
+
+    fn dimension(&self) -> (usize, usize) {
+        (self.len(), self[0].len())
+    }
+}
+impl<T> Matrix for &mut [AVec<T>] {
+    type MatElement = T;
+    type R = AVec<T>;
+
+    fn dimension(&self) -> (usize, usize) {
+        (self.len(), self[0].len())
+    }
+}
+impl<T> MatrixMut for &mut [AVec<T>] {}
 
 #[cfg(test)]
 mod test {
@@ -236,4 +292,14 @@ mod test {
         v3.set(1, 1, 0);
         assert_eq!(0, v3[1][1]);
     }
+
+    #[test]
+    fn test_matrix_split() {
+        let v1 = vec![1_u64, 2_u64];
+        let v2 = vec![3_u64, 4_u64];
+        let mut v3 = vec![v1, v2];
+        let (mut first, mut second) = v3.split_at_row(1);
+        first.set(0, 0, 0);
+        assert_eq!(v3[0][0], 0)
+    }
 }
diff --git a/src/core_crypto/ring.rs b/src/core_crypto/ring.rs
@@ -670,6 +670,7 @@ mod tests {
 
     use crate::{
         core_crypto::{
+            matrix::MatrixEntity,
             modulus::{ModulusBackendConfig, NativeModulusBackend},
             ntt::NativeNTTBackend,
             prime::generate_primes_vec,
@@ -731,7 +732,7 @@ mod tests {
 
         let mut test = DefaultU64SeededRandomGenerator::new();
 
-        let mut poly_q_in = <Vec<Vec<u64>> as Matrix>::zeros(q_chain.len(), n);
+        let mut poly_q_in = <Vec<Vec<u64>> as MatrixEntity>::zeros(q_chain.len(), n);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test, &q_chain, &mut poly_q_in);
         let mut poly_p_out = Vec::<Vec<u64>>::zeros(p_chain.len(), n);
 
@@ -828,10 +829,10 @@ mod tests {
 
         let mut test_rng = DefaultU64SeededRandomGenerator::new();
 
-        let mut poly_q_in = <Vec<Vec<u64>> as Matrix>::zeros(q_chain.len(), n);
+        let mut poly_q_in = <Vec<Vec<u64>> as MatrixEntity>::zeros(q_chain.len(), n);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test_rng, &q_chain, &mut poly_q_in);
 
-        let mut poly_p_out = Vec::<Vec<u64>>::zeros(p_chain.len(), n);
+        let mut poly_p_out = <Vec<Vec<u64>> as MatrixEntity>::zeros(p_chain.len(), n);
 
         switch_crt_basis(
             &mut poly_p_out,
@@ -932,8 +933,8 @@ mod tests {
         let mut test_rng = DefaultU64SeededRandomGenerator::new();
 
         // P U Q
-        let mut poly_q_in = <Vec<Vec<u64>> as Matrix>::zeros(q_chain.len(), n);
-        let mut poly_p_in = <Vec<Vec<u64>> as Matrix>::zeros(p_chain.len(), n);
+        let mut poly_q_in = <Vec<Vec<u64>> as MatrixEntity>::zeros(q_chain.len(), n);
+        let mut poly_p_in = <Vec<Vec<u64>> as MatrixEntity>::zeros(p_chain.len(), n);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test_rng, &q_chain, &mut poly_q_in);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test_rng, &p_chain, &mut poly_p_in);
 
@@ -1040,7 +1041,7 @@ mod tests {
 
         let mut test_rng = DefaultU64SeededRandomGenerator::new();
 
-        let mut poly_q_in = <Vec<Vec<u64>> as Matrix>::zeros(q_chain.len(), n);
+        let mut poly_q_in = <Vec<Vec<u64>> as MatrixEntity>::zeros(q_chain.len(), n);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test_rng, &q_chain, &mut poly_q_in);
 
         let mut poly_t_out = Vec::<Vec<u64>>::zeros(1, n);
@@ -1155,8 +1156,8 @@ mod tests {
         let mut test_rng = DefaultU64SeededRandomGenerator::new();
 
         // Random polynomial in QP
-        let mut poly0_q_part = <Vec<Vec<u64>> as Matrix>::zeros(q_chain.len(), n);
-        let mut poly0_p_part = <Vec<Vec<u64>> as Matrix>::zeros(p_chain.len(), n);
+        let mut poly0_q_part = <Vec<Vec<u64>> as MatrixEntity>::zeros(q_chain.len(), n);
+        let mut poly0_p_part = <Vec<Vec<u64>> as MatrixEntity>::zeros(p_chain.len(), n);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test_rng, &q_chain, &mut poly0_q_part);
         RandomUniformDist::<Vec<Vec<u64>>>::random_fill(&mut test_rng, &p_chain, &mut poly0_p_part);
 

diff --git a/src/schemes/bfv/default_impl/entities.rs b/src/schemes/bfv/default_impl/entities.rs
@@ -1,7 +1,7 @@
 use crate::{
     ciphertext::{Ciphertext, InitLevelledCiphertext, Representation, RlweCiphertext},
     core_crypto::{
-        matrix::{Drop2Dimension, Matrix, MatrixMut, Row, RowMut},
+        matrix::{Drop2Dimension, Matrix, MatrixEntity, MatrixMut, Row, RowMut},
         modulus::{
             BarrettBackend, ModulusArithmeticBackend, ModulusBackendConfig, ModulusVecBackend,
             MontgomeryBackend, MontgomeryScalar, NativeModulusBackend,
@@ -60,7 +60,7 @@ impl BfvSecretKey {
 
 impl<P> Encryptor<[u64], BfvCiphertextScalarU64GenericStorage<P>> for BfvSecretKey
 where
-    P: TryConvertFrom<[i32], Parameters = [u64]> + MatrixMut<MatElement = u64>,
+    P: TryConvertFrom<[i32], Parameters = [u64]> + MatrixMut<MatElement = u64> + MatrixEntity,
     <P as Matrix>::R: RowMut,
 {
     fn encrypt(&self, message: &[u64]) -> BfvCiphertextScalarU64GenericStorage<P> {
@@ -77,6 +77,7 @@ impl<P> Decryptor<Vec<u64>, BfvCiphertextScalarU64GenericStorage<P>> for BfvSecr
 where
     P: Matrix<MatElement = u64>
         + MatrixMut
+        + MatrixEntity
         + Drop2Dimension
         + TryConvertFrom<[i32], Parameters = [u64]>
         + Clone,
@@ -100,7 +101,7 @@ pub(super) struct BfvCiphertextScalarU64GenericStorage<P> {
 impl<P, R> Ciphertext for BfvCiphertextScalarU64GenericStorage<P>
 where
     R: Row<Element = u64> + RowMut,
-    P: Matrix<MatElement = u64> + MatrixMut<MatElement = u64, R = R>,
+    P: Matrix<MatElement = u64> + MatrixMut<MatElement = u64, R = R> + MatrixEntity,
 {
     type Scalar = u64;
     type Poly = P;
@@ -133,7 +134,7 @@ impl<P> InitLevelledCiphertext for BfvCiphertextScalarU64GenericStorage<P> {
 impl<P, R> RlweCiphertext for BfvCiphertextScalarU64GenericStorage<P>
 where
     R: Row<Element = u64> + RowMut,
-    P: Matrix<MatElement = u64> + MatrixMut<MatElement = u64, R = R>,
+    P: Matrix<MatElement = u64> + MatrixMut<MatElement = u64, R = R> + MatrixEntity,
 {
     fn c_partq(&self) -> &[Self::Poly] {
         &self.c_partq

diff --git a/src/schemes/bfv/ops.rs b/src/schemes/bfv/ops.rs
@@ -5,7 +5,7 @@ use rand::{CryptoRng, RngCore};
 use crate::{
     ciphertext::{Ciphertext, InitLevelledCiphertext, Representation, RlweCiphertext},
     core_crypto::{
-        matrix::{Matrix, MatrixMut, RowMut},
+        matrix::{Matrix, MatrixEntity, MatrixMut, RowMut},
         modulus::ModulusVecBackend,
         ntt::Ntt,
         num::UnsignedInteger,
@@ -72,7 +72,7 @@ pub fn simd_decode_message<
 
 pub fn secret_key_encryption<
     Scalar: UnsignedInteger,
-    Poly: MatrixMut<MatElement = Scalar>,
+    Poly: MatrixMut<MatElement = Scalar> + MatrixEntity,
     S: SecretKey<Scalar = i32>,
     P: BfvEncryptionParameters<Scalar = Scalar>,
     C: RlweCiphertext<Poly = Poly> + InitLevelledCiphertext<C = Vec<Poly>>,
@@ -303,7 +303,7 @@ pub fn ciphertext_mul<
     let ring_size = parameters.ring_size();
 
     // Scale and round c1 in basis Q by \frac{P}{Q} and output c0 in basis P
-    let mut poverq_c10_partp = <C::Poly as Matrix>::zeros(p_size, ring_size);
+    let mut poverq_c10_partp = <C::Poly>::zeros(p_size, ring_size);
     fast_convert_p_over_q(
         &mut poverq_c10_partp,
         &c1.c_partq()[0],

diff --git a/src/schemes/ckks/default_impl/entities.rs b/src/schemes/ckks/default_impl/entities.rs
@@ -10,7 +10,7 @@ use rand_chacha::{ChaCha8Core, ChaCha8Rng};
 use crate::{
     ciphertext::{Ciphertext, Representation, RlweCiphertext, SeededCiphertext},
     core_crypto::{
-        matrix::{Matrix, MatrixMut, RowMut},
+        matrix::{Matrix, MatrixEntity, MatrixMut, RowMut},
         modulus::{ModulusBackendConfig, ModulusVecBackend, NativeModulusBackend},
         ntt::{NativeNTTBackend, Ntt, NttConfig},
         num::{big_float::BigFloat, BFloat, ComplexNumber},
@@ -45,7 +45,7 @@ pub static NATIVE_CKKS_CLIENT_PARAMETERS_U64: OnceLock<
     >,
 > = OnceLock::new();
 
-impl<M: MatrixMut<MatElement = u64>> LevelEncoder<M> for Vec<DefaultComplex>
+impl<M: MatrixMut<MatElement = u64> + MatrixEntity> LevelEncoder<M> for Vec<DefaultComplex>
 where
     <M as Matrix>::R: RowMut,
 {
@@ -99,8 +99,12 @@ impl CkksSecretKey {
     }
 }
 
-impl<M: MatrixMut<MatElement = u64> + Clone + TryConvertFrom<[i32], Parameters = [u64]>>
-    Encryptor<[DefaultComplex], CkksCiphertextGenericStorage<M>> for CkksSecretKey
+impl<
+        M: MatrixMut<MatElement = u64>
+            + MatrixEntity
+            + Clone
+            + TryConvertFrom<[i32], Parameters = [u64]>,
+    > Encryptor<[DefaultComplex], CkksCiphertextGenericStorage<M>> for CkksSecretKey
 where
     <M as Matrix>::R: RowMut,
 {
@@ -133,8 +137,12 @@ where
     }
 }
 
-impl<M: MatrixMut<MatElement = u64> + Clone + TryConvertFrom<[i32], Parameters = [u64]>>
-    Decryptor<Vec<DefaultComplex>, CkksCiphertextGenericStorage<M>> for CkksSecretKey
+impl<
+        M: MatrixMut<MatElement = u64>
+            + Clone
+            + MatrixEntity
+            + TryConvertFrom<[i32], Parameters = [u64]>,
+    > Decryptor<Vec<DefaultComplex>, CkksCiphertextGenericStorage<M>> for CkksSecretKey
 where
     <M as Matrix>::R: RowMut,
 {
@@ -289,8 +297,9 @@ pub struct CkksCiphertextGenericStorage<M> {
     scale: DefaultBigFloat,
 }
 
-impl<M: MatrixMut<MatElement = u64>> Ciphertext for CkksCiphertextGenericStorage<M>
+impl<M> Ciphertext for CkksCiphertextGenericStorage<M>
 where
+    M: MatrixMut<MatElement = u64> + MatrixEntity,
     <M as Matrix>::R: RowMut,
 {
     type Poly = M;
@@ -317,8 +326,9 @@ impl<M> SeededCiphertext for CkksCiphertextGenericStorage<M> {
     }
 }
 
-impl<M: MatrixMut<MatElement = u64>> RlweCiphertext for CkksCiphertextGenericStorage<M>
+impl<M> RlweCiphertext for CkksCiphertextGenericStorage<M>
 where
+    M: MatrixMut<MatElement = u64> + MatrixEntity,
     <M as Matrix>::R: RowMut,
 {
     fn c_partq(&self) -> &[Self::Poly] {