From 7bfcf6f7c118db228bd328487ec8d8e40261aacd Mon Sep 17 00:00:00 2001
From: arnaucube <git@arnaucube.com>
Date: Sun, 3 Aug 2025 07:42:29 +0000
Subject: [PATCH] add TGGSW & TGLev impl

---
 README.md         |   4 +-
 tfhe/src/lib.rs   |   1 +
 tfhe/src/tggsw.rs | 134 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 137 insertions(+), 2 deletions(-)
 create mode 100644 tfhe/src/tggsw.rs

diff --git a/README.md b/README.md
index c9d65d9..80bfffb 100644
--- a/README.md
+++ b/README.md
@@ -36,13 +36,13 @@ let m3 = M::rand_u64(&mut rng, msg_dist)?;
 // encode the msgs into the plaintext space
 let p1 = S::encode::<T>(&m1); // plaintext
 let p2 = S::encode::<T>(&m2); // plaintext
-let c3_const: Tn<1> = Tn(array::from_fn(|i| T64(m3.coeffs()[i].0))); // encode it as constant value
+let c3_const: Tn<1> = Tn(array::from_fn(|i| T64(m3.coeffs()[i].0))); // encode it as constant
 
 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
 let c2 = S::encrypt(&mut rng, &pk, &p2)?;
 
 // now we can do encrypted operations (notice that we do them using simple
-// operations by operator overloading):
+// operation notation by rust's operator overloading):
 let c_12 = c1 + c2;
 let c4 = c_12 * c3_const;
 
diff --git a/tfhe/src/lib.rs b/tfhe/src/lib.rs
index 9db005b..f2f32fa 100644
--- a/tfhe/src/lib.rs
+++ b/tfhe/src/lib.rs
@@ -5,6 +5,7 @@
 #![allow(clippy::upper_case_acronyms)]
 #![allow(dead_code)] // TMP
 
+pub mod tggsw;
 pub mod tglwe;
 pub mod tgsw;
 pub mod tlev;
