adapt gfhe to work with Ring trait, so that it can work with Rq & Tn (for TFHE)

arith/src/ring.rs

@@ -3,7 +3,10 @@ use std::fmt::Debug;
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
 
-/// Represents a ring element. Currently implemented by ring_n.rs#R and ring_nq.rs#Rq.
+/// Represents a ring element. Currently implemented by ring_n.rs#R and
+/// ring_nq.rs#Rq. Is not a 'pure algebraic ring', but more a custom trait
+/// definition which includes methods like `mod_switch`.
+// assumed to be mod (X^N +1)
 pub trait Ring:
     Sized
     + Add<Output = Self>
@@ -11,17 +14,22 @@ pub trait Ring:
     + Sum
     + Sub<Output = Self>
     + SubAssign
-    + Mul<Output = Self>
+    + Mul<Output = Self> // internal product
-    + Mul<u64, Output = Self> // scalar mul
+    + Mul<u64, Output = Self> // scalar mul, external product
+    + Mul<Self::C, Output = Self>
     + PartialEq
     + Debug
     + Clone
+    + Copy
     + Sum<<Self as Add>::Output>
     + Sum<<Self as Mul>::Output>
 {
     /// C defines the coefficient type
     type C: Debug + Clone;
 
+    const Q: u64;
+    const N: usize;
+
     fn coeffs(&self) -> Vec<Self::C>;
     fn zero() -> Self;
     // note/wip/warning: dist (0,q) with f64, will output more '0=q' elements than other values
@@ -31,6 +39,9 @@ pub trait Ring:
 
     fn decompose(&self, beta: u32, l: u32) -> Vec<Self>;
 
+    fn remodule<const P: u64>(&self) -> impl Ring;
+    fn mod_switch<const P: u64>(&self) -> impl Ring;
+
     /// returns [ [(num/den) * self].round() ] mod q
     /// ie. performs the multiplication and division over f64, and then it
     /// rounds the result, only applying the mod Q (if the ring is mod Q) at the

arith/src/ring_n.rs

@@ -15,9 +15,13 @@ use crate::Ring;
 #[derive(Clone, Copy)]
 pub struct R<const N: usize>(pub [i64; N]);
 
-impl<const N: usize> Ring for R<N> {
+// impl<const N: usize> Ring for R<N> {
-    type C = i64;
+impl<const N: usize> R<N> {
-    fn coeffs(&self) -> Vec<Self::C> {
+    // type C = i64;
+    // const Q: u64 = i64::MAX as u64; // WIP
+    // const N: usize = N;
+
+    pub fn coeffs(&self) -> Vec<i64> {
         self.0.to_vec()
     }
     fn zero() -> Self {
@@ -32,20 +36,28 @@ impl<const N: usize> Ring for R<N> {
         // Self(coeffs)
     }
 
-    fn from_vec(coeffs: Vec<Self::C>) -> Self {
+    pub fn from_vec(coeffs: Vec<i64>) -> Self {
         let mut p = coeffs;
         modulus::<N>(&mut p);
         Self(array::from_fn(|i| p[i]))
     }
 
+    /*
     // returns the decomposition of each polynomial coefficient
     fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
         unimplemented!();
         // array::from_fn(|i| self.coeffs[i].decompose(beta, l))
     }
 
-    // performs the multiplication and division over f64, and then it rounds the
+    fn remodule<const P: u64>(&self) -> impl Ring {
-    // result, only applying the mod Q at the end
+        unimplemented!()
+    }
+    fn mod_switch<const P: u64, const M: usize>(&self) -> R<N> {
+        unimplemented!()
+    }
+
+    /// performs the multiplication and division over f64, and then it rounds the
+    /// result, only applying the mod Q at the end
     fn mul_div_round(&self, num: u64, den: u64) -> Self {
         unimplemented!()
         // fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> {
@@ -56,6 +68,7 @@ impl<const N: usize> Ring for R<N> {
         //     .collect();
         // crate::Rq::<Q, N>::from_vec_f64(r)
     }
+    */
 }
 
 impl<const Q: u64, const N: usize> From<crate::ring_nq::Rq<Q, N>> for R<N> {
@@ -65,9 +78,9 @@ impl<const Q: u64, const N: usize> From<crate::ring_nq::Rq<Q, N>> for R<N> {
 }
 
 impl<const N: usize> R<N> {
-    pub fn coeffs(&self) -> [i64; N] {
+    // pub fn coeffs(&self) -> [i64; N] {
-        self.0
+    //     self.0
-    }
+    // }
     pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
         crate::Rq::<Q, N>::from(self)
     }
@@ -318,6 +331,13 @@ pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
     R::<N>::from_vec(r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
 }
 
+impl<const N: usize> Mul<i64> for R<N> {
+    type Output = Self;
+
+    fn mul(self, s: i64) -> Self {
+        self.mul_by_i64(s)
+    }
+}
 // mul by u64
 impl<const N: usize> Mul<u64> for R<N> {
     type Output = Self;

arith/src/ring_nq.rs

@@ -30,6 +30,9 @@ pub struct Rq<const Q: u64, const N: usize> {
 impl<const Q: u64, const N: usize> Ring for Rq<Q, N> {
     type C = Zq<Q>;
 
+    const Q: u64 = Q;
+    const N: usize = N;
+
     fn coeffs(&self) -> Vec<Self::C> {
         self.coeffs.to_vec()
     }
@@ -71,9 +74,26 @@ impl<const Q: u64, const N: usize> Ring for Rq<Q, N> {
         r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
     }
 
-    // returns [ [(num/den) * self].round() ] mod q
+    // Warning: this method will behave differently depending on the values P and Q:
-    // ie. performs the multiplication and division over f64, and then it rounds the
+    // if Q<P, it just 'renames' the modulus parameter to P
-    // result, only applying the mod Q at the end
+    // if Q>=P, it crops to mod P
+    fn remodule<const P: u64>(&self) -> Rq<P, N> {
+        Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
+    }
+
+    /// perform the mod switch operation from Q to Q', where Q2=Q'
+    // fn mod_switch<const P: u64, const M: usize>(&self) -> impl Ring {
+    fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
+        // assert_eq!(N, M); // sanity check
+        Rq::<P, N> {
+            coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<P>()),
+            evals: None,
+        }
+    }
+
+    /// returns [ [(num/den) * self].round() ] mod q
+    /// ie. performs the multiplication and division over f64, and then it rounds the
+    /// result, only applying the mod Q at the end
     fn mul_div_round(&self, num: u64, den: u64) -> Self {
         let r: Vec<f64> = self
             .coeffs()
@@ -183,17 +203,9 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
     // Warning: this method will behave differently depending on the values P and Q:
     // if Q<P, it just 'renames' the modulus parameter to P
     // if Q>=P, it crops to mod P
-    pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
+    // pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
-        Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
+    //     Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
-    }
+    // }
-
-    /// perform the mod switch operation from Q to Q', where Q2=Q'
-    pub fn mod_switch<const Q2: u64>(&self) -> Rq<Q2, N> {
-        Rq::<Q2, N> {
-            coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<Q2>()),
-            evals: None,
-        }
-    }
 
     // applies mod(T) to all coefficients of self
     pub fn coeffs_mod<const T: u64>(&self) -> Self {

arith/src/ring_torus.rs

@@ -12,7 +12,7 @@ use std::array;
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
 
-use crate::{ring::Ring, torus::T64};
+use crate::{ring::Ring, torus::T64, Rq, Zq};
 
 /// 𝕋_<N,Q>[X] = 𝕋<Q>[X]/(X^N +1), polynomials modulo X^N+1 with coefficients in
 /// 𝕋, where Q=2^64.
@@ -22,6 +22,9 @@ pub struct Tn<const N: usize>(pub [T64; N]);
 impl<const N: usize> Ring for Tn<N> {
     type C = T64;
 
+    const Q: u64 = u64::MAX; // WIP
+    const N: usize = N;
+
     fn coeffs(&self) -> Vec<T64> {
         self.0.to_vec()
     }
@@ -50,6 +53,22 @@ impl<const N: usize> Ring for Tn<N> {
         r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
     }
 
+    fn remodule<const P: u64>(&self) -> Tn<N> {
+        todo!()
+        // Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
+    }
+
+    // fn mod_switch<const P: u64>(&self) -> impl Ring {
+    fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
+        // unimplemented!()
+        // TODO WIP
+        let coeffs = array::from_fn(|i| Zq::<P>::from_u64(self.0[i].mod_switch::<P>().0));
+        Rq::<P, N> {
+            coeffs,
+            evals: None,
+        }
+    }
+
     /// returns [ [(num/den) * self].round() ] mod q
     /// ie. performs the multiplication and division over f64, and then it rounds the
     /// result, only applying the mod Q at the end
@@ -174,12 +193,19 @@ fn modulus_u128<const N: usize>(p: &mut Vec<u128>) {
         return;
     }
     for i in N..p.len() {
-        p[i - N] = p[i - N].clone() - p[i].clone();
+        // p[i - N] = p[i - N].clone() - p[i].clone();
+        p[i - N] = p[i - N].wrapping_sub(p[i]);
         p[i] = 0;
     }
     p.truncate(N);
 }
 
+impl<const N: usize> Mul<T64> for Tn<N> {
+    type Output = Self;
+    fn mul(self, s: T64) -> Self {
+        Self(array::from_fn(|i| self.0[i] * s))
+    }
+}
 // mul by u64
 impl<const N: usize> Mul<u64> for Tn<N> {
     type Output = Self;

arith/src/torus.rs

@@ -33,6 +33,9 @@ impl T64 {
             .map(|i| T64(((self.0 >> i) & 1) as u64))
             .collect()
     }
+    pub fn mod_switch<const Q2: u64>(&self) -> T64 {
+        todo!()
+    }
 }
 
 impl Add<T64> for T64 {

gfhe/src/glev.rs

@@ -3,7 +3,7 @@ use rand::Rng;
 use rand_distr::{Normal, Uniform};
 use std::ops::{Add, Mul};
 
-use arith::{Ring, Rq, TR};
+use arith::{Ring, TR};
 
 use crate::glwe::{PublicKey, SecretKey, GLWE};
 
@@ -11,25 +11,19 @@ const ERR_SIGMA: f64 = 3.2;
 
 // l GLWEs
 #[derive(Clone, Debug)]
-pub struct GLev<const Q: u64, const N: usize, const K: usize>(pub(crate) Vec<GLWE<Q, N, K>>);
+pub struct GLev<R: Ring, const K: usize>(pub(crate) Vec<GLWE<R, K>>);
 
-impl<const Q: u64, const N: usize, const K: usize> GLev<Q, N, K> {
+impl<R: Ring, const K: usize> GLev<R, K> {
-    pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> {
-        m.remodule::<Q>()
-    }
-    pub fn decode<const T: u64>(p: &Rq<Q, N>) -> Rq<T, N> {
-        p.remodule::<T>()
-    }
     pub fn encrypt(
         mut rng: impl Rng,
         beta: u32,
         l: u32,
-        pk: &PublicKey<Q, N, K>,
+        pk: &PublicKey<R, K>,
-        m: &Rq<Q, N>,
+        m: &R,
     ) -> Result<Self> {
-        let glev: Vec<GLWE<Q, N, K>> = (1..l + 1)
+        let glev: Vec<GLWE<R, K>> = (0..l)
             .map(|i| {
-                GLWE::<Q, N, K>::encrypt(&mut rng, pk, &(*m * (Q / beta.pow(i as u32) as u64)))
+                GLWE::<R, K>::encrypt(&mut rng, pk, &(*m * (R::Q / beta.pow(i as u32) as u64)))
             })
             .collect::<Result<Vec<_>>>()?;
 
@@ -39,21 +33,22 @@ impl<const Q: u64, const N: usize, const K: usize> GLev<Q, N, K> {
         mut rng: impl Rng,
         beta: u32,
         l: u32,
-        sk: &SecretKey<Q, N, K>,
+        sk: &SecretKey<R, K>,
-        m: &Rq<Q, N>,
+        m: &R,
+        // delta: u64,
     ) -> Result<Self> {
-        let glev: Vec<GLWE<Q, N, K>> = (1..l + 1)
+        let glev: Vec<GLWE<R, K>> = (1..l + 1)
             .map(|i| {
-                GLWE::<Q, N, K>::encrypt_s(&mut rng, sk, &(*m * (Q / beta.pow(i as u32) as u64)))
+                GLWE::<R, K>::encrypt_s(&mut rng, sk, &(*m * (R::Q / beta.pow(i as u32) as u64)))
             })
             .collect::<Result<Vec<_>>>()?;
 
         Ok(Self(glev))
     }
 
-    pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, beta: u32) -> Rq<Q, N> {
+    pub fn decrypt<const T: u64>(&self, sk: &SecretKey<R, K>, beta: u32) -> R {
-        let pt = self.0[0].decrypt(sk);
+        let pt = self.0[1].decrypt(sk);
-        pt.mul_div_round(beta as u64, Q)
+        pt.mul_div_round(beta as u64, R::Q)
     }
 }
 
@@ -63,6 +58,7 @@ mod tests {
     use rand::distributions::Uniform;
 
     use super::*;
+    use arith::Rq;
 
     #[test]
     fn test_encrypt_decrypt() -> Result<()> {
@@ -70,25 +66,25 @@ mod tests {
         const N: usize = 128;
         const T: u64 = 2; // plaintext modulus
         const K: usize = 16;
-        type S = GLev<Q, N, K>;
+        type S = GLev<Rq<Q, N>, K>;
 
         let beta: u32 = 2;
         let l: u32 = 16;
 
+        // let delta: u64 = Q / T; // floored
         let mut rng = rand::thread_rng();
 
         for _ in 0..200 {
-            let (sk, pk) = GLWE::<Q, N, K>::new_key(&mut rng)?;
+            let (sk, pk) = GLWE::<Rq<Q, N>, K>::new_key(&mut rng)?;
 
             let msg_dist = Uniform::new(0_u64, T);
             let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
-            let p: Rq<Q, N> = S::encode::<T>(&m); // plaintext
+            let m: Rq<Q, N> = m.remodule::<Q>();
 
-            let c = S::encrypt(&mut rng, beta, l, &pk, &p)?;
+            let c = S::encrypt(&mut rng, beta, l, &pk, &m)?;
-            let p_recovered = c.decrypt::<T>(&sk, beta);
+            let m_recovered = c.decrypt::<T>(&sk, beta);
-            let m_recovered = S::decode::<T>(&p_recovered);
 
-            assert_eq!(m, m_recovered);
+            assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
         }
 
         Ok(())

gfhe/src/glwe.rs

@@ -1,3 +1,6 @@
+//! Generalized LWE.
+//!
+
 use anyhow::Result;
 use itertools::zip_eq;
 use rand::Rng;
@@ -11,32 +14,34 @@ use crate::glev::GLev;
 
 const ERR_SIGMA: f64 = 3.2;
 
+/// 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).
 #[derive(Clone, Debug)]
-pub struct GLWE<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>, Rq<Q, N>);
+pub struct GLWE<R: Ring, const K: usize>(TR<R, K>, R);
 
 #[derive(Clone, Debug)]
-pub struct SecretKey<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>);
+pub struct SecretKey<R: Ring, const K: usize>(TR<R, K>);
 #[derive(Clone, Debug)]
-pub struct PublicKey<const Q: u64, const N: usize, const K: usize>(Rq<Q, N>, TR<Rq<Q, N>, K>);
+pub struct PublicKey<R: Ring, const K: usize>(R, TR<R, K>);
 
 // K GLevs, each KSK_i=l GLWEs
 #[derive(Clone, Debug)]
-pub struct KSK<const Q: u64, const N: usize, const K: usize>(Vec<GLev<Q, N, K>>);
+pub struct KSK<R: Ring, const K: usize>(Vec<GLev<R, K>>);
 
-impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> GLWE<R, K> {
     pub fn zero() -> Self {
-        Self(TR::zero(), Rq::zero())
+        Self(TR::zero(), R::zero())
     }
 
-    pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<Q, N, K>, PublicKey<Q, N, K>)> {
+    pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<R, K>, PublicKey<R, K>)> {
         let Xi_key = Uniform::new(0_f64, 2_f64);
         let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
 
-        let s: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key);
+        let s: TR<R, K> = TR::rand(&mut rng, Xi_key);
-        let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Uniform::new(0_f64, Q as f64));
+        let a: TR<R, K> = TR::rand(&mut rng, Uniform::new(0_f64, R::Q as f64));
-        let e = Rq::<Q, N>::rand(&mut rng, Xi_err);
+        let e = R::rand(&mut rng, Xi_err);
 
-        let pk: PublicKey<Q, N, K> = PublicKey((&a * &s) + e, a);
+        let pk: PublicKey<R, K> = PublicKey((&a * &s) + e, a);
         Ok((SecretKey(s), pk))
     }
 
@@ -44,42 +49,85 @@ impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
         mut rng: impl Rng,
         beta: u32,
         l: u32,
-        sk: &SecretKey<Q, N, K>,
+        sk: &SecretKey<R, K>,
-        new_sk: &SecretKey<Q, N, K>,
+        new_sk: &SecretKey<R, K>,
-    ) -> Result<KSK<Q, N, K>> {
+    ) -> Result<KSK<R, K>> {
-        let r: Vec<GLev<Q, N, K>> = (0..K)
+        let r: Vec<GLev<R, K>> = (0..K)
             .into_iter()
             .map(|i|
                 // treat sk_i as the msg being encrypted
-                GLev::<Q, N, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
+                GLev::<R, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
             .collect::<Result<Vec<_>>>()?;
 
         Ok(KSK(r))
     }
-    pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<Q, N, K>) -> Self {
+    pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<R, K>) -> Self {
-        let (a, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1);
+        let (a, b): (TR<R, K>, R) = (self.0.clone(), self.1);
 
-        let lhs: GLWE<Q, N, K> = GLWE(TR::zero(), b);
+        let lhs: GLWE<R, K> = GLWE(TR::zero(), b);
 
         // K iterations, ksk.0 contains K times GLev
-        let rhs: GLWE<Q, N, K> = zip_eq(a.0, ksk.0.clone())
+        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))
             .sum();
 
         lhs - rhs
     }
     // note: a_decomp is of length N
-    fn dot_prod(a_decomp: Vec<Rq<Q, N>>, ksk_i: GLev<Q, N, K>) -> GLWE<Q, N, K> {
+    fn dot_prod(a_decomp: Vec<R>, ksk_i: GLev<R, K>) -> GLWE<R, K> {
         // l times GLWES
-        let glwes: Vec<GLWE<Q, N, K>> = ksk_i.0;
+        let glwes: Vec<GLWE<R, K>> = ksk_i.0;
 
         // l iterations
-        let r: GLWE<Q, N, K> = zip_eq(a_decomp, glwes)
+        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(
+        mut rng: impl Rng,
+        sk: &SecretKey<R, K>,
+        m: &R, // already scaled
+    ) -> Result<Self> {
+        let Xi_key = Uniform::new(0_f64, 2_f64);
+        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
+
+        let a: TR<R, K> = TR::rand(&mut rng, Xi_key);
+        let e = R::rand(&mut rng, Xi_err);
+
+        let b: R = (&a * &sk.0) + *m + e;
+        Ok(Self(a, b))
+    }
+    pub fn encrypt(
+        mut rng: impl Rng,
+        pk: &PublicKey<R, K>,
+        m: &R, // already scaled
+    ) -> Result<Self> {
+        let Xi_key = Uniform::new(0_f64, 2_f64);
+        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
+
+        let u: R = R::rand(&mut rng, Xi_key);
+
+        let e0 = R::rand(&mut rng, Xi_err);
+        let e1 = TR::<R, K>::rand(&mut rng, Xi_err);
+
+        let b: R = pk.0.clone() * u.clone() + *m + e0;
+        let d: TR<R, K> = &pk.1 * &u + e1;
+
+        Ok(Self(d, b))
+    }
+    // returns m' not downscaled
+    pub fn decrypt(&self, sk: &SecretKey<R, K>) -> R {
+        let (d, b): (TR<R, K>, R) = (self.0.clone(), self.1);
+        let p: R = b - &d * &sk.0;
+        p
+    }
+}
+
+// Methods for when Ring=Rq<Q,N>
+impl<const Q: u64, const N: usize, const K: usize> GLWE<Rq<Q, N>, K> {
     // scale up
     pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> {
         let m = m.remodule::<Q>();
@@ -87,80 +135,54 @@ impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
         m * delta
     }
     // scale down
-    pub fn decode<const T: u64>(p: &Rq<Q, N>) -> Rq<T, N> {
+    pub fn decode<const T: u64>(m: &Rq<Q, N>) -> Rq<T, N> {
-        let r = p.mul_div_round(T, Q);
+        let r = m.mul_div_round(T, Q);
-        r.remodule::<T>()
+        let r: Rq<T, N> = r.remodule::<T>();
-    }
-
-    // encrypts with the given SecretKey (instead of PublicKey)
-    pub fn encrypt_s(mut rng: impl Rng, sk: &SecretKey<Q, N, K>, m: &Rq<Q, N>) -> Result<Self> {
-        let Xi_key = Uniform::new(0_f64, 2_f64);
-        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
-
-        let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key);
-        let e = Rq::<Q, N>::rand(&mut rng, Xi_err);
-
-        let b: Rq<Q, N> = (&a * &sk.0) + *m + e;
-        Ok(Self(a, b))
-    }
-    pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<Q, N, K>, m: &Rq<Q, N>) -> Result<Self> {
-        let Xi_key = Uniform::new(0_f64, 2_f64);
-        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
-
-        let u: Rq<Q, N> = Rq::rand(&mut rng, Xi_key);
-
-        let e0 = Rq::<Q, N>::rand(&mut rng, Xi_err);
-        let e1 = TR::<Rq<Q, N>, K>::rand(&mut rng, Xi_err);
-
-        let b: Rq<Q, N> = pk.0 * u + *m + e0;
-        let d: TR<Rq<Q, N>, K> = &pk.1 * &u + e1;
-
-        Ok(Self(d, b))
-    }
-    pub fn decrypt(&self, sk: &SecretKey<Q, N, K>) -> Rq<Q, N> {
-        let (d, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1);
-        let r: Rq<Q, N> = b - &d * &sk.0;
         r
     }
-
+    pub fn mod_switch<const P: u64>(&self) -> GLWE<Rq<P, N>, K> {
-    pub fn mod_switch<const P: u64>(&self) -> GLWE<P, N, K> {
+        let a: TR<Rq<P, N>, K> = TR(self
-        let a: TR<Rq<P, N>, K> = TR(self.0 .0.iter().map(|r| r.mod_switch::<P>()).collect());
+            .0
+             .0
+            .iter()
+            .map(|r| r.mod_switch::<P>())
+            .collect::<Vec<_>>());
         let b: Rq<P, N> = self.1.mod_switch::<P>();
         GLWE(a, b)
     }
 }
 
-impl<const Q: u64, const N: usize, const K: usize> Add<GLWE<Q, N, K>> for GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> Add<GLWE<R, K>> for GLWE<R, K> {
     type Output = Self;
     fn add(self, other: Self) -> Self {
-        let a: TR<Rq<Q, N>, K> = self.0 + other.0;
+        let a: TR<R, K> = self.0 + other.0;
-        let b: Rq<Q, N> = self.1 + other.1;
+        let b: R = self.1 + other.1;
         Self(a, b)
     }
 }
 
-impl<const Q: u64, const N: usize, const K: usize> Add<Rq<Q, N>> for GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> Add<R> for GLWE<R, K> {
     type Output = Self;
-    fn add(self, plaintext: Rq<Q, N>) -> Self {
+    fn add(self, plaintext: R) -> Self {
-        let a: TR<Rq<Q, N>, K> = self.0;
+        let a: TR<R, K> = self.0;
-        let b: Rq<Q, N> = self.1 + plaintext;
+        let b: R = self.1 + plaintext;
         Self(a, b)
     }
 }
-impl<const Q: u64, const N: usize, const K: usize> AddAssign for GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> AddAssign for GLWE<R, K> {
     fn add_assign(&mut self, rhs: Self) {
         for i in 0..K {
-            self.0 .0[i] = self.0 .0[i] + rhs.0 .0[i];
+            self.0 .0[i] = self.0 .0[i].clone() + rhs.0 .0[i].clone();
         }
-        self.1 = self.1 + rhs.1;
+        self.1 = self.1.clone() + rhs.1.clone();
     }
 }
-impl<const Q: u64, const N: usize, const K: usize> Sum<GLWE<Q, N, K>> for GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> Sum<GLWE<R, K>> for GLWE<R, K> {
     fn sum<I>(iter: I) -> Self
     where
         I: Iterator<Item = Self>,
     {
-        let mut acc = GLWE::<Q, N, K>::zero();
+        let mut acc = GLWE::<R, K>::zero();
         for e in iter {
             acc += e;
         }
@@ -168,37 +190,60 @@ impl<const Q: u64, const N: usize, const K: usize> Sum<GLWE<Q, N, K>> for GLWE<Q
     }
 }
 
-impl<const Q: u64, const N: usize, const K: usize> Sub<GLWE<Q, N, K>> for GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> Sub<GLWE<R, K>> for GLWE<R, K> {
     type Output = Self;
     fn sub(self, other: Self) -> Self {
-        let a: TR<Rq<Q, N>, K> = self.0 - other.0;
+        let a: TR<R, K> = self.0 - other.0;
-        let b: Rq<Q, N> = self.1 - other.1;
+        let b: R = self.1 - other.1;
         Self(a, b)
     }
 }
 
-impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Q, N, K> {
+impl<R: Ring, const K: usize> Mul<R> for GLWE<R, K> {
     type Output = Self;
-    fn mul(self, plaintext: Rq<Q, N>) -> Self {
+    fn mul(self, plaintext: R) -> Self {
-        // first compute the NTT for plaintext, to avoid computing it at each
+        let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
-        // iteration, speeding up the multiplications
+        let b: R = self.1 * plaintext;
-        let mut plaintext = plaintext.clone();
-        plaintext.compute_evals();
-
-        let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
-        let b: Rq<Q, N> = self.1 * plaintext;
         Self(a, b)
     }
 }
+// for when R = Rq<Q,N>
+// impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Rq<Q, N>, K> {
+//     type Output = Self;
+//     fn mul(self, plaintext: Rq<Q, N>) -> Self {
+//         // first compute the NTT for plaintext, to avoid computing it at each
+//         // iteration, speeding up the multiplications
+//         let mut plaintext = plaintext.clone();
+//         plaintext.compute_evals();
+//
+//         let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
+//         let b: Rq<Q, N> = self.1 * plaintext;
+//         Self(a, b)
+//     }
+// }
 
-impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
+// impl<R: Ring, const K: usize> Mul<R::C> for GLWE<R, K>
-    type Output = Self;
+// // where
-    fn mul(self, e: Zq<Q>) -> Self {
+// //     // R: std::ops::Mul<<R as arith::Ring>::C>,
-        let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
+// //     // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
-        let b: Rq<Q, N> = self.1 * e;
+// //     Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
-        Self(a, b)
+// {
-    }
+//     type Output = Self;
-}
+//     fn mul(self, e: R::C) -> Self {
+//         let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e.clone()).collect());
+//         let b: R = self.1 * e.clone();
+//         Self(a, b)
+//     }
+// }
+
+// impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
+//     type Output = Self;
+//     fn mul(self, e: Zq<Q>) -> Self {
+//         let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
+//         let b: Rq<Q, N> = self.1 * e;
+//         Self(a, b)
+//     }
+// }
 
 #[cfg(test)]
 mod tests {
@@ -213,7 +258,7 @@ mod tests {
         const N: usize = 128;
         const T: u64 = 32; // plaintext modulus
         const K: usize = 16;
-        type S = GLWE<Q, N, K>;
+        type S = GLWE<Rq<Q, N>, K>;
 
         let mut rng = rand::thread_rng();
 
@@ -221,10 +266,11 @@ mod tests {
             let (sk, pk) = S::new_key(&mut rng)?;
 
             let msg_dist = Uniform::new(0_u64, T);
-            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; // msg
-            let p = S::encode::<T>(&m); // plaintext
+                                                               // let m: Rq<Q, N> = m.remodule::<Q>();
 
-            let c = S::encrypt(&mut rng, &pk, &p)?;
+            let p = S::encode::<T>(&m); // plaintext
+            let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
             let p_recovered = c.decrypt(&sk);
             let m_recovered = S::decode::<T>(&p_recovered);
 
@@ -241,13 +287,57 @@ mod tests {
         Ok(())
     }
 
+    use arith::{Tn, T64};
+    use std::array;
+    pub fn t_encode<const P: u64>(m: &Rq<P, 4>) -> Tn<4> {
+        let delta = u64::MAX / P; // floored
+        let coeffs = m.coeffs();
+        Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
+    }
+    pub fn t_decode<const P: u64>(p: &Tn<4>) -> Rq<P, 4> {
+        let p = p.mul_div_round(P, u64::MAX);
+        Rq::<P, 4>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
+    }
+    #[test]
+    fn test_encrypt_decrypt_torus() -> Result<()> {
+        const N: usize = 128;
+        const T: u64 = 32; // plaintext modulus
+        const K: usize = 16;
+        type S = GLWE<Tn<4>, K>;
+
+        let mut rng = rand::thread_rng();
+
+        for _ in 0..200 {
+            let (sk, pk) = S::new_key(&mut rng)?;
+
+            let msg_dist = Uniform::new(0_f64, T as f64);
+            let m = Rq::<T, 4>::rand(&mut rng, msg_dist); // msg
+
+            let p = t_encode::<T>(&m); // plaintext
+            let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
+            let p_recovered = c.decrypt(&sk);
+            let m_recovered = t_decode::<T>(&p_recovered);
+
+            assert_eq!(m, m_recovered);
+
+            // same but using encrypt_s (with sk instead of pk))
+            let c = S::encrypt_s(&mut rng, &sk, &p)?;
+            let p_recovered = c.decrypt(&sk);
+            let m_recovered = t_decode::<T>(&p_recovered);
+
+            assert_eq!(m, m_recovered);
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn test_addition() -> Result<()> {
         const Q: u64 = 2u64.pow(16) + 1;
         const N: usize = 128;
         const T: u64 = 20;
         const K: usize = 16;
-        type S = GLWE<Q, N, K>;
+        type S = GLWE<Rq<Q, N>, K>;
 
         let mut rng = rand::thread_rng();
 
@@ -280,7 +370,7 @@ mod tests {
         const N: usize = 128;
         const T: u64 = 32;
         const K: usize = 16;
-        type S = GLWE<Q, N, K>;
+        type S = GLWE<Rq<Q, N>, K>;
 
         let mut rng = rand::thread_rng();
 
@@ -300,7 +390,7 @@ mod tests {
             let p3_recovered = c3.decrypt(&sk);
             let m3_recovered = S::decode::<T>(&p3_recovered);
 
-            assert_eq!((m1 + m2).remodule::<T>(), m3_recovered);
+            assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
         }
 
         Ok(())
@@ -312,7 +402,7 @@ mod tests {
         const N: usize = 16;
         const T: u64 = 4;
         const K: usize = 16;
-        type S = GLWE<Q, N, K>;
+        type S = GLWE<Rq<Q, N>, K>;
 
         let mut rng = rand::thread_rng();
 
@@ -323,14 +413,14 @@ mod tests {
             let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
             let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
-            let p2: Rq<Q, N> = m2.remodule::<Q>();
+            let p2 = m2.remodule::<Q>(); // notice we don't encode (scale by delta)
 
             let c1 = S::encrypt(&mut rng, &pk, &p1)?;
 
             let c3 = c1 * p2;
 
             let p3_recovered: Rq<Q, N> = c3.decrypt(&sk);
-            let m3_recovered = S::decode::<T>(&p3_recovered);
+            let m3_recovered: Rq<T, N> = S::decode::<T>(&p3_recovered);
             assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
         }
 
@@ -343,35 +433,27 @@ mod tests {
         const P: u64 = 2u64.pow(8) + 1;
         // note: wip, Q and P chosen so that P/Q is an integer
         const N: usize = 8;
-        const T: u64 = 8; // plaintext modulus, must be a prime or power of a prime
+        const T: u64 = 4; // plaintext modulus, must be a prime or power of a prime
         const K: usize = 16;
-        type S = GLWE<Q, N, K>;
+        type S = GLWE<Rq<Q, N>, K>;
 
-        let delta: u64 = Q / T; // floored
         let mut rng = rand::thread_rng();
 
-        dbg!(P as f64 / Q as f64);
-        dbg!(delta);
-        dbg!(delta as f64 * P as f64 / Q as f64);
-        dbg!(delta as f64 * (P as f64 / Q as f64));
-
         for _ in 0..200 {
             let (sk, pk) = S::new_key(&mut rng)?;
 
             let msg_dist = Uniform::new(0_u64, T);
             let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
-            let p = S::encode::<T>(&m); // plaintext
 
+            let p = S::encode::<T>(&m);
             let c = S::encrypt(&mut rng, &pk, &p)?;
-            // let c = S::encrypt_s(&mut rng, &sk, &m, delta)?;
 
-            let c2 = c.mod_switch::<P>();
+            let c2: GLWE<Rq<P, N>, K> = c.mod_switch::<P>();
-            let sk2: SecretKey<P, N, K> =
+            let sk2: SecretKey<Rq<P, N>, K> =
                 SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect()));
-            // let delta2: u64 = ((P as f64 * delta as f64) / Q as f64).round() as u64;
 
             let p_recovered = c2.decrypt(&sk2);
-            let m_recovered = GLWE::<P, N, K>::decode::<T>(&p_recovered);
+            let m_recovered = GLWE::<Rq<P, N>, K>::decode::<T>(&p_recovered);
 
             assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
         }
@@ -385,7 +467,7 @@ mod tests {
         const N: usize = 128;
         const T: u64 = 2; // plaintext modulus
         const K: usize = 16;
-        type S = GLWE<Q, N, K>;
+        type S = GLWE<Rq<Q, N>, K>;
 
         let beta: u32 = 2;
         let l: u32 = 16;
@@ -399,8 +481,8 @@ mod tests {
 
         let msg_dist = Uniform::new(0_u64, T);
         let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
-        let p: Rq<Q, N> = S::encode::<T>(&m); // plaintext
+        let p = S::encode::<T>(&m); // plaintext
-
+                                    //
         let c = S::encrypt_s(&mut rng, &sk, &p)?;
 
         let c2 = c.key_switch(beta, l, &ksk);
@@ -408,14 +490,14 @@ mod tests {
         // decrypt with the 2nd secret key
         let p_recovered = c2.decrypt(&sk2);
         let m_recovered = S::decode::<T>(&p_recovered);
-        assert_eq!(m, m_recovered);
+        assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
 
         // do the same but now encrypting with pk
-        // let c = S::encrypt(&mut rng, &pk, &p)?;
+        let c = S::encrypt(&mut rng, &pk, &p)?;
-        // let c2 = c.key_switch(beta, l, &ksk);
+        let c2 = c.key_switch(beta, l, &ksk);
-        // let p_recovered = c2.decrypt(&sk2);
+        let p_recovered = c2.decrypt(&sk2);
-        // let m_recovered = S::decode::<T>(&p_recovered);
+        let m_recovered = S::decode::<T>(&p_recovered);
-        // assert_eq!(m, m_recovered);
+        assert_eq!(m, m_recovered);
 
         Ok(())
     }

gfhe/src/lib.rs

@@ -7,3 +7,6 @@
 
 pub mod glev;
 pub mod glwe;
+
+pub use glev::GLev;
+pub use glwe::GLWE;

tfhe/Cargo.toml

@@ -10,3 +10,4 @@ rand_distr = { workspace = true }
 itertools = { workspace = true }
 
 arith = { path="../arith" }
+gfhe = { path="../gfhe" }

tfhe/src/lib.rs

@@ -5,5 +5,94 @@
 #![allow(clippy::upper_case_acronyms)]
 #![allow(dead_code)] // TMP
 
+use anyhow::Result;
+use rand::Rng;
+use rand_distr::{Normal, Uniform};
+use std::array;
+
+use arith::{Ring, Rq, Tn, T64};
+use gfhe::{glwe, GLWE};
+
 pub mod tlev;
 pub mod tlwe;
+
+#[derive(Clone, Debug)]
+pub struct SecretKey<const K: usize>(glwe::SecretKey<Tn<1>, K>);
+#[derive(Clone, Debug)]
+pub struct PublicKey<const K: usize>(glwe::PublicKey<Tn<1>, K>);
+
+#[derive(Clone, Debug)]
+pub struct TLWE<const K: usize>(pub GLWE<Tn<1>, K>);
+
+impl<const K: usize> TLWE<K> {
+    pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
+        let (sk, pk) = GLWE::new_key(rng)?;
+        Ok((SecretKey(sk), PublicKey(pk)))
+    }
+
+    pub fn encode<const P: u64>(m: &Rq<P, 1>) -> Tn<1> {
+        let delta = u64::MAX / P; // floored
+        let coeffs = m.coeffs();
+        Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
+    }
+    pub fn decode<const P: u64>(p: &Tn<1>) -> Rq<P, 1> {
+        let p = p.mul_div_round(P, u64::MAX);
+        Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
+    }
+    pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &Tn<1>) -> Result<Self> {
+        let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
+        Ok(Self(glwe))
+    }
+
+    pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &Tn<1>) -> Result<Self> {
+        let glwe = GLWE::encrypt(rng, &pk.0, p)?;
+        Ok(Self(glwe))
+    }
+
+    pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> {
+        self.0.decrypt(&sk.0)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use rand::distributions::Uniform;
+
+    use super::*;
+
+    #[test]
+    fn test_encrypt_decrypt() -> Result<()> {
+        const T: u64 = 128; // plaintext modulus
+        const K: usize = 16;
+        type S = TLWE<K>;
+
+        // let delta: u64 = Q / T; // floored
+        let mut rng = rand::thread_rng();
+
+        for _ in 0..200 {
+            let (sk, pk) = S::new_key(&mut rng)?;
+
+            let msg_dist = Uniform::new(0_f64, T as f64);
+            let m = Rq::<T, 1>::rand(&mut rng, msg_dist); // msg
+
+            // let m: Rq<Q, N> = m.remodule::<Q>();
+
+            let p = S::encode::<T>(&m); // plaintext
+            let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
+            let p_recovered = c.decrypt(&sk);
+            let m_recovered = S::decode::<T>(&p_recovered);
+
+            assert_eq!(m, m_recovered);
+
+            // same but using encrypt_s (with sk instead of pk))
+            let c = S::encrypt_s(&mut rng, &sk, &p)?;
+            let p_recovered = c.decrypt(&sk);
+            let m_recovered = S::decode::<T>(&p_recovered);
+
+            assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
+        }
+
+        Ok(())
+    }
+}