phantom-zone/src/lwe.rs
src/lwe.rsRaw
1use std::{
2 cell::RefCell,
3 collections::btree_map::Values,
4 fmt::{Debug, Display},
5 marker::PhantomData,
6};
7
8use itertools::{izip, Itertools};
9use num_traits::{abs, PrimInt, ToPrimitive, Zero};
10
11use crate::{
12 backend::{ArithmeticOps, VectorOps},
13 decomposer::Decomposer,
14 lwe,
15 num::UnsignedInteger,
16 random::{DefaultSecureRng, NewWithSeed, RandomGaussianDist, RandomUniformDist, DEFAULT_RNG},
17 utils::{fill_random_ternary_secret_with_hamming_weight, TryConvertFrom, WithLocal},
18 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
19};
20
21struct SeededLweKeySwitchingKey<Ro, S>
22where
23 Ro: Row,
24{
25 data: Ro,
26 seed: S,
27 to_lwe_n: usize,
28 modulus: Ro::Element,
29}
30
31impl<Ro: RowEntity, S> SeededLweKeySwitchingKey<Ro, S> {
32 pub(crate) fn empty(
33 from_lwe_n: usize,
34 to_lwe_n: usize,
35 d: usize,
36 seed: S,
37 modulus: Ro::Element,
38 ) -> Self {
39 let data = Ro::zeros(from_lwe_n * d);
40 SeededLweKeySwitchingKey {
41 data,
42 to_lwe_n,
43 seed,
44 modulus,
45 }
46 }
47}
48
49struct LweKeySwitchingKey<M, R> {
50 data: M,
51 _phantom: PhantomData<R>,
52}
53
54impl<
55 M: MatrixMut + MatrixEntity,
56 R: NewWithSeed + RandomUniformDist<[M::MatElement], Parameters = M::MatElement>,
57 > From<&SeededLweKeySwitchingKey<M::R, R::Seed>> for LweKeySwitchingKey<M, R>
58where
59 M::R: RowMut,
60 R::Seed: Clone,
61 M::MatElement: Copy,
62{
63 fn from(value: &SeededLweKeySwitchingKey<M::R, R::Seed>) -> Self {
64 let mut p_rng = R::new_with_seed(value.seed.clone());
65 let mut data = M::zeros(value.data.as_ref().len(), value.to_lwe_n + 1);
66 izip!(value.data.as_ref().iter(), data.iter_rows_mut()).for_each(|(bi, lwe_i)| {
67 RandomUniformDist::random_fill(&mut p_rng, &value.modulus, &mut lwe_i.as_mut()[1..]);
68 lwe_i.as_mut()[0] = *bi;
69 });
70 LweKeySwitchingKey {
71 data,
72 _phantom: PhantomData,
73 }
74 }
75}
76
77trait LweCiphertext<M: Matrix> {}
78
79#[derive(Clone)]
80pub struct LweSecret {
81 pub(crate) values: Vec<i32>,
82}
83
84impl Secret for LweSecret {
85 type Element = i32;
86 fn values(&self) -> &[Self::Element] {
87 &self.values
88 }
89}
90
91impl LweSecret {
92 pub(crate) fn random(hw: usize, n: usize) -> LweSecret {
93 DefaultSecureRng::with_local_mut(|rng| {
94 let mut out = vec![0i32; n];
95 fill_random_ternary_secret_with_hamming_weight(&mut out, hw, rng);
96
97 LweSecret { values: out }
98 })
99 }
100}
101
102pub(crate) fn lwe_key_switch<
103 M: Matrix,
104 Ro: AsMut<[M::MatElement]> + AsRef<[M::MatElement]>,
105 Op: VectorOps<Element = M::MatElement> + ArithmeticOps<Element = M::MatElement>,
106 D: Decomposer<Element = M::MatElement>,
107>(
108 lwe_out: &mut Ro,
109 lwe_in: &Ro,
110 lwe_ksk: &M,
111 operator: &Op,
112 decomposer: &D,
113) {
114 assert!(lwe_ksk.dimension().0 == ((lwe_in.as_ref().len() - 1) * decomposer.decomposition_count()));
115 assert!(lwe_out.as_ref().len() == lwe_ksk.dimension().1);
116
117 let lwe_in_a_decomposed = lwe_in
118 .as_ref()
119 .iter()
120 .skip(1)
121 .flat_map(|ai| decomposer.decompose(ai));
122 izip!(lwe_in_a_decomposed, lwe_ksk.iter_rows()).for_each(|(ai_j, beta_ij_lwe)| {
123 operator.elwise_fma_scalar_mut(lwe_out.as_mut(), beta_ij_lwe.as_ref(), &ai_j);
124 });
125
126 let out_b = operator.add(&lwe_out.as_ref()[0], &lwe_in.as_ref()[0]);
127 lwe_out.as_mut()[0] = out_b;
128}
129
130pub fn lwe_ksk_keygen<
131 Ro: Row + RowMut + RowEntity,
132 S,
133 Op: VectorOps<Element = Ro::Element> + ArithmeticOps<Element = Ro::Element>,
134 R: RandomGaussianDist<Ro::Element, Parameters = Ro::Element>,
135 PR: RandomUniformDist<[Ro::Element], Parameters = Ro::Element>,
136>(
137 from_lwe_sk: &[S],
138 to_lwe_sk: &[S],
139 ksk_out: &mut Ro,
140 gadget: &[Ro::Element],
141 operator: &Op,
142 p_rng: &mut PR,
143 rng: &mut R,
144) where
145 Ro: TryConvertFrom<[S], Parameters = Ro::Element>,
146 Ro::Element: Zero + Debug,
147{
148 assert!(ksk_out.as_ref().len() == (from_lwe_sk.len() * gadget.len()));
149
150 let d = gadget.len();
151
152 let modulus = VectorOps::modulus(operator);
153 let mut neg_sk_in_m = Ro::try_convert_from(from_lwe_sk, &modulus);
154 operator.elwise_neg_mut(neg_sk_in_m.as_mut());
155 let sk_out_m = Ro::try_convert_from(to_lwe_sk, &modulus);
156
157 let mut scratch = Ro::zeros(to_lwe_sk.len());
158
159 izip!(neg_sk_in_m.as_ref(), ksk_out.as_mut().chunks_mut(d)).for_each(
160 |(neg_sk_in_si, d_lwes_partb)| {
161 izip!(gadget.iter(), d_lwes_partb.into_iter()).for_each(|(f, lwe_b)| {
162 // sample `a`
163 RandomUniformDist::random_fill(p_rng, &modulus, scratch.as_mut());
164
165 // a * z
166 let mut az = Ro::Element::zero();
167 izip!(scratch.as_ref().iter(), sk_out_m.as_ref()).for_each(|(ai, si)| {
168 let ai_si = operator.mul(ai, si);
169 az = operator.add(&az, &ai_si);
170 });
171
172 // a*z + (-s_i)*\beta^j + e
173 let mut b = operator.add(&az, &operator.mul(f, neg_sk_in_si));
174 let mut e = Ro::Element::zero();
175 RandomGaussianDist::random_fill(rng, &modulus, &mut e);
176 b = operator.add(&b, &e);
177
178 *lwe_b = b;
179 })
180 },
181 );
182}
183
184/// Encrypts encoded message m as LWE ciphertext
185pub fn encrypt_lwe<
186 Ro: Row + RowMut,
187 R: RandomGaussianDist<Ro::Element, Parameters = Ro::Element>
188 + RandomUniformDist<[Ro::Element], Parameters = Ro::Element>,
189 S,
190 Op: ArithmeticOps<Element = Ro::Element>,
191>(
192 lwe_out: &mut Ro,
193 m: &Ro::Element,
194 s: &[S],
195 operator: &Op,
196 rng: &mut R,
197) where
198 Ro: TryConvertFrom<[S], Parameters = Ro::Element>,
199 Ro::Element: Zero,
200{
201 let s = Ro::try_convert_from(s, &operator.modulus());
202 assert!(s.as_ref().len() == (lwe_out.as_ref().len() - 1));
203
204 // a*s
205 RandomUniformDist::random_fill(rng, &operator.modulus(), &mut lwe_out.as_mut()[1..]);
206 let mut sa = Ro::Element::zero();
207 izip!(lwe_out.as_mut().iter().skip(1), s.as_ref()).for_each(|(ai, si)| {
208 let tmp = operator.mul(ai, si);
209 sa = operator.add(&tmp, &sa);
210 });
211
212 // b = a*s + e + m
213 let mut e = Ro::Element::zero();
214 RandomGaussianDist::random_fill(rng, &operator.modulus(), &mut e);
215 let b = operator.add(&operator.add(&sa, &e), m);
216 lwe_out.as_mut()[0] = b;
217}
218
219pub fn decrypt_lwe<Ro: Row, Op: ArithmeticOps<Element = Ro::Element>, S>(
220 lwe_ct: &Ro,
221 s: &[S],
222 operator: &Op,
223) -> Ro::Element
224where
225 Ro: TryConvertFrom<[S], Parameters = Ro::Element>,
226 Ro::Element: Zero,
227{
228 let s = Ro::try_convert_from(s, &operator.modulus());
229
230 let mut sa = Ro::Element::zero();
231 izip!(lwe_ct.as_ref().iter().skip(1), s.as_ref()).for_each(|(ai, si)| {
232 let tmp = operator.mul(ai, si);
233 sa = operator.add(&tmp, &sa);
234 });
235
236 let b = &lwe_ct.as_ref()[0];
237 operator.sub(b, &sa)
238}
239
240/// Measures noise in input LWE ciphertext with reference of `ideal_m`
241///
242/// - ct: Input LWE ciphertext
243/// - s: corresponding secret
244/// - ideal_m: Ideal `encoded` message
245pub(crate) fn measure_noise_lwe<Ro: Row, Op: ArithmeticOps<Element = Ro::Element>, S>(
246 ct: &Ro,
247 s: &[S],
248 operator: &Op,
249 ideal_m: &Ro::Element,
250) -> f64
251where
252 Ro: TryConvertFrom<[S], Parameters = Ro::Element>,
253 Ro::Element: Zero + ToPrimitive + PrimInt + Display,
254{
255 assert!(s.len() == ct.as_ref().len() - 1,);
256
257 let s = Ro::try_convert_from(s, &operator.modulus());
258 let mut sa = Ro::Element::zero();
259 izip!(s.as_ref().iter(), ct.as_ref().iter().skip(1)).for_each(|(si, ai)| {
260 sa = operator.add(&sa, &operator.mul(si, ai));
261 });
262 let m = operator.sub(&ct.as_ref()[0], &sa);
263
264 let mut diff = operator.sub(&m, ideal_m);
265 let q = operator.modulus();
266 if diff > (q >> 1) {
267 diff = q - diff;
268 }
269 return diff.to_f64().unwrap().log2();
270}
271
272#[cfg(test)]
273mod tests {
274
275 use crate::{
276 backend::{ModInit, ModularOpsU64},
277 decomposer::{Decomposer, DefaultDecomposer},
278 lwe::{lwe_key_switch, measure_noise_lwe},
279 random::DefaultSecureRng,
280 rgsw::measure_noise,
281 Secret,
282 };
283
284 use super::{
285 decrypt_lwe, encrypt_lwe, lwe_ksk_keygen, LweKeySwitchingKey, LweSecret,
286 SeededLweKeySwitchingKey,
287 };
288
289 const K: usize = 50;
290
291 #[test]
292 fn encrypt_decrypt_works() {
293 let logq = 16;
294 let q = 1u64 << logq;
295 let lwe_n = 1024;
296 let logp = 3;
297
298 let modq_op = ModularOpsU64::new(q);
299 let lwe_sk = LweSecret::random(lwe_n >> 1, lwe_n);
300
301 let mut rng = DefaultSecureRng::new();
302
303 // encrypt
304 for m in 0..1u64 << logp {
305 let encoded_m = m << (logq - logp);
306 let mut lwe_ct = vec![0u64; lwe_n + 1];
307 encrypt_lwe(
308 &mut lwe_ct,
309 &encoded_m,
310 &lwe_sk.values(),
311 &modq_op,
312 &mut rng,
313 );
314 let encoded_m_back = decrypt_lwe(&lwe_ct, &lwe_sk.values(), &modq_op);
315 let m_back = ((((encoded_m_back as f64) * ((1 << logp) as f64)) / q as f64).round()
316 as u64)
317 % (1u64 << logp);
318 assert_eq!(m, m_back, "Expected {m} but got {m_back}");
319 }
320 }
321
322 #[test]
323 fn key_switch_works() {
324 let logq = 18;
325 let logp = 2;
326 let q = 1u64 << logq;
327 let lwe_in_n = 2048;
328 let lwe_out_n = 493;
329 let d_ks = 3;
330 let logb = 6;
331
332 let lwe_sk_in = LweSecret::random(lwe_in_n >> 1, lwe_in_n);
333 let lwe_sk_out = LweSecret::random(lwe_out_n >> 1, lwe_out_n);
334
335 let mut rng = DefaultSecureRng::new();
336 let modq_op = ModularOpsU64::new(q);
337
338 // genrate ksk
339 for _ in 0..K {
340 let mut ksk_seed = [0u8; 32];
341 rng.fill_bytes(&mut ksk_seed);
342 let mut seeded_ksk =
343 SeededLweKeySwitchingKey::empty(lwe_in_n, lwe_out_n, d_ks, ksk_seed, q);
344 let mut p_rng = DefaultSecureRng::new_seeded(ksk_seed);
345 let decomposer = DefaultDecomposer::new(q, logb, d_ks);
346 let gadget = decomposer.gadget_vector();
347 lwe_ksk_keygen(
348 &lwe_sk_in.values(),
349 &lwe_sk_out.values(),
350 &mut seeded_ksk.data,
351 &gadget,
352 &modq_op,
353 &mut p_rng,
354 &mut rng,
355 );
356 // println!("{:?}", ksk);
357 let ksk = LweKeySwitchingKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_ksk);
358
359 for m in 0..(1 << logp) {
360 // encrypt using lwe_sk_in
361 let encoded_m = m << (logq - logp);
362 let mut lwe_in_ct = vec![0u64; lwe_in_n + 1];
363 encrypt_lwe(
364 &mut lwe_in_ct,
365 &encoded_m,
366 lwe_sk_in.values(),
367 &modq_op,
368 &mut rng,
369 );
370
371 // key switch from lwe_sk_in to lwe_sk_out
372 let decomposer = DefaultDecomposer::new(1u64 << logq, logb, d_ks);
373 let mut lwe_out_ct = vec![0u64; lwe_out_n + 1];
374 lwe_key_switch(
375 &mut lwe_out_ct,
376 &lwe_in_ct,
377 &ksk.data,
378 &modq_op,
379 &decomposer,
380 );
381
382 // decrypt lwe_out_ct using lwe_sk_out
383 let encoded_m_back = decrypt_lwe(&lwe_out_ct, &lwe_sk_out.values(), &modq_op);
384 let m_back = ((((encoded_m_back as f64) * ((1 << logp) as f64)) / q as f64).round()
385 as u64)
386 % (1u64 << logp);
387 let noise =
388 measure_noise_lwe(&lwe_out_ct, lwe_sk_out.values(), &modq_op, &encoded_m);
389 println!("Noise: {noise}");
390 assert_eq!(m, m_back, "Expected {m} but got {m_back}");
391 // dbg!(m, m_back);
392 // dbg!(encoded_m, encoded_m_back);
393 }
394 }
395 }
396}