diff --git a/tfhe/src/tggsw.rs b/tfhe/src/tggsw.rs
new file mode 100644
index 0000000..fac7c22
--- /dev/null
+++ b/tfhe/src/tggsw.rs
@@ -0,0 +1,134 @@
+use anyhow::Result;
+use itertools::zip_eq;
+use rand::Rng;
+use std::array;
+use std::ops::{Add, Mul};
+
+use arith::{Ring, Rq, Tn, T64, TR};
+
+use crate::tglwe::{PublicKey, SecretKey, TGLWE};
+use gfhe::glwe::GLWE;
+
+/// vector of length K+1 = ([K * TGLev], [1 * TGLev])
+#[derive(Clone, Debug)]
+pub struct TGGSW<const N: usize, const K: usize>(pub(crate) Vec<TGLev<N, K>>, TGLev<N, K>);
+
+impl<const N: usize, const K: usize> TGGSW<N, K> {
+    pub fn encrypt_s(
+        mut rng: impl Rng,
+        beta: u32,
+        l: u32,
+        sk: &SecretKey<N, K>,
+        m: &Tn<N>,
+    ) -> Result<Self> {
+        let a: Vec<TGLev<N, K>> = (0..K)
+            .map(|i| TGLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0 .0[i] * *m)))
+            .collect::<Result<Vec<_>>>()?;
+        let b: TGLev<N, K> = TGLev::encrypt_s(&mut rng, beta, l, sk, m)?;
+        Ok(Self(a, b))
+    }
+
+    pub fn decrypt(&self, sk: &SecretKey<N, K>, beta: u32) -> Tn<N> {
+        self.1.decrypt(sk, beta)
+    }
+
+    pub fn cmux(bit: Self, ct1: TGLWE<N, K>, ct2: TGLWE<N, K>) -> TGLWE<N, K> {
+        ct1.clone() + (bit * (ct2 - ct1))
+    }
+}
+
+/// External product TGGSW x TGLWE
+impl<const N: usize, const K: usize> Mul<TGLWE<N, K>> for TGGSW<N, K> {
+    type Output = TGLWE<N, K>;
+
+    fn mul(self, tglwe: TGLWE<N, K>) -> TGLWE<N, K> {
+        let beta: u32 = 2;
+        let l: u32 = 64; // TODO wip
+
+        let tglwe_ab: Vec<Tn<N>> = [tglwe.0 .0 .0.clone(), vec![tglwe.0 .1]].concat();
+
+        let tgsw_ab: Vec<TGLev<N, K>> = [self.0.clone(), vec![self.1]].concat();
+        assert_eq!(tgsw_ab.len(), tglwe_ab.len());
+
+        let r: TGLWE<N, K> = zip_eq(tgsw_ab, tglwe_ab)
+            .map(|(tlev_i, tglwe_i)| tlev_i * tglwe_i.decompose(beta, l))
+            .sum();
+        r
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct TGLev<const N: usize, const K: usize>(pub(crate) Vec<TGLWE<N, K>>);
+
+impl<const N: usize, const K: usize> TGLev<N, K> {
+    pub fn encode<const T: u64>(m: &Rq<T, N>) -> Tn<N> {
+        let coeffs = m.coeffs();
+        Tn(array::from_fn(|i| T64(coeffs[i].0)))
+    }
+    pub fn decode<const T: u64>(p: &Tn<N>) -> Rq<T, N> {
+        Rq::<T, N>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
+    }
+    pub fn encrypt(
+        mut rng: impl Rng,
+        beta: u32,
+        l: u32,
+        pk: &PublicKey<N, K>,
+        m: &Tn<N>,
+    ) -> Result<Self> {
+        let tlev: Vec<TGLWE<N, K>> = (1..l + 1)
+            .map(|i| {
+                TGLWE::<N, K>::encrypt(&mut rng, pk, &(*m * (u64::MAX / beta.pow(i as u32) as u64)))
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        Ok(Self(tlev))
+    }
+    pub fn encrypt_s(
+        mut rng: impl Rng,
+        _beta: u32, // TODO rm, and make beta=2 always
+        l: u32,
+        sk: &SecretKey<N, K>,
+        m: &Tn<N>,
+    ) -> Result<Self> {
+        let tlev: Vec<TGLWE<N, K>> = (1..l as u64 + 1)
+            .map(|i| {
+                let aux = if i < 64 {
+                    *m * (u64::MAX / (1u64 << i))
+                } else {
+                    // 1<<64 would overflow, and anyways we're dividing u64::MAX
+                    // by it, which would be equal to 1
+                    *m
+                };
+                TGLWE::<N, K>::encrypt_s(&mut rng, sk, &aux)
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        Ok(Self(tlev))
+    }
+
+    pub fn decrypt(&self, sk: &SecretKey<N, K>, beta: u32) -> Tn<N> {
+        let pt = self.0[0].decrypt(sk);
+        pt.mul_div_round(beta as u64, u64::MAX)
+    }
+}
+
+impl<const N: usize, const K: usize> TGLev<N, K> {
+    pub fn iter(&self) -> std::slice::Iter<TGLWE<N, K>> {
+        self.0.iter()
+    }
+}
+
+// dot product between a TGLev and Vec<Tn<N>>, usually Vec<Tn<N>> comes from a
+// decomposition of Tn<N>
+// TGLev * Vec<Tn<N>> --> TGLWE
+impl<const N: usize, const K: usize> Mul<Vec<Tn<N>>> for TGLev<N, K> {
+    type Output = TGLWE<N, K>;
+    fn mul(self, v: Vec<Tn<N>>) -> Self::Output {
+        assert_eq!(self.0.len(), v.len());
+
+        // l TGLWES
+        let tlwes: Vec<TGLWE<N, K>> = self.0;
+        let r: TGLWE<N, K> = zip_eq(v, tlwes).map(|(a_d_i, glwe_i)| glwe_i * a_d_i).sum();
+        r
+    }
+}