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