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