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