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};
11
12const ERR_SIGMA: f64 = 3.2;
13
14#[derive(Clone, Debug)]
15pub struct TLWE<const K: usize>(TR<Tn<1>, K>, Tn<1>);
16
17#[derive(Clone, Debug)]
18pub struct SecretKey<const K: usize>(TR<Tn<1>, K>);
19#[derive(Clone, Debug)]
20pub struct PublicKey<const K: usize>(Tn<1>, TR<Tn<1>, K>);
21
22impl<const K: usize> TLWE<K> {
23 pub fn zero() -> Self {
24 Self(TR::zero(), Tn::zero())
25 }
26
27 pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
28 let Xi_key = Uniform::new(0_f64, 2_f64);
29 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
30
31 let s: TR<Tn<1>, K> = TR::rand(&mut rng, Xi_key);
32 let a: TR<Tn<1>, K> = TR::rand(&mut rng, Standard);
33 let e = Tn::rand(&mut rng, Xi_err);
34
35 let pk: PublicKey<K> = PublicKey((&a * &s) + e, a);
36 Ok((SecretKey(s), pk))
37 }
38
39 pub fn encode<const P: u64>(m: &Rq<P, 1>) -> Tn<1> {
40 let delta = u64::MAX / P; // floored
41 let coeffs = m.coeffs();
42 Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
43 }
44 pub fn decode<const P: u64>(p: &Tn<1>) -> Rq<P, 1> {
45 let p = p.mul_div_round(P, u64::MAX);
46 Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
47 }
48
49 // encrypts with the given SecretKey (instead of PublicKey)
50 pub fn encrypt_s(mut rng: impl Rng, sk: &SecretKey<K>, m: &Tn<1>) -> Result<Self> {
51 let Xi_key = Uniform::new(0_f64, 2_f64);
52 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
53
54 let a: TR<Tn<1>, K> = TR::rand(&mut rng, Xi_key);
55 let e = Tn::rand(&mut rng, Xi_err);
56
57 let b: Tn<1> = (&a * &sk.0) + *m + e;
58 Ok(Self(a, b))
59 }
60 pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<K>, m: &Tn<1>) -> Result<Self> {
61 let Xi_key = Uniform::new(0_f64, 2_f64);
62 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
63
64 let u: Tn<1> = Tn::rand(&mut rng, Xi_key);
65
66 let e0: Tn<1> = Tn::rand(&mut rng, Xi_err);
67 let e1 = TR::<Tn<1>, K>::rand(&mut rng, Xi_err);
68
69 let b: Tn<1> = pk.0 * u + *m + e0;
70 let d: TR<Tn<1>, K> = &pk.1 * &u + e1;
71
72 Ok(Self(d, b))
73 }
74 pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> {
75 let (d, b): (TR<Tn<1>, K>, Tn<1>) = (self.0.clone(), self.1);
76 b - &d * &sk.0
77 }
78}
79
80impl<const K: usize> Add<TLWE<K>> for TLWE<K> {
81 type Output = Self;
82 fn add(self, other: Self) -> Self {
83 let a: TR<Tn<1>, K> = self.0 + other.0;
84 let b: Tn<1> = self.1 + other.1;
85 Self(a, b)
86 }
87}
88impl<const K: usize> AddAssign for TLWE<K> {
89 fn add_assign(&mut self, rhs: Self) {
90 for i in 0..K {
91 self.0 .0[i] = self.0 .0[i] + rhs.0 .0[i];
92 }
93 self.1 = self.1 + rhs.1;
94 }
95}
96impl<const K: usize> Sum<TLWE<K>> for TLWE<K> {
97 fn sum<I>(iter: I) -> Self
98 where
99 I: Iterator<Item = Self>,
100 {
101 let mut acc = TLWE::<K>::zero();
102 for e in iter {
103 acc += e;
104 }
105 acc
106 }
107}
108
109impl<const K: usize> Sub<TLWE<K>> for TLWE<K> {
110 type Output = Self;
111 fn sub(self, other: Self) -> Self {
112 let a: TR<Tn<1>, K> = self.0 - other.0;
113 let b: Tn<1> = self.1 - other.1;
114 Self(a, b)
115 }
116}
117
118// plaintext addition
119impl<const K: usize> Add<Tn<1>> for TLWE<K> {
120 type Output = Self;
121 fn add(self, plaintext: Tn<1>) -> Self {
122 let a: TR<Tn<1>, K> = self.0;
123 let b: Tn<1> = self.1 + plaintext;
124 Self(a, b)
125 }
126}
127// plaintext substraction
128impl<const K: usize> Sub<Tn<1>> for TLWE<K> {
129 type Output = Self;
130 fn sub(self, plaintext: Tn<1>) -> Self {
131 let a: TR<Tn<1>, K> = self.0;
132 let b: Tn<1> = self.1 - plaintext;
133 Self(a, b)
134 }
135}
136// plaintext multiplication
137impl<const K: usize> Mul<Tn<1>> for TLWE<K> {
138 type Output = Self;
139 fn mul(self, plaintext: Tn<1>) -> Self {
140 let a: TR<Tn<1>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
141 let b: Tn<1> = self.1 * plaintext;
142 Self(a, b)
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use anyhow::Result;
149 use rand::distributions::Uniform;
150
151 use super::*;
152
153 #[test]
154 fn test_encrypt_decrypt() -> Result<()> {
155 const T: u64 = 128; // msg space (msg modulus)
156 const K: usize = 16;
157 type S = TLWE<K>;
158
159 let mut rng = rand::thread_rng();
160
161 for _ in 0..200 {
162 let (sk, pk) = S::new_key(&mut rng)?;
163
164 let msg_dist = Uniform::new(0_u64, T);
165 let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
166 dbg!(&m);
167 let p: Tn<1> = S::encode::<T>(&m);
168 dbg!(&p);
169
170 let c = S::encrypt(&mut rng, &pk, &p)?;
171 let p_recovered = c.decrypt(&sk);
172 let m_recovered = S::decode::<T>(&p_recovered);
173
174 assert_eq!(m, m_recovered);
175
176 // same but using encrypt_s (with sk instead of pk))
177 let c = S::encrypt_s(&mut rng, &sk, &p)?;
178 let p_recovered = c.decrypt(&sk);
179 let m_recovered = S::decode::<T>(&p_recovered);
180
181 assert_eq!(m, m_recovered);
182 }
183
184 Ok(())
185 }
186
187 #[test]
188 fn test_addition() -> Result<()> {
189 const T: u64 = 128;
190 const K: usize = 16;
191 type S = TLWE<K>;
192
193 let mut rng = rand::thread_rng();
194
195 for _ in 0..200 {
196 let (sk, pk) = S::new_key(&mut rng)?;
197
198 let msg_dist = Uniform::new(0_u64, T);
199 let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
200 let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
201 let p1: Tn<1> = S::encode::<T>(&m1); // plaintext
202 let p2: Tn<1> = S::encode::<T>(&m2); // plaintext
203
204 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
205 let c2 = S::encrypt(&mut rng, &pk, &p2)?;
206
207 let c3 = c1 + c2;
208
209 let p3_recovered = c3.decrypt(&sk);
210 let m3_recovered = S::decode::<T>(&p3_recovered);
211
212 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
213 }
214
215 Ok(())
216 }
217
218 #[test]
219 fn test_add_plaintext() -> Result<()> {
220 const T: u64 = 128;
221 const K: usize = 16;
222 type S = TLWE<K>;
223
224 let mut rng = rand::thread_rng();
225
226 for _ in 0..200 {
227 let (sk, pk) = S::new_key(&mut rng)?;
228
229 let msg_dist = Uniform::new(0_u64, T);
230 let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
231 let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
232 let p1: Tn<1> = S::encode::<T>(&m1); // plaintext
233 let p2: Tn<1> = S::encode::<T>(&m2); // plaintext
234
235 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
236
237 let c3 = c1 + p2;
238
239 let p3_recovered = c3.decrypt(&sk);
240 let m3_recovered = S::decode::<T>(&p3_recovered);
241
242 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered);
243 }
244
245 Ok(())
246 }
247
248 #[test]
249 fn test_mul_plaintext() -> Result<()> {
250 const T: u64 = 128;
251 const K: usize = 16;
252 type S = TLWE<K>;
253
254 let mut rng = rand::thread_rng();
255
256 for _ in 0..200 {
257 let (sk, pk) = S::new_key(&mut rng)?;
258
259 let msg_dist = Uniform::new(0_u64, T);
260 let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
261 let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
262 let p1: Tn<1> = S::encode::<T>(&m1);
263 // don't scale up p2, set it directly from m2
264 let p2: Tn<1> = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
265
266 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
267
268 let c3 = c1 * p2;
269
270 let p3_recovered: Tn<1> = c3.decrypt(&sk);
271 let m3_recovered = S::decode::<T>(&p3_recovered);
272 assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
273 }
274
275 Ok(())
276 }
277}