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