fhe-study/tfhe/src/tlwe.rs
tfhe/src/tlwe.rsRaw
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;
14
15// #[derive(Clone, Debug)]
16pub struct SecretKey<const K: usize>(pub glwe::SecretKey<T64, K>);
17// pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>;
18
19impl<const KN: usize> SecretKey<KN> {
20 /// from TFHE [2018-421] paper: A TLWE key k \in B^n, can be interpreted as a
21 /// TRLWE key K \in B_N[X]^k having the same sequence of coefficients and
22 /// vice-versa.
23 pub fn to_tglwe<const N: usize, const K: usize>(self) -> crate::tglwe::SecretKey<N, K> {
24 let s: TR<T64, KN> = self.0 .0;
25 // split into K vectors, and interpret each of them as a T_N[X]/(X^N+1)
26 // polynomial
27 let r: Vec<Tn<N>> =
28 s.0.chunks(N)
29 .map(|v| Tn::<N>::from_vec(v.to_vec()))
30 .collect();
31 crate::tglwe::SecretKey(glwe::SecretKey::<Tn<N>, K>(TR(r)))
32 }
33}
34
35// #[derive(Clone, Debug)]
36// pub struct PublicKey<const K: usize>(glwe::PublicKey<T64, K>);
37pub type PublicKey<const K: usize> = glwe::PublicKey<T64, K>;
38
39#[derive(Clone, Debug)]
40pub struct KSK<const K: usize>(Vec<TLev<K>>);
41
42#[derive(Clone, Debug)]
43pub struct TLWE<const K: usize>(pub GLWE<T64, K>);
44
45impl<const K: usize> TLWE<K> {
46 pub fn zero() -> Self {
47 Self(GLWE::<T64, K>::zero())
48 }
49
50 pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
51 let (sk, pk): (glwe::SecretKey<T64, K>, glwe::PublicKey<T64, K>) = GLWE::new_key(rng)?;
52 // Ok((SecretKey(sk), PublicKey(pk)))
53 Ok((SecretKey(sk), pk))
54 }
55
56 pub fn encode<const P: u64>(m: &Rq<P, 1>) -> T64 {
57 let delta = u64::MAX / P; // floored
58 let coeffs = m.coeffs();
59 // Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
60 T64(coeffs[0].0 * delta)
61 }
62 pub fn decode<const P: u64>(p: &T64) -> Rq<P, 1> {
63 let p = p.mul_div_round(P, u64::MAX);
64 Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
65 }
66
67 // encrypts with the given SecretKey (instead of PublicKey)
68 pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &T64) -> Result<Self> {
69 let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
70 Ok(Self(glwe))
71 }
72 pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &T64) -> Result<Self> {
73 let glwe = GLWE::encrypt(rng, &pk, p)?;
74 Ok(Self(glwe))
75 }
76 pub fn decrypt(&self, sk: &SecretKey<K>) -> T64 {
77 self.0.decrypt(&sk.0)
78 }
79
80 pub fn new_ksk(
81 mut rng: impl Rng,
82 beta: u32,
83 l: u32,
84 sk: &SecretKey<K>,
85 new_sk: &SecretKey<K>,
86 ) -> Result<KSK<K>> {
87 let r: Vec<TLev<K>> = (0..K)
88 .into_iter()
89 .map(|i|
90 // treat sk_i as the msg being encrypted
91 TLev::<K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0.0 .0[i]))
92 .collect::<Result<Vec<_>>>()?;
93
94 Ok(KSK(r))
95 }
96 pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<K>) -> Self {
97 let (a, b): (TR<T64, K>, T64) = (self.0 .0.clone(), self.0 .1);
98
99 let lhs: TLWE<K> = TLWE(GLWE(TR::zero(), b));
100
101 // K iterations, ksk.0 contains K times GLev
102 let rhs: TLWE<K> = zip_eq(a.0, 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}
109
110impl<const K: usize> Add<TLWE<K>> for TLWE<K> {
111 type Output = Self;
112 fn add(self, other: Self) -> Self {
113 Self(self.0 + other.0)
114 }
115}
116impl<const K: usize> AddAssign for TLWE<K> {
117 fn add_assign(&mut self, rhs: Self) {
118 self.0 += rhs.0
119 }
120}
121impl<const K: usize> Sum<TLWE<K>> for TLWE<K> {
122 fn sum<I>(iter: I) -> Self
123 where
124 I: Iterator<Item = Self>,
125 {
126 let mut acc = TLWE::<K>::zero();
127 for e in iter {
128 acc += e;
129 }
130 acc
131 }
132}
133
134impl<const K: usize> Sub<TLWE<K>> for TLWE<K> {
135 type Output = Self;
136 fn sub(self, other: Self) -> Self {
137 Self(self.0 - other.0)
138 }
139}
140
141// plaintext addition
142impl<const K: usize> Add<T64> for TLWE<K> {
143 type Output = Self;
144 fn add(self, plaintext: T64) -> Self {
145 let a: TR<T64, K> = self.0 .0;
146 let b: T64 = self.0 .1 + plaintext;
147 Self(GLWE(a, b))
148 }
149}
150// plaintext substraction
151impl<const K: usize> Sub<T64> for TLWE<K> {
152 type Output = Self;
153 fn sub(self, plaintext: T64) -> Self {
154 let a: TR<T64, K> = self.0 .0;
155 let b: T64 = self.0 .1 - plaintext;
156 Self(GLWE(a, b))
157 }
158}
159// plaintext multiplication
160impl<const K: usize> Mul<T64> for TLWE<K> {
161 type Output = Self;
162 fn mul(self, plaintext: T64) -> Self {
163 let a: TR<T64, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect());
164 let b: T64 = self.0 .1 * plaintext;
165 Self(GLWE(a, b))
166 }
167}
168
169#[cfg(test)]
170mod tests {
171 use anyhow::Result;
172 use rand::distributions::Uniform;
173
174 use super::*;
175
176 #[test]
177 fn test_encrypt_decrypt() -> Result<()> {
178 const T: u64 = 128; // msg space (msg modulus)
179 const K: usize = 16;
180 type S = TLWE<K>;
181
182 let mut rng = rand::thread_rng();
183 let msg_dist = Uniform::new(0_u64, T);
184
185 for _ in 0..200 {
186 let (sk, pk) = S::new_key(&mut rng)?;
187
188 let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
189 dbg!(&m);
190 let p: T64 = S::encode::<T>(&m);
191 dbg!(&p);
192
193 let c = S::encrypt(&mut rng, &pk, &p)?;
194 let p_recovered = c.decrypt(&sk);
195 let m_recovered = S::decode::<T>(&p_recovered);
196
197 assert_eq!(m, m_recovered);
198
199 // same but using encrypt_s (with sk instead of pk))
200 let c = S::encrypt_s(&mut rng, &sk, &p)?;
201 let p_recovered = c.decrypt(&sk);
202 let m_recovered = S::decode::<T>(&p_recovered);
203
204 assert_eq!(m, m_recovered);
205 }
206
207 Ok(())
208 }
209
210 #[test]
211 fn test_addition() -> Result<()> {
212 const T: u64 = 128;
213 const K: usize = 16;
214 type S = TLWE<K>;
215
216 let mut rng = rand::thread_rng();
217 let msg_dist = Uniform::new(0_u64, T);
218
219 for _ in 0..200 {
220 let (sk, pk) = S::new_key(&mut rng)?;
221
222 let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
223 let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
224 let p1: T64 = S::encode::<T>(&m1); // plaintext
225 let p2: T64 = S::encode::<T>(&m2); // plaintext
226
227 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
228 let c2 = S::encrypt(&mut rng, &pk, &p2)?;
229
230 let c3 = c1 + c2;
231
232 let p3_recovered = c3.decrypt(&sk);
233 let m3_recovered = S::decode::<T>(&p3_recovered);
234
235 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
236 }
237
238 Ok(())
239 }
240
241 #[test]
242 fn test_add_plaintext() -> Result<()> {
243 const T: u64 = 128;
244 const K: usize = 16;
245 type S = TLWE<K>;
246
247 let mut rng = rand::thread_rng();
248 let msg_dist = Uniform::new(0_u64, T);
249
250 for _ in 0..200 {
251 let (sk, pk) = S::new_key(&mut rng)?;
252
253 let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
254 let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
255 let p1: T64 = S::encode::<T>(&m1); // plaintext
256 let p2: T64 = S::encode::<T>(&m2); // plaintext
257
258 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
259
260 let c3 = c1 + p2;
261
262 let p3_recovered = c3.decrypt(&sk);
263 let m3_recovered = S::decode::<T>(&p3_recovered);
264
265 assert_eq!(m1 + m2, m3_recovered);
266 }
267
268 Ok(())
269 }
270
271 #[test]
272 fn test_mul_plaintext() -> Result<()> {
273 const T: u64 = 128;
274 const K: usize = 16;
275 type S = TLWE<K>;
276
277 let mut rng = rand::thread_rng();
278 let msg_dist = Uniform::new(0_u64, T);
279
280 for _ in 0..200 {
281 let (sk, pk) = S::new_key(&mut rng)?;
282
283 let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
284 let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
285 let p1: T64 = S::encode::<T>(&m1);
286 // don't scale up p2, set it directly from m2
287 // let p2: T64 = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
288 let p2: T64 = T64(m2.coeffs()[0].0);
289
290 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
291
292 let c3 = c1 * p2;
293
294 let p3_recovered: T64 = c3.decrypt(&sk);
295 let m3_recovered = S::decode::<T>(&p3_recovered);
296 assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
297 }
298
299 Ok(())
300 }
301
302 #[test]
303 fn test_key_switch() -> Result<()> {
304 const T: u64 = 128; // plaintext modulus
305 const K: usize = 16;
306 type S = TLWE<K>;
307
308 let beta: u32 = 2;
309 let l: u32 = 64;
310
311 let mut rng = rand::thread_rng();
312
313 let (sk, pk) = S::new_key(&mut rng)?;
314 let (sk2, _) = S::new_key(&mut rng)?;
315 // ksk to switch from sk to sk2
316 let ksk = S::new_ksk(&mut rng, beta, l, &sk, &sk2)?;
317
318 let msg_dist = Uniform::new(0_u64, T);
319 let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
320 let p = S::encode::<T>(&m); // plaintext
321 //
322 let c = S::encrypt_s(&mut rng, &sk, &p)?;
323
324 let c2 = c.key_switch(beta, l, &ksk);
325
326 // decrypt with the 2nd secret key
327 let p_recovered = c2.decrypt(&sk2);
328 let m_recovered = S::decode::<T>(&p_recovered);
329 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
330
331 // do the same but now encrypting with pk
332 let c = S::encrypt(&mut rng, &pk, &p)?;
333 let c2 = c.key_switch(beta, l, &ksk);
334 let p_recovered = c2.decrypt(&sk2);
335 let m_recovered = S::decode::<T>(&p_recovered);
336 assert_eq!(m, m_recovered);
337
338 Ok(())
339 }
340}