1use anyhow::Result;
  2use itertools::zip_eq;
  3use rand::Rng;
  4use std::iter::Sum;
  5use std::ops::{Add, AddAssign, Mul, Sub};
  6
  7use arith::{Ring, RingParam, Rq, Tn, Zq, T64, TR};
  8use gfhe::{glwe, glwe::Param, GLWE};
  9
 10use crate::tggsw::TGGSW;
 11use crate::tlev::TLev;
 12use crate::{tglwe, tglwe::TGLWE};
 13
 14pub struct SecretKey(pub glwe::SecretKey<T64>);
 15
 16impl SecretKey {
 17    /// from TFHE [2018-421] paper: A TLWE key k \in B^n, can be interpreted as a
 18    /// TRLWE key K \in B_N[X]^k having the same sequence of coefficients and
 19    /// vice-versa.
 20    pub fn to_tglwe(self, param: &Param) -> crate::tglwe::SecretKey {
 21        let s: TR<T64> = self.0 .0; // of length K*N
 22        assert_eq!(s.r.len(), param.k * param.ring.n); // sanity check
 23
 24        // split into K vectors, and interpret each of them as a T_N[X]/(X^N+1)
 25        // polynomial
 26        let r: Vec<Tn> =
 27            s.r.chunks(param.ring.n)
 28                .map(|v| Tn::from_vec(&param.ring, v.to_vec()))
 29                .collect();
 30        crate::tglwe::SecretKey(glwe::SecretKey::<Tn>(TR { k: param.k, r }))
 31    }
 32}
 33
 34pub type PublicKey = glwe::PublicKey<T64>;
 35
 36#[derive(Clone, Debug)]
 37pub struct KSK(Vec<TLev>);
 38
 39#[derive(Clone, Debug)]
 40pub struct TLWE(pub GLWE<T64>);
 41
 42impl TLWE {
 43    pub fn zero(k: usize, ring_param: &RingParam) -> Self {
 44        Self(GLWE::<T64>::zero(k, ring_param))
 45    }
 46
 47    pub fn new_key(rng: impl Rng, param: &Param) -> Result<(SecretKey, PublicKey)> {
 48        let (sk, pk): (glwe::SecretKey<T64>, glwe::PublicKey<T64>) = GLWE::new_key(rng, param)?;
 49        Ok((SecretKey(sk), pk))
 50    }
 51
 52    // TODO use param instead of p:u64
 53    pub fn encode(param: &Param, m: &Rq) -> T64 {
 54        assert_eq!(param.ring.n, 1);
 55        debug_assert_eq!(param.t, m.param.q); // plaintext modulus
 56                                              //
 57        let delta = u64::MAX / param.t; // floored
 58        let coeffs = m.coeffs();
 59        T64(coeffs[0].v * delta)
 60    }
 61    pub fn decode(param: &Param, p: &T64) -> Rq {
 62        let p = p.mul_div_round(param.t, u64::MAX);
 63        Rq::from_vec_u64(&param.pt(), p.coeffs().iter().map(|c| c.0).collect())
 64    }
 65
 66    // encrypts with the given SecretKey (instead of PublicKey)
 67    pub fn encrypt_s(rng: impl Rng, param: &Param, sk: &SecretKey, p: &T64) -> Result<Self> {
 68        let glwe = GLWE::encrypt_s(rng, param, &sk.0, p)?;
 69        Ok(Self(glwe))
 70    }
 71    pub fn encrypt(rng: impl Rng, param: &Param, pk: &PublicKey, p: &T64) -> Result<Self> {
 72        let glwe = GLWE::encrypt(rng, param, pk, p)?;
 73        Ok(Self(glwe))
 74    }
 75    pub fn decrypt(&self, sk: &SecretKey) -> T64 {
 76        self.0.decrypt(&sk.0)
 77    }
 78
 79    pub fn new_ksk(
 80        mut rng: impl Rng,
 81        param: &Param,
 82        beta: u32,
 83        l: u32,
 84        sk: &SecretKey,
 85        new_sk: &SecretKey,
 86    ) -> Result<KSK> {
 87        let r: Vec<TLev> = (0..param.k)
 88            .into_iter()
 89            .map(|i|
 90                // treat sk_i as the msg being encrypted
 91                TLev::encrypt_s(&mut rng, param, beta, l, &new_sk, &sk.0.0 .r[i]))
 92            .collect::<Result<Vec<_>>>()?;
 93
 94        Ok(KSK(r))
 95    }
 96    pub fn key_switch(&self, param: &Param, beta: u32, l: u32, ksk: &KSK) -> Self {
 97        let (a, b): (TR<T64>, T64) = (self.0 .0.clone(), self.0 .1);
 98
 99        let lhs: TLWE = TLWE(GLWE(TR::zero(param.k * param.ring.n, &param.ring), b));
100
101        // K iterations, ksk.0 contains K times GLev
102        let rhs: TLWE = zip_eq(a.r, ksk.0.clone())
103            .map(|(a_i, ksk_i)| ksk_i * a_i.decompose(beta, l)) // dot_product
104            .sum();
105
106        lhs - rhs
107    }
108    // modulus switch from Q (2^64) to Q2 (in blind_rotation Q2=K*N)
109    pub fn mod_switch(&self, q2: u64) -> Self {
110        let a: TR<T64> = self.0 .0.mod_switch(q2);
111        let b: T64 = self.0 .1.mod_switch(q2);
112        Self(GLWE(a, b))
113    }
114}
115// NOTE: the ugly const generics are temporary
116pub fn blind_rotation(
117    param: &Param,
118    c: TLWE, // kn
119    btk: BootstrappingKey,
120    table: TGLWE, // n,k
121) -> TGLWE {
122    debug_assert_eq!(c.0 .0.k, param.k); // TODO or k*n?
123
124    // TODO replace `param.k*param.ring.n` by `param.kn()`
125    let c_kn: TLWE = c.mod_switch((param.k * param.ring.n) as u64);
126    let (a, b): (TR<T64>, T64) = (c_kn.0 .0, c_kn.0 .1);
127    // two main parts: rotate by a known power of X, rotate by a secret
128    // power of X (using the C gate)
129
130    // table * X^-b, ie. left rotate
131    let v_xb: TGLWE = table.left_rotate(b.0 as usize);
132
133    // rotate by a secret power of X using the cmux gate
134    let mut c_j: TGLWE = v_xb.clone();
135    let _ = (1..param.k).map(|j| {
136        c_j = TGGSW::cmux(
137            btk.0[j].clone(),
138            c_j.clone(),
139            c_j.clone().left_rotate(a.r[j].0 as usize),
140        );
141        // dbg!(&c_j);
142    });
143    c_j
144}
145
146pub fn bootstrapping(
147    param: &Param,
148    btk: BootstrappingKey,
149    table: TGLWE,
150    c: TLWE, // kn
151) -> TLWE {
152    // kn
153    let rotated: TGLWE = blind_rotation(param, c, btk.clone(), table);
154    let c_h: TLWE = rotated.sample_extraction(&param, 0);
155    let r = c_h.key_switch(param, 2, 64, &btk.1);
156    r
157}
158
159#[derive(Clone, Debug)]
160pub struct BootstrappingKey(
161    pub Vec<TGGSW>,
162    pub KSK, // kn
163);
164impl BootstrappingKey {
165    pub fn from_sk(mut rng: impl Rng, param: &Param, sk: &tglwe::SecretKey) -> Result<Self> {
166        let (beta, l) = (2u32, 64u32); // TMP
167                                       //
168        let s: TR<Tn> = sk.0 .0.clone();
169        let (sk2, _) = TLWE::new_key(&mut rng, &param.lwe())?; // TLWE<KN> compatible with TGLWE<N,K>
170
171        // each btk_j = TGGSW_sk(s_i)
172        let btk: Vec<TGGSW> = s
173            .iter()
174            .map(|s_i| TGGSW::encrypt_s(&mut rng, param, beta, l, sk, s_i))
175            .collect::<Result<Vec<_>>>()?;
176
177        let ksk = TLWE::new_ksk(
178            &mut rng,
179            &param.lwe(),
180            beta,
181            l,
182            &sk.to_tlwe(&param.lwe()), // converted to length k*n
183            &sk2,                      // created with length k*n
184        )?;
185        debug_assert_eq!(ksk.0.len(), param.lwe().k);
186        debug_assert_eq!(ksk.0.len(), param.k * param.ring.n);
187
188        Ok(Self(btk, ksk))
189    }
190}
191
192pub fn compute_lookup_table(param: &Param) -> TGLWE {
193    // from 2021-1402:
194    // v(x) = \sum_j^{N-1} [(p_j / 2N mod p)/p] X^j
195
196    // matrix of coefficients with size K*N = delta x T
197    let delta: usize = param.ring.n / param.t as usize;
198    let values: Vec<Zq> = (0..param.t).map(|v| Zq::from_u64(param.t, v)).collect();
199    let coeffs: Vec<Zq> = (0..param.t as usize)
200        .flat_map(|i| vec![values[i]; delta])
201        .collect();
202    let table = Rq::from_vec(&param.pt(), coeffs);
203
204    // encode the table as plaintext
205    let v: Tn = TGLWE::encode(param, &table);
206
207    // encode the table as TGLWE ciphertext
208    let v: TGLWE = TGLWE::from_plaintext(param.k, &param.ring, v);
209    v
210}
211
212impl Add<TLWE> for TLWE {
213    type Output = Self;
214    fn add(self, other: Self) -> Self {
215        debug_assert_eq!(self.0 .0.k, other.0 .0.k);
216        debug_assert_eq!(self.0 .1.param(), other.0 .1.param());
217        Self(self.0 + other.0)
218    }
219}
220impl AddAssign for TLWE {
221    fn add_assign(&mut self, rhs: Self) {
222        debug_assert_eq!(self.0 .0.k, rhs.0 .0.k);
223        debug_assert_eq!(self.0 .1.param(), rhs.0 .1.param());
224        self.0 += rhs.0
225    }
226}
227impl Sum<TLWE> for TLWE {
228    fn sum<I>(mut iter: I) -> Self
229    where
230        I: Iterator<Item = Self>,
231    {
232        // let mut acc = TLWE::<K>::zero();
233        // for e in iter {
234        //     acc += e;
235        // }
236        // acc
237        let first = iter.next().unwrap();
238        iter.fold(first, |acc, e| acc + e)
239    }
240}
241
242impl Sub<TLWE> for TLWE {
243    type Output = Self;
244    fn sub(self, other: Self) -> Self {
245        debug_assert_eq!(self.0 .0.k, other.0 .0.k);
246        debug_assert_eq!(self.0 .1.param(), other.0 .1.param());
247        Self(self.0 - other.0)
248    }
249}
250
251// plaintext addition
252impl Add<T64> for TLWE {
253    type Output = Self;
254    fn add(self, plaintext: T64) -> Self {
255        let a: TR<T64> = self.0 .0;
256        let b: T64 = self.0 .1 + plaintext;
257        Self(GLWE(a, b))
258    }
259}
260// plaintext substraction
261impl Sub<T64> for TLWE {
262    type Output = Self;
263    fn sub(self, plaintext: T64) -> Self {
264        let a: TR<T64> = self.0 .0;
265        let b: T64 = self.0 .1 - plaintext;
266        Self(GLWE(a, b))
267    }
268}
269// plaintext multiplication
270impl Mul<T64> for TLWE {
271    type Output = Self;
272    fn mul(self, plaintext: T64) -> Self {
273        let a: TR<T64> = TR {
274            k: self.0 .0.k,
275            r: self.0 .0.r.iter().map(|r_i| *r_i * plaintext).collect(),
276        };
277        let b: T64 = self.0 .1 * plaintext;
278        Self(GLWE(a, b))
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use anyhow::Result;
285    use rand::distributions::Uniform;
286    use std::time::Instant;
287
288    use super::*;
289
290    #[test]
291    fn test_encrypt_decrypt() -> Result<()> {
292        let param = Param {
293            ring: RingParam { q: u64::MAX, n: 1 },
294            k: 16,
295            t: 128, // plaintext modulus
296        };
297
298        let mut rng = rand::thread_rng();
299        let msg_dist = Uniform::new(0_u64, param.t);
300
301        for _ in 0..200 {
302            let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
303
304            let m = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
305            let p: T64 = TLWE::encode(&param, &m);
306
307            let c = TLWE::encrypt(&mut rng, &param, &pk, &p)?;
308            let p_recovered = c.decrypt(&sk);
309            let m_recovered = TLWE::decode(&param, &p_recovered);
310
311            assert_eq!(m, m_recovered);
312
313            // same but using encrypt_s (with sk instead of pk))
314            let c = TLWE::encrypt_s(&mut rng, &param, &sk, &p)?;
315            let p_recovered = c.decrypt(&sk);
316            let m_recovered = TLWE::decode(&param, &p_recovered);
317
318            assert_eq!(m, m_recovered);
319        }
320
321        Ok(())
322    }
323
324    #[test]
325    fn test_addition() -> Result<()> {
326        let param = Param {
327            ring: RingParam { q: u64::MAX, n: 1 },
328            k: 16,
329            t: 128, // plaintext modulus
330        };
331
332        let mut rng = rand::thread_rng();
333        let msg_dist = Uniform::new(0_u64, param.t);
334
335        for _ in 0..200 {
336            let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
337
338            let m1 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
339            let m2 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
340            let p1: T64 = TLWE::encode(&param, &m1); // plaintext
341            let p2: T64 = TLWE::encode(&param, &m2); // plaintext
342
343            let c1 = TLWE::encrypt(&mut rng, &param, &pk, &p1)?;
344            let c2 = TLWE::encrypt(&mut rng, &param, &pk, &p2)?;
345
346            let c3 = c1 + c2;
347
348            let p3_recovered = c3.decrypt(&sk);
349            let m3_recovered = TLWE::decode(&param, &p3_recovered);
350
351            assert_eq!((m1 + m2).remodule(param.t), m3_recovered.remodule(param.t));
352        }
353
354        Ok(())
355    }
356
357    #[test]
358    fn test_add_plaintext() -> Result<()> {
359        let param = Param {
360            ring: RingParam { q: u64::MAX, n: 1 },
361            k: 16,
362            t: 128, // plaintext modulus
363        };
364
365        let mut rng = rand::thread_rng();
366        let msg_dist = Uniform::new(0_u64, param.t);
367
368        for _ in 0..200 {
369            let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
370
371            let m1 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
372            let m2 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
373            let p1: T64 = TLWE::encode(&param, &m1); // plaintext
374            let p2: T64 = TLWE::encode(&param, &m2); // plaintext
375
376            let c1 = TLWE::encrypt(&mut rng, &param, &pk, &p1)?;
377
378            let c3 = c1 + p2;
379
380            let p3_recovered = c3.decrypt(&sk);
381            let m3_recovered = TLWE::decode(&param, &p3_recovered);
382
383            assert_eq!(m1 + m2, m3_recovered);
384        }
385
386        Ok(())
387    }
388
389    #[test]
390    fn test_mul_plaintext() -> Result<()> {
391        let param = Param {
392            ring: RingParam { q: u64::MAX, n: 1 },
393            k: 16,
394            t: 128, // plaintext modulus
395        };
396
397        let mut rng = rand::thread_rng();
398        let msg_dist = Uniform::new(0_u64, param.t);
399
400        for _ in 0..200 {
401            let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
402
403            let m1 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
404            let m2 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
405            let p1: T64 = TLWE::encode(&param, &m1);
406            // don't scale up p2, set it directly from m2
407            // let p2: T64 = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
408            let p2: T64 = T64(m2.coeffs()[0].v);
409
410            let c1 = TLWE::encrypt(&mut rng, &param, &pk, &p1)?;
411
412            let c3 = c1 * p2;
413
414            let p3_recovered: T64 = c3.decrypt(&sk);
415            let m3_recovered = TLWE::decode(&param, &p3_recovered);
416            assert_eq!((m1.to_r() * m2.to_r()).to_rq(param.t), m3_recovered);
417        }
418
419        Ok(())
420    }
421
422    #[test]
423    fn test_key_switch() -> Result<()> {
424        let param = Param {
425            ring: RingParam { q: u64::MAX, n: 1 },
426            k: 16,
427            t: 128, // plaintext modulus
428        };
429
430        let beta: u32 = 2;
431        let l: u32 = 64;
432
433        let mut rng = rand::thread_rng();
434
435        let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
436        let (sk2, _) = TLWE::new_key(&mut rng, &param)?;
437        // ksk to switch from sk to sk2
438        let ksk = TLWE::new_ksk(&mut rng, &param, beta, l, &sk, &sk2)?;
439
440        let msg_dist = Uniform::new(0_u64, param.t);
441        let m = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
442        let p = TLWE::encode(&param, &m); // plaintext
443                                          //
444        let c = TLWE::encrypt_s(&mut rng, &param, &sk, &p)?;
445
446        let c2 = c.key_switch(&param, beta, l, &ksk);
447
448        // decrypt with the 2nd secret key
449        let p_recovered = c2.decrypt(&sk2);
450        let m_recovered = TLWE::decode(&param, &p_recovered);
451        assert_eq!(m.remodule(param.t), m_recovered.remodule(param.t));
452
453        // do the same but now encrypting with pk
454        let c = TLWE::encrypt(&mut rng, &param, &pk, &p)?;
455        let c2 = c.key_switch(&param, beta, l, &ksk);
456        let p_recovered = c2.decrypt(&sk2);
457        let m_recovered = TLWE::decode(&param, &p_recovered);
458        assert_eq!(m, m_recovered);
459
460        Ok(())
461    }
462
463    #[test]
464    fn test_bootstrapping() -> Result<()> {
465        let param = Param {
466            ring: RingParam {
467                q: u64::MAX,
468                n: 1024,
469            },
470            k: 1,
471            t: 128, // plaintext modulus
472        };
473        // const T: u64 = 128; // plaintext modulus
474        // const K: usize = 1;
475        // const N: usize = 1024;
476        // const KN: usize = K * N;
477        let mut rng = rand::thread_rng();
478
479        let start = Instant::now();
480        let table: TGLWE = compute_lookup_table(&param);
481        println!("table took: {:?}", start.elapsed());
482
483        let (sk, _) = TGLWE::new_key(&mut rng, &param)?;
484        let sk_tlwe: SecretKey = sk.to_tlwe(&param);
485
486        let start = Instant::now();
487        let btk = BootstrappingKey::from_sk(&mut rng, &param, &sk)?;
488        println!("btk took: {:?}", start.elapsed());
489
490        let msg_dist = Uniform::new(0_u64, param.t);
491        let m = Rq::rand_u64(&mut rng, msg_dist, &param.lwe().pt())?; // q=t, n=1
492        let p = TLWE::encode(&param.lwe(), &m); // plaintext
493
494        let c = TLWE::encrypt_s(&mut rng, &param.lwe(), &sk_tlwe, &p)?;
495
496        let start = Instant::now();
497        // the ugly const generics are temporary
498        let bootstrapped: TLWE = bootstrapping(&param, btk, table, c);
499        println!("bootstrapping took: {:?}", start.elapsed());
500
501        let p_recovered: T64 = bootstrapped.decrypt(&sk_tlwe);
502        let m_recovered = TLWE::decode(&param.lwe(), &p_recovered);
503        assert_eq!(m_recovered, m);
504
505        Ok(())
506    }
507}