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