tfhe: add blind_rotation & lookup table computation

gfhe/src/glwe.rs

@@ -33,6 +33,9 @@ impl GLWE {
     pub fn zero() -> Self {
         Self(TR::zero(), R::zero())
     }
+    pub fn from_plaintext(p: R) -> Self {
+        Self(TR::zero(), p)
+    }
 
     pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<R, K>, PublicKey<R, K>)> {
         let Xi_key = Uniform::new(0_f64, 2_f64);

tfhe/src/tglwe.rs

@@ -19,12 +19,12 @@ pub struct SecretKey(pub glwe::SecretKey,
 // 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;
+    pub fn to_tlwe<const KN: usize>(&self) -> tlwe::SecretKey<KN> {
+        let s: TR<Tn<N>, K> = self.0 .0.clone();
 
         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)))
+        tlwe::SecretKey::<KN>(glwe::SecretKey::<T64, KN>(TR::<T64, KN>::new(r)))
     }
 }
 
@@ -37,15 +37,18 @@ impl TGLWE {
     pub fn zero() -> Self {
         Self(GLWE::<Tn<N>, K>::zero())
     }
+    pub fn from_plaintext(p: Tn<N>) -> Self {
+        Self(GLWE::<Tn<N>, K>::from_plaintext(p))
+    }
 
     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
+        // 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())?;
+        let pk: PublicKey<N, K> = GLWE::pk_from_sk(rng, sk.0.clone())?;
         Ok((sk, pk))
     }
 
@@ -73,7 +76,7 @@ impl TGLWE {
     }
 
     /// Sample extraction / Coefficient extraction
-    pub fn sample_extraction(&self, h: usize) -> TLWE<K> {
+    pub fn sample_extraction<const KN: usize>(&self, h: usize) -> TLWE<KN> {
         assert!(h < N);
 
         let a: TR<Tn<N>, K> = self.0 .0.clone();
@@ -91,6 +94,11 @@ impl TGLWE {
 
         TLWE(GLWE(TR(new_a), self.0 .1.coeffs()[h]))
     }
+    pub fn left_rotate(&self, h: usize) -> Self {
+        dbg!(&h);
+        let (a, b): (TR<Tn<N>, K>, Tn<N>) = (self.0 .0.clone(), self.0 .1);
+        Self(GLWE(a.left_rotate(h), b.left_rotate(h)))
+    }
 }
 
 impl<const N: usize, const K: usize> Add<TGLWE<N, K>> for TGLWE<N, K> {
@@ -291,13 +299,14 @@ mod tests {
         const T: u64 = 128; // msg space (msg modulus)
         const N: usize = 64;
         const K: usize = 16;
+        const KN: usize = K * N;
 
         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 (sk, pk) = TGLWE::<N, K>::new_key::<KN>(&mut rng)?;
+            let sk_tlwe = sk.to_tlwe::<KN>();
 
             let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let p: Tn<N> = TGLWE::<N, K>::encode::<T>(&m);
@@ -305,10 +314,10 @@ mod tests {
             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 c_h: TLWE<KN> = c.sample_extraction(h);
 
                 let p_recovered = c_h.decrypt(&sk_tlwe);
-                let m_recovered = TLWE::<K>::decode::<T>(&p_recovered);
+                let m_recovered = TLWE::<KN>::decode::<T>(&p_recovered);
                 assert_eq!(m.coeffs()[h], m_recovered.coeffs()[0]);
             }
         }

tfhe/src/tgsw.rs

@@ -7,7 +7,10 @@ use std::ops::{Add, Mul};
 use arith::{Ring, Rq, Tn, T64, TR};
 
 use crate::tlev::TLev;
-use crate::tlwe::{PublicKey, SecretKey, TLWE};
+use crate::{
+    tglwe::TGLWE,
+    tlwe::{PublicKey, SecretKey, TLWE},
+};
 use gfhe::glwe::GLWE;
 
 /// vector of length K+1 = [K], [1]

tfhe/src/tlev.rs

@@ -6,6 +6,7 @@ use std::ops::{Add, Mul};
 
 use arith::{Ring, Rq, Tn, T64, TR};
 
+use crate::tglwe::TGLWE;
 use crate::tlwe::{PublicKey, SecretKey, TLWE};
 
 #[derive(Clone, Debug)]

tfhe/src/tlwe.rs

@@ -1,20 +1,17 @@
 use anyhow::Result;
 use itertools::zip_eq;
-use rand::distributions::Standard;
 use rand::Rng;
-use rand_distr::{Normal, Uniform};
-use std::array;
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Mul, Sub};
 
-use arith::{Ring, Rq, Tn, T64, TR};
+use arith::{Ring, Rq, Tn, Zq, T64, TR};
 use gfhe::{glwe, GLWE};
 
+use crate::tggsw::TGGSW;
 use crate::tlev::TLev;
+use crate::{tglwe, tglwe::TGLWE};
 
-// #[derive(Clone, Debug)]
 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
@@ -32,8 +29,6 @@ impl SecretKey {
     }
 }
 
