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