tfhe: add external prod TGSW * TLWE, also TLev * Vec<T64>

README.md

@@ -1,14 +1,15 @@
 # fhe-study
 Implementations from scratch done while studying some FHE papers; do not use in production.
 
-- `arith`: contains $\mathbb{Z}_q$, $R_q=\mathbb{Z}_q[X]/(X^N+1)$, $R=\mathbb{Z}[X]/(X^N+1)$, $\mathbb{T}_Q[X]/(X^N +1)$ arithmetic implementations, together with the NTT implementation.
+- `arith`: contains $\mathbb{Z}_q$, $R_q=\mathbb{Z}_q[X]/(X^N+1)$, $R=\mathbb{Z}[X]/(X^N+1)$, $\mathbb{T}_q[X]/(X^N +1)$ arithmetic implementations, together with the NTT implementation.
 - `gfhe`: (gfhe=generalized-fhe) contains the structs and logic for RLWE, GLWE, GLev, GGSW, RGSW cryptosystems, and modulus switching and key switching methods, which can be used by concrete FHE schemes.
 - `bfv`: https://eprint.iacr.org/2012/144.pdf scheme implementation
 - `ckks`: https://eprint.iacr.org/2016/421.pdf scheme implementation
 - `tfhe`: https://eprint.iacr.org/2018/421.pdf scheme implementation
 
-`cargo test --release`
 
+## Run tests
+`cargo test --release`
 
 ## Example of usage
 > the repo is a work in progress, interfaces will change.

gfhe/src/glev.rs

@@ -1,4 +1,5 @@
 use anyhow::Result;
+use itertools::zip_eq;
 use rand::Rng;
 use rand_distr::{Normal, Uniform};
 use std::ops::{Add, Mul};
@@ -7,8 +8,6 @@ use arith::{Ring, TR};
 
 use crate::glwe::{PublicKey, SecretKey, GLWE};
 
-const ERR_SIGMA: f64 = 3.2;
-
 // l GLWEs
 #[derive(Clone, Debug)]
 pub struct GLev<R: Ring, const K: usize>(pub(crate) Vec<GLWE<R, K>>);
@@ -52,6 +51,21 @@ impl<R: Ring, const K: usize> GLev<R, K> {
     }
 }
 
