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