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