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