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