fhe-study/arith/src/zq.rs
arith/src/zq.rsRaw
1use rand::{distributions::Distribution, Rng};
2use std::borrow::Borrow;
3use std::fmt;
4use std::ops::{Add, AddAssign, Div, Mul, Neg, Sub, SubAssign};
5
6/// Z_q, integers modulus q, not necessarily prime
7#[derive(Clone, Copy, PartialEq)]
8pub struct Zq {
9 pub q: u64,
10 pub v: u64,
11}
12
13// WIP
14// impl<const Q: u64> From<Vec<u64>> for Vec<Zq<Q>> {
15// fn from(v: Vec<u64>) -> Self {
16// v.into_iter().map(Zq::new).collect()
17// }
18// }
19
20pub(crate) fn modulus_u64(q: u64, e: u64) -> u64 {
21 (e % q + q) % q
22}
23impl Zq {
24 pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, q: u64) -> Self {
25 // TODO WIP
26 let r: f64 = dist.sample(&mut rng);
27 Self::from_f64(q, r)
28 // Self::from_u64(r.round() as u64)
29 }
30 pub fn from_u64(q: u64, v: u64) -> Self {
31 if v >= q {
32 // (v % Q + Q) % Q
33 return Zq {
34 q,
35 v: modulus_u64(q, v),
36 };
37 // return Zq(v % Q);
38 }
39 Zq { q, v }
40 }
41 pub fn from_f64(q: u64, e: f64) -> Self {
42 // WIP method
43 let e: i64 = e.round() as i64;
44 let q_i64 = q as i64;
45 if e < 0 || e >= q_i64 {
46 return Zq::from_u64(q, ((e % q_i64 + q_i64) % q_i64) as u64);
47 }
48 Zq { q, v: e as u64 }
49
50 // if e < 0 {
51 // // dbg!(&e);
52 // // dbg!(Zq::<Q>(((Q as i64 + e) % Q as i64) as u64));
53 // // return Zq(((Q as i64 + e) % Q as i64) as u64);
54 // return Zq(e as u64 % Q);
55 // } else if e >= Q as i64 {
56 // return Zq((e % Q as i64) as u64);
57 // }
58 // Zq(e as u64)
59 }
60 pub fn from_bool(q: u64, b: bool) -> Self {
61 if b {
62 Zq { q, v: 1 }
63 } else {
64 Zq { q, v: 0 }
65 }
66 }
67 pub fn zero(q: u64) -> Self {
68 Self { q, v: 0u64 }
69 }
70 pub fn one(q: u64) -> Self {
71 Self { q, v: 1u64 }
72 }
73 pub fn square(self) -> Self {
74 self * self
75 }
76 // modular exponentiation
77 pub fn exp(self, e: Self) -> Self {
78 // mul-square approach
79 let mut res = Self::one(self.q);
80 let mut rem = e.clone();
81 let mut exp = self;
82 // for rem != Self(0) {
83 while rem != Self::zero(self.q) {
84 // if odd
85 // TODO use a more readible expression
86 if 1 - ((rem.v & 1) << 1) as i64 == -1 {
87 res = res * exp;
88 }
89 exp = exp.square();
90 rem = Self {
91 q: self.q,
92 v: rem.v >> 1,
93 };
94 }
95 res
96 }
97 // multiplicative inverse
98 // WARNING: if this is needed, it means that 'Zq' is a Finite Field. For the moment we assume
99 // we work in a Finite Field
100 pub fn inv_OLD(self) -> Self {
101 // TODO
102 // let a = self.0;
103 // let q = Q;
104 let mut t = 0;
105 let mut r = self.q;
106 let mut new_t = 0;
107 let mut new_r = self.v.clone();
108 while new_r != 0 {
109 let q = r / new_r;
110
111 t = new_t.clone();
112 new_t = t - q;
113
114 r = new_r.clone();
115 new_r = r - (q * new_r);
116 }
117 // if t < 0 {
118 // t = t + q;
119 // }
120 return Zq::from_u64(self.q, t);
121 }
122 pub fn inv(self) -> Zq {
123 let (g, x, _) = Self::egcd(self.v as i128, self.q as i128);
124 if g != 1 {
125 // None
126 panic!("E");
127 } else {
128 let q = self.q as i128;
129 Zq::from_u64(self.q, ((x % q + q) % q) as u64) // TODO maybe just Zq::new(x)
130 }
131 }
132 fn egcd(a: i128, b: i128) -> (i128, i128, i128) {
133 if a == 0 {
134 (b, 0, 1)
135 } else {
136 let (g, x, y) = Self::egcd(b % a, a);
137 (g, y - (b / a) * x, x)
138 }
139 }
140
141 /// perform the mod switch operation from Q to Q', where Q2=Q'
142 pub fn mod_switch(&self, q2: u64) -> Zq {
143 Zq::from_u64(
144 q2,
145 ((self.v as f64 * q2 as f64) / self.q as f64).round() as u64,
146 )
147 }
148
149 pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
150 if beta == 2 {
151 self.decompose_base2(l)
152 } else {
153 self.decompose_base_beta(beta, l)
154 }
155 }
156 pub fn decompose_base_beta(&self, beta: u32, l: u32) -> Vec<Self> {
157 let mut rem: u64 = self.v;
158 // next if is for cases in which beta does not divide Q (concretely
159 // beta^l!=Q). round to the nearest multiple of q/beta^l
160 if rem >= beta.pow(l) as u64 {
161 // rem = Q - 1 - (Q / beta as u64); // floor
162 return vec![
163 Zq {
164 q: self.q,
165 v: beta as u64 - 1
166 };
167 l as usize
168 ];
169 }
170
171 let mut x: Vec<Self> = vec![];
172 for i in 1..l + 1 {
173 let den = self.q / beta.pow(i) as u64;
174 let x_i = rem / den; // division between u64 already does floor
175 x.push(Self::from_u64(self.q, x_i));
176 if x_i != 0 {
177 rem = rem % den;
178 }
179 }
180 x
181 }
182 /// decompose when beta=2
183 pub fn decompose_base2(&self, l: u32) -> Vec<Self> {
184 // next if is for cases in which beta does not divide Q (concretely
185 // beta^l!=Q). round to the nearest multiple of q/beta^l
186 if self.v >= 1 << l as u64 {
187 // rem = Q - 1 - (Q / beta as u64); // floor
188 // (where beta=2)
189 return vec![Zq::one(self.q); l as usize];
190 }
191
192 (0..l)
193 .rev()
194 .map(|i| Self::from_u64(self.q, ((self.v >> i) & 1) as u64))
195 .collect()
196
197 // naive ver:
198 // let mut rem: u64 = self.0;
199 // // next if is for cases in which beta does not divide Q (concretely
200 // // beta^l!=Q). round to the nearest multiple of q/beta^l
201 // if rem >= 1 << l as u64 {
202 // // rem = Q - 1 - (Q / beta as u64); // floor
203 // return vec![Zq(1); l as usize];
204 // }
205 //
206 // let mut x: Vec<Self> = vec![];
207 // for i in 1..l + 1 {
208 // let den = Q / (1 << i as u64);
209 // let x_i = rem / den; // division between u64 already does floor
210 // x.push(Self::from_u64(x_i));
211 // if x_i != 0 {
212 // rem = rem % den;
213 // }
214 // }
215 // x
216 }
217}
218
219impl Zq {
220 fn r#mod(self) -> Self {
221 if self.v >= self.q {
222 return Zq::from_u64(self.q, self.v % self.q);
223 }
224 self
225 }
226}
227
228impl Add<Zq> for Zq {
229 type Output = Self;
230
231 fn add(self, rhs: Self) -> Self::Output {
232 assert_eq!(self.q, rhs.q);
233
234 let mut v = self.v + rhs.v;
235 if v >= self.q {
236 v -= self.q;
237 }
238 Zq { q: self.q, v }
239 }
240}
241impl Add<&Zq> for &Zq {
242 type Output = Zq;
243
244 fn add(self, rhs: &Zq) -> Self::Output {
245 assert_eq!(self.q, rhs.q);
246
247 let mut v = self.v + rhs.v;
248 if v >= self.q {
249 v -= self.q;
250 }
251 Zq { q: self.q, v }
252 }
253}
254impl AddAssign<Zq> for Zq {
255 fn add_assign(&mut self, rhs: Self) {
256 *self = *self + rhs
257 }
258}
259impl std::iter::Sum for Zq {
260 fn sum<I>(mut iter: I) -> Self
261 where
262 I: Iterator<Item = Self>,
263 {
264 let first: Zq = iter.next().unwrap();
265 iter.fold(first, |acc, x| acc + x)
266 }
267}
268impl Sub<Zq> for Zq {
269 type Output = Self;
270
271 fn sub(self, rhs: Self) -> Zq {
272 assert_eq!(self.q, rhs.q);
273
274 if self.v >= rhs.v {
275 Zq {
276 q: self.q,
277 v: self.v - rhs.v,
278 }
279 } else {
280 Zq {
281 q: self.q,
282 v: (self.q + self.v) - rhs.v,
283 }
284 }
285 }
286}
287impl Sub<&Zq> for &Zq {
288 type Output = Zq;
289
290 fn sub(self, rhs: &Zq) -> Self::Output {
291 assert_eq!(self.q, rhs.q);
292
293 if self.q >= rhs.q {
294 Zq {
295 q: self.q,
296 v: self.v - rhs.v,
297 }
298 } else {
299 Zq {
300 q: self.q,
301 v: (self.q + self.v) - rhs.v,
302 }
303 }
304 }
305}
306impl SubAssign<Zq> for Zq {
307 fn sub_assign(&mut self, rhs: Self) {
308 *self = *self - rhs
309 }
310}
311impl Neg for Zq {
312 type Output = Self;
313
314 fn neg(self) -> Self::Output {
315 if self.v == 0 {
316 return self;
317 }
318 Zq {
319 q: self.q,
320 v: self.q - self.v,
321 }
322 }
323}
324impl Mul<Zq> for Zq {
325 type Output = Self;
326
327 fn mul(self, rhs: Self) -> Zq {
328 assert_eq!(self.q, rhs.q);
329
330 // TODO non-naive way
331 Zq {
332 q: self.q,
333 v: ((self.v as u128 * rhs.v as u128) % self.q as u128) as u64,
334 }
335 // Zq((self.0 * rhs.0) % Q)
336 }
337}
338impl Div<Zq> for Zq {
339 type Output = Self;
340
341 fn div(self, rhs: Self) -> Zq {
342 // TODO non-naive way
343 // Zq((self.0 / rhs.0) % Q)
344 self * rhs.inv()
345 }
346}
347
348impl fmt::Display for Zq {
349 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
350 write!(f, "{}", self.v)
351 }
352}
353impl fmt::Debug for Zq {
354 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
355 write!(f, "{}", self.v)
356 }
357}
358
359#[cfg(test)]
360mod tests {
361 use super::*;
362 use rand::distributions::Uniform;
363
364 #[test]
365 fn exp() {
366 const q: u64 = 1021;
367 let a = Zq::from_u64(q, 3);
368 let b = Zq::from_u64(q, 3);
369 assert_eq!(a.exp(b), Zq::from_u64(q, 27));
370
371 let a = Zq::from_u64(q, 1000);
372 let b = Zq::from_u64(q, 3);
373 assert_eq!(a.exp(b), Zq::from_u64(q, 949));
374 }
375 #[test]
376 fn neg() {
377 let q: u64 = 1021;
378 let a = Zq::from_f64(q, 101.0);
379 let b = Zq::from_f64(q, -1.0);
380 assert_eq!(-a, a * b);
381 }
382
383 fn recompose(q: u64, beta: u32, l: u32, d: Vec<Zq>) -> Zq {
384 let mut x = 0u64;
385 for i in 0..l {
386 x += d[i as usize].v * q / beta.pow(i + 1) as u64;
387 }
388 Zq::from_u64(q, x)
389 }
390
391 #[test]
392 fn test_decompose() {
393 let q1: u64 = 16;
394 let beta: u32 = 2;
395 let l: u32 = 4;
396 let x = Zq::from_u64(q1, 9);
397 let d = x.decompose(beta, l);
398
399 assert_eq!(recompose(q1, beta, l, d), x);
400
401 let q: u64 = 5u64.pow(3);
402 let beta: u32 = 5;
403 let l: u32 = 3;
404
405 let dist = Uniform::new(0_u64, q);
406 let mut rng = rand::thread_rng();
407
408 for _ in 0..1000 {
409 let x = Zq::from_u64(q, dist.sample(&mut rng));
410 let d = x.decompose(beta, l);
411 assert_eq!(d.len(), l as usize);
412 assert_eq!(recompose(q, beta, l, d), x);
413 }
414 }
415
416 #[test]
417 fn test_decompose_approx() {
418 let q: u64 = 2u64.pow(4) + 1;
419 let beta: u32 = 2;
420 let l: u32 = 4;
421 let x = Zq::from_u64(q, 16); // in q, but bigger than beta^l
422 let d = x.decompose(beta, l);
423 assert_eq!(d.len(), l as usize);
424 assert_eq!(recompose(q, beta, l, d), Zq::from_u64(q, 15));
425
426 let q2: u64 = 5u64.pow(3) + 1;
427 let beta: u32 = 5;
428 let l: u32 = 3;
429 let x = Zq::from_u64(q2, 125); // in q, but bigger than beta^l
430 let d = x.decompose(beta, l);
431 assert_eq!(d.len(), l as usize);
432 assert_eq!(recompose(q2, beta, l, d), Zq::from_u64(q2, 124));
433
434 let q3: u64 = 2u64.pow(16) + 1;
435 let beta: u32 = 2;
436 let l: u32 = 16;
437 let x = Zq::from_u64(q3, q3 - 1); // in q, but bigger than beta^l
438 let d = x.decompose(beta, l);
439 assert_eq!(d.len(), l as usize);
440 assert_eq!(
441 recompose(q3, beta, l, d),
442 Zq::from_u64(q3, beta.pow(l) as u64 - 1)
443 );
444 }
445}