1use anyhow::Result;
  2use itertools::zip_eq;
  3use rand::distributions::Standard;
  4use rand::Rng;
  5use rand_distr::{Normal, Uniform};
  6use std::array;
  7use std::iter::Sum;
  8use std::ops::{Add, AddAssign, Mul, Sub};
  9
 10use arith::{Ring, Rq, Tn, T64, TR};
 11use gfhe::{glwe, GLWE};
 12
 13use crate::tlev::TLev;
 14use crate::{tlwe, tlwe::TLWE};
 15
 16// pub type SecretKey<const N: usize, const K: usize> = glwe::SecretKey<Tn<N>, K>;
 17#[derive(Clone, Debug)]
 18pub struct SecretKey<const N: usize, const K: usize>(pub glwe::SecretKey<Tn<N>, K>);
 19// pub struct SecretKey<const K: usize>(pub tlwe::SecretKey<K>);
 20
 21impl<const N: usize, const K: usize> SecretKey<N, K> {
 22    pub fn to_tlwe<const KN: usize>(self) -> tlwe::SecretKey<K> {
 23        let s: TR<Tn<N>, K> = self.0 .0;
 24
 25        let r: Vec<Vec<T64>> = s.0.iter().map(|s_i| s_i.coeffs()).collect();
 26        let r: Vec<T64> = r.into_iter().flatten().collect();
 27        tlwe::SecretKey(glwe::SecretKey(TR(r)))
 28    }
 29}
 30
 31pub type PublicKey<const N: usize, const K: usize> = glwe::PublicKey<Tn<N>, K>;
 32
 33#[derive(Clone, Debug)]
 34pub struct TGLWE<const N: usize, const K: usize>(pub GLWE<Tn<N>, K>);
 35
 36impl<const N: usize, const K: usize> TGLWE<N, K> {
 37    pub fn zero() -> Self {
 38        Self(GLWE::<Tn<N>, K>::zero())
 39    }
 40
 41    pub fn new_key<const KN: usize>(
 42        mut rng: impl Rng,
 43    ) -> Result<(SecretKey<N, K>, PublicKey<N, K>)> {
 44        assert_eq!(KN, K * N); // this is wip, while not being able to compute K*N
 45        let (sk_tlwe, _) = TLWE::<KN>::new_key(&mut rng)?;
 46        // let sk = crate::tlwe::sk_to_tglwe::<N, K, KN>(sk_tlwe);
 47        let sk = sk_tlwe.to_tglwe::<N, K>();
 48        let pk = GLWE::pk_from_sk(rng, sk.0.clone())?;
 49        Ok((sk, pk))
 50    }
 51
 52    pub fn encode<const P: u64>(m: &Rq<P, N>) -> Tn<N> {
 53        let delta = u64::MAX / P; // floored
 54        let coeffs = m.coeffs();
 55        Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
 56    }
 57    pub fn decode<const P: u64>(p: &Tn<N>) -> Rq<P, N> {
 58        let p = p.mul_div_round(P, u64::MAX);
 59        Rq::<P, N>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
 60    }
 61
 62    // encrypts with the given SecretKey (instead of PublicKey)
 63    pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<N, K>, p: &Tn<N>) -> Result<Self> {
 64        let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
 65        Ok(Self(glwe))
 66    }
 67    pub fn encrypt(rng: impl Rng, pk: &PublicKey<N, K>, p: &Tn<N>) -> Result<Self> {
 68        let glwe = GLWE::encrypt(rng, &pk, p)?;
 69        Ok(Self(glwe))
 70    }
 71    pub fn decrypt(&self, sk: &SecretKey<N, K>) -> Tn<N> {
 72        self.0.decrypt(&sk.0)
 73    }
 74
 75    /// Sample extraction / Coefficient extraction
 76    pub fn sample_extraction(&self, h: usize) -> TLWE<K> {
 77        assert!(h < N);
 78
 79        let a: TR<Tn<N>, K> = self.0 .0.clone();
 80        // set a_{n*i+j} = a_{i, h-j}     if j \in {0, h}
 81        //                 -a_{i, n+h-j}  if j \in {h+1, n-1}
 82        let new_a: Vec<T64> = a
 83            .iter()
 84            .flat_map(|a_i| {
 85                let a_i = a_i.coeffs();
 86                (0..N)
 87                    .map(|j| if j <= h { a_i[h - j] } else { -a_i[N + h - j] })
 88                    .collect::<Vec<T64>>()
 89            })
 90            .collect::<Vec<T64>>();
 91
 92        TLWE(GLWE(TR(new_a), self.0 .1.coeffs()[h]))
 93    }
 94}
 95
 96impl<const N: usize, const K: usize> Add<TGLWE<N, K>> for TGLWE<N, K> {
 97    type Output = Self;
 98    fn add(self, other: Self) -> Self {
 99        Self(self.0 + other.0)
100    }
101}
102impl<const N: usize, const K: usize> AddAssign for TGLWE<N, K> {
103    fn add_assign(&mut self, rhs: Self) {
104        self.0 += rhs.0
105    }
106}
107impl<const N: usize, const K: usize> Sum<TGLWE<N, K>> for TGLWE<N, K> {
108    fn sum<I>(iter: I) -> Self
109    where
110        I: Iterator<Item = Self>,
111    {
112        let mut acc = TGLWE::<N, K>::zero();
113        for e in iter {
114            acc += e;
115        }
116        acc
117    }
118}
119
120impl<const N: usize, const K: usize> Sub<TGLWE<N, K>> for TGLWE<N, K> {
121    type Output = Self;
122    fn sub(self, other: Self) -> Self {
123        Self(self.0 - other.0)
124    }
125}
126
127// plaintext addition
128impl<const N: usize, const K: usize> Add<Tn<N>> for TGLWE<N, K> {
129    type Output = Self;
130    fn add(self, plaintext: Tn<N>) -> Self {
131        let a: TR<Tn<N>, K> = self.0 .0;
132        let b: Tn<N> = self.0 .1 + plaintext;
133        Self(GLWE(a, b))
134    }
135}
136// plaintext substraction
137impl<const N: usize, const K: usize> Sub<Tn<N>> for TGLWE<N, K> {
138    type Output = Self;
139    fn sub(self, plaintext: Tn<N>) -> Self {
140        let a: TR<Tn<N>, K> = self.0 .0;
141        let b: Tn<N> = self.0 .1 - plaintext;
142        Self(GLWE(a, b))
143    }
144}
145// plaintext multiplication
146impl<const N: usize, const K: usize> Mul<Tn<N>> for TGLWE<N, K> {
147    type Output = Self;
148    fn mul(self, plaintext: Tn<N>) -> Self {
149        let a: TR<Tn<N>, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect());
150        let b: Tn<N> = self.0 .1 * plaintext;
151        Self(GLWE(a, b))
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use anyhow::Result;
158    use rand::distributions::Uniform;
159
160    use super::*;
161
162    #[test]
163    fn test_encrypt_decrypt() -> Result<()> {
164        const T: u64 = 128; // msg space (msg modulus)
165        const N: usize = 64;
166        const K: usize = 16;
167        type S = TGLWE<N, K>;
168
169        let mut rng = rand::thread_rng();
170        let msg_dist = Uniform::new(0_u64, T);
171
172        for _ in 0..200 {
173            let (sk, pk) = TGLWE::<N, K>::new_key::<{ K * N }>(&mut rng)?;
174
175            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
176            let p: Tn<N> = S::encode::<T>(&m);
177
178            let c = S::encrypt(&mut rng, &pk, &p)?;
179            let p_recovered = c.decrypt(&sk);
180            let m_recovered = S::decode::<T>(&p_recovered);
181
182            assert_eq!(m, m_recovered);
183
184            // same but using encrypt_s (with sk instead of pk))
185            let c = S::encrypt_s(&mut rng, &sk, &p)?;
186            let p_recovered = c.decrypt(&sk);
187            let m_recovered = S::decode::<T>(&p_recovered);
188
189            assert_eq!(m, m_recovered);
190        }
191
192        Ok(())
193    }
194
195    #[test]
196    fn test_addition() -> Result<()> {
197        const T: u64 = 128;
198        const N: usize = 64;
199        const K: usize = 16;
200        type S = TGLWE<N, K>;
201
202        let mut rng = rand::thread_rng();
203        let msg_dist = Uniform::new(0_u64, T);
204
205        for _ in 0..200 {
206            let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
207
208            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
209            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
210            let p1: Tn<N> = S::encode::<T>(&m1); // plaintext
211            let p2: Tn<N> = S::encode::<T>(&m2); // plaintext
212
213            let c1 = S::encrypt(&mut rng, &pk, &p1)?;
214            let c2 = S::encrypt(&mut rng, &pk, &p2)?;
215
216            let c3 = c1 + c2;
217
218            let p3_recovered = c3.decrypt(&sk);
219            let m3_recovered = S::decode::<T>(&p3_recovered);
220
221            assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
222        }
223
224        Ok(())
225    }
226
227    #[test]
228    fn test_add_plaintext() -> Result<()> {
229        const T: u64 = 128;
230        const N: usize = 64;
231        const K: usize = 16;
232        type S = TGLWE<N, K>;
233
234        let mut rng = rand::thread_rng();
235        let msg_dist = Uniform::new(0_u64, T);
236
237        for _ in 0..200 {
238            let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
239
240            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
241            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
242            let p1: Tn<N> = S::encode::<T>(&m1); // plaintext
243            let p2: Tn<N> = S::encode::<T>(&m2); // plaintext
244
245            let c1 = S::encrypt(&mut rng, &pk, &p1)?;
246
247            let c3 = c1 + p2;
248
249            let p3_recovered = c3.decrypt(&sk);
250            let m3_recovered = S::decode::<T>(&p3_recovered);
251
252            assert_eq!(m1 + m2, m3_recovered);
253        }
254
255        Ok(())
256    }
257
258    #[test]
259    fn test_mul_plaintext() -> Result<()> {
260        const T: u64 = 128;
261        const N: usize = 64;
262        const K: usize = 16;
263        type S = TGLWE<N, K>;
264
265        let mut rng = rand::thread_rng();
266        let msg_dist = Uniform::new(0_u64, T);
267
268        for _ in 0..200 {
269            let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
270
271            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
272            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
273            let p1: Tn<N> = S::encode::<T>(&m1);
274            // don't scale up p2, set it directly from m2
275            let p2: Tn<N> = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
276
277            let c1 = S::encrypt(&mut rng, &pk, &p1)?;
278
279            let c3 = c1 * p2;
280
281            let p3_recovered: Tn<N> = c3.decrypt(&sk);
282            let m3_recovered = S::decode::<T>(&p3_recovered);
283            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
284        }
285
286        Ok(())
287    }
288
289    #[test]
290    fn test_sample_extraction() -> Result<()> {
291        const T: u64 = 128; // msg space (msg modulus)
292        const N: usize = 64;
293        const K: usize = 16;
294
295        let mut rng = rand::thread_rng();
296        let msg_dist = Uniform::new(0_u64, T);
297
298        for _ in 0..20 {
299            let (sk, pk) = TGLWE::<N, K>::new_key::<{ K * N }>(&mut rng)?;
300            let sk_tlwe = sk.to_tlwe::<{ K * N }>();
301
302            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
303            let p: Tn<N> = TGLWE::<N, K>::encode::<T>(&m);
304
305            let c = TGLWE::<N, K>::encrypt(&mut rng, &pk, &p)?;
306
307            for h in 0..N {
308                let c_h: TLWE<K> = c.sample_extraction(h);
309
310                let p_recovered = c_h.decrypt(&sk_tlwe);
311                let m_recovered = TLWE::<K>::decode::<T>(&p_recovered);
312                assert_eq!(m.coeffs()[h], m_recovered.coeffs()[0]);
313            }
314        }
315
316        Ok(())
317    }
318}