+// dot product between a GLev and Vec<R>.
+// Used for operating decompositions with KSK_i.
+// GLev * Vec<R> --> GLWE
+impl<R: Ring, const K: usize> Mul<Vec<R>> for GLev<R, K> {
+    type Output = GLWE<R, K>;
+    fn mul(self, v: Vec<R>) -> GLWE<R, K> {
+        // l times GLWES
+        let glwes: Vec<GLWE<R, K>> = self.0;
+
+        // l iterations
+        let r: GLWE<R, K> = zip_eq(v, glwes).map(|(v_i, glwe_i)| glwe_i * v_i).sum();
+        r
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;

gfhe/src/glwe.rs

@@ -12,7 +12,8 @@ use arith::{Ring, Rq, Zq, TR};
 
 use crate::glev::GLev;
 
-const ERR_SIGMA: f64 = 3.2;
+// const ERR_SIGMA: f64 = 3.2;
+const ERR_SIGMA: f64 = 0.0; // TODO WIP
 
 /// GLWE implemented over the `Ring` trait, so that it can be also instantiated
 /// over the Torus polynomials 𝕋_<N,q>[X] = 𝕋_q[X]/ (X^N+1).
@@ -68,22 +69,11 @@ impl<R: Ring, const K: usize> GLWE<R, K> {
 
         // K iterations, ksk.0 contains K times GLev
         let rhs: GLWE<R, K> = zip_eq(a.0, ksk.0.clone())
-            .map(|(a_i, ksk_i)| Self::dot_prod(a_i.decompose(beta, l), ksk_i))
+            .map(|(a_i, ksk_i)| ksk_i * a_i.decompose(beta, l)) // dot_product
             .sum();
 
         lhs - rhs
     }
-    // note: a_decomp is of length N
-    fn dot_prod(a_decomp: Vec<R>, ksk_i: GLev<R, K>) -> GLWE<R, K> {
-        // l times GLWES
-        let glwes: Vec<GLWE<R, K>> = ksk_i.0;
-
-        // l iterations
-        let r: GLWE<R, K> = zip_eq(a_decomp, glwes)
-            .map(|(a_d_i, glwe_i)| glwe_i * a_d_i)
-            .sum();
-        r
-    }
 
     // encrypts with the given SecretKey (instead of PublicKey)
     pub fn encrypt_s(

tfhe/src/tgsw.rs

@@ -37,6 +37,28 @@ impl<const K: usize> TGSW<K> {
     }
 }
 
+// external product TGSW x TLWE
+impl<const K: usize> Mul<TLWE<K>> for TGSW<K> {
+    type Output = TLWE<K>;
+
+    fn mul(self, tlwe: TLWE<K>) -> TLWE<K> {
+        let beta: u32 = 2;
+        let l: u32 = 64; // TODO wip
+
+        // since N=1, each tlwe element is a vector of length=1, decomposed into
+        // l elements, and we have K of them
+        let tlwe_ab: Vec<T64> = [tlwe.0 .0 .0.clone(), vec![tlwe.0 .1]].concat();
+
+        let tgsw_ab: Vec<TLev<K>> = [self.0.clone(), vec![self.1]].concat();
+        assert_eq!(tgsw_ab.len(), tlwe_ab.len());
+
+        let r: TLWE<K> = zip_eq(tgsw_ab, tlwe_ab)
+            .map(|(tlev_i, tlwe_i)| tlev_i * tlwe_i.decompose(beta, l))
+            .sum();
+        r
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;
@@ -71,4 +93,42 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_external_product() -> Result<()> {
+        const T: u64 = 2; // plaintext modulus
+        const K: usize = 32;
+
+        let beta: u32 = 2;
+        let l: u32 = 64;
+
+        let mut rng = rand::thread_rng();
+        let msg_dist = Uniform::new(0_u64, T);
+
+        for i in 0..50 {
+            dbg!(&i);
+            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
+
+            let m1: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
+            let p1: T64 = TLev::<K>::encode::<T>(&m1);
+
+            let m2: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
+            let p2: T64 = TLWE::<K>::encode::<T>(&m2); // scaled by delta
+
+            let tgsw = TGSW::<K>::encrypt_s(&mut rng, beta, l, &sk, &p1)?;
+            let tlwe = TLWE::<K>::encrypt_s(&mut rng, &sk, &p2)?;
+
+            let res: TLWE<K> = tgsw * tlwe;
+
+            // let p_recovered = res.decrypt(&sk, beta);
+            let p_recovered = res.decrypt(&sk);
+            // downscaled by delta^-1
+            let res_recovered = TLWE::<K>::decode::<T>(&p_recovered);
+
+            // assert_eq!(m1 * m2, m_recovered);
+            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), res_recovered);
+        }
+
+        Ok(())
+    }
 }

tfhe/src/tlev.rs

@@ -64,6 +64,27 @@ impl<const K: usize> TLev<K> {
 }
 // TODO review u64::MAX, since is -1 of the value we actually want
 
+impl<const K: usize> TLev<K> {
+    pub fn iter(&self) -> std::slice::Iter<TLWE<K>> {
+        self.0.iter()
+    }
+}
+
+// dot product between a TLev and Vec<T64>, usually Vec<T64> comes from a
+// decomposition of T64
+// TLev * Vec<T64> --> TLWE
+impl<const K: usize> Mul<Vec<T64>> for TLev<K> {
+    type Output = TLWE<K>;
+    fn mul(self, v: Vec<T64>) -> Self::Output {
+        assert_eq!(self.0.len(), v.len());
+
+        // l TLWES
+        let tlwes: Vec<TLWE<K>> = self.0;
+        let r: TLWE<K> = zip_eq(v, tlwes).map(|(a_d_i, glwe_i)| glwe_i * a_d_i).sum();
+        r
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;
@@ -98,4 +119,37 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_tlev_vect64_product() -> Result<()> {
+        const T: u64 = 2; // plaintext modulus
+        const K: usize = 16;
+
+        let beta: u32 = 2;
+        let l: u32 = 16;
+
+        let mut rng = rand::thread_rng();
+        let msg_dist = Uniform::new(0_u64, T);
+
+        for _ in 0..200 {
+            let (sk, pk) = TLWE::<K>::new_key(&mut rng)?;
+
+            let m1: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
+            let m2: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
+            let p1: T64 = TLev::<K>::encode::<T>(&m1);
+            let p2: T64 = TLev::<K>::encode::<T>(&m2);
+
+            let c1 = TLev::<K>::encrypt(&mut rng, beta, l, &pk, &p1)?;
+            let c2 = p2.decompose(beta, l);
+
+            let c3 = c1 * c2;
+
+            let p_recovered = c3.decrypt(&sk);
+            let m_recovered = TLev::<K>::decode::<T>(&p_recovered);
+
+            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m_recovered);
+        }
+
+        Ok(())
+    }
 }