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