tfhe: add TGLWE.sample_extraction

gfhe/src/glwe.rs

@@ -45,6 +45,15 @@ impl GLWE {
         let pk: PublicKey<R, K> = PublicKey((&a * &s) + e, a);
         Ok((SecretKey(s), pk))
     }
+    pub fn pk_from_sk(mut rng: impl Rng, sk: SecretKey<R, K>) -> Result<PublicKey<R, K>> {
+        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
+
+        let a: TR<R, K> = TR::rand(&mut rng, Uniform::new(0_f64, R::Q as f64));
+        let e = R::rand(&mut rng, Xi_err);
+
+        let pk: PublicKey<R, K> = PublicKey((&a * &sk.0) + e, a);
+        Ok(pk)
+    }
 
     pub fn new_ksk(
         mut rng: impl Rng,

tfhe/src/tglwe.rs

@@ -11,8 +11,23 @@ use arith::{Ring, Rq, Tn, T64, TR};
 use gfhe::{glwe, GLWE};
 
 use crate::tlev::TLev;
+use crate::{tlwe, tlwe::TLWE};
+
+// pub type SecretKey<const N: usize, const K: usize> = glwe::SecretKey<Tn<N>, K>;
+#[derive(Clone, Debug)]
+pub struct SecretKey<const N: usize, const K: usize>(pub glwe::SecretKey<Tn<N>, K>);
+// pub struct SecretKey<const K: usize>(pub tlwe::SecretKey<K>);
+
+impl<const N: usize, const K: usize> SecretKey<N, K> {
+    pub fn to_tlwe<const KN: usize>(self) -> tlwe::SecretKey<K> {
+        let s: TR<Tn<N>, K> = self.0 .0;
+
+        let r: Vec<Vec<T64>> = s.0.iter().map(|s_i| s_i.coeffs()).collect();
+        let r: Vec<T64> = r.into_iter().flatten().collect();
+        tlwe::SecretKey(glwe::SecretKey(TR(r)))
+    }
+}
 
-pub type SecretKey<const N: usize, const K: usize> = glwe::SecretKey<Tn<N>, K>;
 pub type PublicKey<const N: usize, const K: usize> = glwe::PublicKey<Tn<N>, K>;
 
 #[derive(Clone, Debug)]
@@ -23,9 +38,14 @@ impl TGLWE {
         Self(GLWE::<Tn<N>, K>::zero())
     }
 
-    pub fn new_key(rng: impl Rng) -> Result<(SecretKey<N, K>, PublicKey<N, K>)> {
-        let (sk, pk) = GLWE::new_key(rng)?;
-        // Ok((SecretKey(sk), PublicKey(pk)))
+    pub fn new_key<const KN: usize>(
+        mut rng: impl Rng,
+    ) -> Result<(SecretKey<N, K>, PublicKey<N, K>)> {
+        assert_eq!(KN, K * N); // this is wip, while not being able to compute K*N
+        let (sk_tlwe, _) = TLWE::<KN>::new_key(&mut rng)?;
+        // let sk = crate::tlwe::sk_to_tglwe::<N, K, KN>(sk_tlwe);
+        let sk = sk_tlwe.to_tglwe::<N, K>();
+        let pk = GLWE::pk_from_sk(rng, sk.0.clone())?;
         Ok((sk, pk))
     }
 
@@ -41,7 +61,7 @@ impl TGLWE {
 
     // encrypts with the given SecretKey (instead of PublicKey)
     pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<N, K>, p: &Tn<N>) -> Result<Self> {
-        let glwe = GLWE::encrypt_s(rng, &sk, p)?;
+        let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
         Ok(Self(glwe))
     }
     pub fn encrypt(rng: impl Rng, pk: &PublicKey<N, K>, p: &Tn<N>) -> Result<Self> {
@@ -49,7 +69,27 @@ impl TGLWE {
         Ok(Self(glwe))
     }
     pub fn decrypt(&self, sk: &SecretKey<N, K>) -> Tn<N> {
-        self.0.decrypt(&sk)
+        self.0.decrypt(&sk.0)
+    }
+
+    /// Sample extraction / Coefficient extraction
+    pub fn sample_extraction(&self, h: usize) -> TLWE<K> {
+        assert!(h < N);
+
+        let a: TR<Tn<N>, K> = self.0 .0.clone();
+        // set a_{n*i+j} = a_{i, h-j}     if j \in {0, h}
+        //                 -a_{i, n+h-j}  if j \in {h+1, n-1}
+        let new_a: Vec<T64> = a
+            .iter()
+            .flat_map(|a_i| {
+                let a_i = a_i.coeffs();
+                (0..N)
+                    .map(|j| if j <= h { a_i[h - j] } else { -a_i[N + h - j] })
+                    .collect::<Vec<T64>>()
+            })
+            .collect::<Vec<T64>>();
+
+        TLWE(GLWE(TR(new_a), self.0 .1.coeffs()[h]))
     }
 }
 
@@ -130,7 +170,7 @@ mod tests {
         let msg_dist = Uniform::new(0_u64, T);
 
         for _ in 0..200 {
-            let (sk, pk) = S::new_key(&mut rng)?;
+            let (sk, pk) = TGLWE::<N, K>::new_key::<{ K * N }>(&mut rng)?;
 
             let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let p: Tn<N> = S::encode::<T>(&m);
@@ -163,7 +203,7 @@ mod tests {
         let msg_dist = Uniform::new(0_u64, T);
 
         for _ in 0..200 {
-            let (sk, pk) = S::new_key(&mut rng)?;
+            let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
 
             let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
@@ -195,7 +235,7 @@ mod tests {
         let msg_dist = Uniform::new(0_u64, T);
 
         for _ in 0..200 {
-            let (sk, pk) = S::new_key(&mut rng)?;
+            let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
 
             let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
@@ -226,7 +266,7 @@ mod tests {
         let msg_dist = Uniform::new(0_u64, T);
 
         for _ in 0..200 {
-            let (sk, pk) = S::new_key(&mut rng)?;
+            let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
 
             let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
@@ -245,4 +285,34 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_sample_extraction() -> Result<()> {
+        const T: u64 = 128; // msg space (msg modulus)
+        const N: usize = 64;
+        const K: usize = 16;
+
+        let mut rng = rand::thread_rng();
+        let msg_dist = Uniform::new(0_u64, T);
+
+        for _ in 0..20 {
+            let (sk, pk) = TGLWE::<N, K>::new_key::<{ K * N }>(&mut rng)?;
+            let sk_tlwe = sk.to_tlwe::<{ K * N }>();
+
+            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            let p: Tn<N> = TGLWE::<N, K>::encode::<T>(&m);
+
+            let c = TGLWE::<N, K>::encrypt(&mut rng, &pk, &p)?;
+
+            for h in 0..N {
+                let c_h: TLWE<K> = c.sample_extraction(h);
+
+                let p_recovered = c_h.decrypt(&sk_tlwe);
+                let m_recovered = TLWE::<K>::decode::<T>(&p_recovered);
+                assert_eq!(m.coeffs()[h], m_recovered.coeffs()[0]);
+            }
+        }
+
+        Ok(())
+    }
 }

tfhe/src/tgsw.rs

@@ -23,7 +23,7 @@ impl TGSW {
         m: &T64,
     ) -> Result<Self> {
         let a: Vec<TLev<K>> = (0..K)
-            .map(|i| TLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0[i] * *m)))
+            .map(|i| TLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0 .0[i] * *m)))
             .collect::<Result<Vec<_>>>()?;
         let b: TLev<K> = TLev::encrypt_s(&mut rng, beta, l, sk, m)?;
         Ok(Self(a, b))

tfhe/src/tlwe.rs

@@ -13,8 +13,24 @@ use gfhe::{glwe, GLWE};
 use crate::tlev::TLev;
 
 // #[derive(Clone, Debug)]
-// pub struct SecretKey<const K: usize>(glwe::SecretKey<T64, K>);
-pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>;
+pub struct SecretKey<const K: usize>(pub glwe::SecretKey<T64, K>);
+// pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>;
+
+impl<const KN: usize> SecretKey<KN> {
+    /// from TFHE [2018-421] paper: A TLWE key k \in B^n, can be interpreted as a
+    /// TRLWE key K \in B_N[X]^k having the same sequence of coefficients and
+    /// vice-versa.
+    pub fn to_tglwe<const N: usize, const K: usize>(self) -> crate::tglwe::SecretKey<N, K> {
+        let s: TR<T64, KN> = self.0 .0;
+        // split into K vectors, and interpret each of them as a T_N[X]/(X^N+1)
+        // polynomial
+        let r: Vec<Tn<N>> =
+            s.0.chunks(N)
+                .map(|v| Tn::<N>::from_vec(v.to_vec()))
+                .collect();
+        crate::tglwe::SecretKey(glwe::SecretKey::<Tn<N>, K>(TR(r)))
+    }
+}
 
 // #[derive(Clone, Debug)]
 // pub struct PublicKey<const K: usize>(glwe::PublicKey<T64, K>);
@@ -32,9 +48,9 @@ impl TLWE {
     }
 
     pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
-        let (sk, pk) = GLWE::new_key(rng)?;
+        let (sk, pk): (glwe::SecretKey<T64, K>, glwe::PublicKey<T64, K>) = GLWE::new_key(rng)?;
         // Ok((SecretKey(sk), PublicKey(pk)))
-        Ok((sk, pk))
+        Ok((SecretKey(sk), pk))
     }
 
     pub fn encode<const P: u64>(m: &Rq<P, 1>) -> T64 {
@@ -50,7 +66,7 @@ impl TLWE {
 
     // encrypts with the given SecretKey (instead of PublicKey)
     pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &T64) -> Result<Self> {
-        let glwe = GLWE::encrypt_s(rng, &sk, p)?;
+        let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
         Ok(Self(glwe))
     }
     pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &T64) -> Result<Self> {
@@ -58,7 +74,7 @@ impl TLWE {
         Ok(Self(glwe))
     }
     pub fn decrypt(&self, sk: &SecretKey<K>) -> T64 {
-        self.0.decrypt(&sk)
+        self.0.decrypt(&sk.0)
     }
 
     pub fn new_ksk(
@@ -72,7 +88,7 @@ impl TLWE {
             .into_iter()
             .map(|i|
                 // treat sk_i as the msg being encrypted
-                TLev::<K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
+                TLev::<K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0.0 .0[i]))
             .collect::<Result<Vec<_>>>()?;
 
         Ok(KSK(r))