-// #[derive(Clone, Debug)]
-// pub struct PublicKey<const K: usize>(glwe::PublicKey<T64, K>);
 pub type PublicKey<const K: usize> = glwe::PublicKey<T64, K>;
 
 #[derive(Clone, Debug)]
@@ -49,14 +44,12 @@ impl TLWE {
 
     pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
         let (sk, pk): (glwe::SecretKey<T64, K>, glwe::PublicKey<T64, K>) = GLWE::new_key(rng)?;
-        // Ok((SecretKey(sk), PublicKey(pk)))
         Ok((SecretKey(sk), pk))
     }
 
     pub fn encode<const P: u64>(m: &Rq<P, 1>) -> T64 {
         let delta = u64::MAX / P; // floored
         let coeffs = m.coeffs();
-        // Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
         T64(coeffs[0].0 * delta)
     }
     pub fn decode<const P: u64>(p: &T64) -> Rq<P, 1> {
@@ -105,6 +98,76 @@ impl TLWE {
 
         lhs - rhs
     }
+    // modulus switch from Q (2^64) to Q2 (in blind_rotation Q2=K*N)
+    pub fn mod_switch<const Q2: u64>(&self) -> Self {
+        let a: TR<T64, K> = self.0 .0.mod_switch::<Q2>();
+        let b: T64 = self.0 .1.mod_switch::<Q2>();
+        Self(GLWE(a, b))
+    }
+}
+// NOTE: the ugly const generics are temporary
+pub fn blind_rotation<const N: usize, const K: usize, const KN: usize, const KN2: u64>(
+    c: TLWE<KN>,
+    btk: BootstrappingKey<N, K, KN>,
+    table: TGLWE<N, K>,
+) -> TGLWE<N, K> {
+    let c_kn: TLWE<KN> = c.mod_switch::<KN2>();
+    let (a, b): (TR<T64, KN>, T64) = (c_kn.0 .0, c_kn.0 .1);
+    // two main parts: rotate by a known power of X, rotate by a secret
+    // power of X (using the C gate)
+
+    // table * X^-b, ie. left rotate
+    let v_xb: TGLWE<N, K> = table.left_rotate(b.0 as usize);
+
+    // rotate by a secret power of X using the cmux gate
+    let mut c_j: TGLWE<N, K> = v_xb.clone();
+    let _ = (1..K).map(|j| {
+        c_j = TGGSW::<N, K>::cmux(
+            btk.0[j].clone(),
+            c_j.clone(),
+            c_j.clone().left_rotate(a.0[j].0 as usize),
+        );
+        dbg!(&c_j);
+    });
+    c_j
+}
+
+#[derive(Clone, Debug)]
+pub struct BootstrappingKey<const N: usize, const K: usize, const KN: usize>(pub Vec<TGGSW<N, K>>);
+impl<const N: usize, const K: usize, const KN: usize> BootstrappingKey<N, K, KN> {
+    pub fn from_sk(mut rng: impl Rng, sk: &tglwe::SecretKey<N, K>) -> Result<Self> {
+        let (beta, l) = (2u32, 64u32); // TMP
+                                       //
+        let s: TR<Tn<N>, K> = sk.0 .0.clone();
+
+        // each btk_j = TGGSW_sk(s_i)
+        let btk: Vec<TGGSW<N, K>> = s
+            .iter()
+            .map(|s_i| TGGSW::<N, K>::encrypt_s(&mut rng, beta, l, sk, s_i))
+            .collect::<Result<Vec<_>>>()?;
+
+        Ok(Self(btk))
+    }
+}
+
+pub fn compute_lookup_table<const T: u64, const K: usize, const N: usize>() -> TGLWE<N, K> {
+    // from 2021-1402:
+    // v(x) = \sum_j^{N-1} [(p_j / 2N mod p)/p] X^j
+
+    // matrix of coefficients with size K*N = delta x T
+    let delta: usize = N / T as usize;
+    let values: Vec<Zq<T>> = (0..T).map(|v| Zq::<T>::from_u64(v)).collect();
+    let coeffs: Vec<Zq<T>> = (0..T as usize)
+        .flat_map(|i| vec![values[i]; delta])
+        .collect();
+    let table = Rq::<T, N>::from_vec(coeffs);
+
+    // encode the table as plaintext
+    let v: Tn<N> = TGLWE::<N, K>::encode::<T>(&table);
+
+    // encode the table as TGLWE ciphertext
+    let v: TGLWE<N, K> = TGLWE::<N, K>::from_plaintext(v);
+    v
 }
 
 impl<const K: usize> Add<TLWE<K>> for TLWE<K> {
@@ -170,6 +233,7 @@ impl Mul for TLWE {
 mod tests {
     use anyhow::Result;
     use rand::distributions::Uniform;
+    use std::time::Instant;
 
     use super::*;
 
@@ -186,9 +250,7 @@ mod tests {
             let (sk, pk) = S::new_key(&mut rng)?;
 
             let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
-            dbg!(&m);
             let p: T64 = S::encode::<T>(&m);
-            dbg!(&p);
 
             let c = S::encrypt(&mut rng, &pk, &p)?;
             let p_recovered = c.decrypt(&sk);