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
135impl<const Q: u64> Zq<Q> {
136 fn r#mod(self) -> Self {
137 if self.0 >= Q {
138 return Zq(self.0 % Q);
139 }
140 self
141 }
142}
143
144impl<const Q: u64> ops::Add<Zq<Q>> for Zq<Q> {
145 type Output = Self;
146
147 fn add(self, rhs: Self) -> Self::Output {
148 let mut r = self.0 + rhs.0;
149 if r >= Q {
150 r -= Q;
151 }
152 Zq(r)
153 }
154}
155impl<const Q: u64> ops::Add<&Zq<Q>> for &Zq<Q> {
156 type Output = Zq<Q>;
157
158 fn add(self, rhs: &Zq<Q>) -> Self::Output {
159 let mut r = self.0 + rhs.0;
160 if r >= Q {
161 r -= Q;
162 }
163 Zq(r)
164 }
165}
166impl<const Q: u64> ops::AddAssign<Zq<Q>> for Zq<Q> {
167 fn add_assign(&mut self, rhs: Self) {
168 *self = *self + rhs
169 }
170}
171impl<const Q: u64> std::iter::Sum for Zq<Q> {
172 fn sum<I>(iter: I) -> Self
173 where
174 I: Iterator<Item = Self>,
175 {
176 iter.fold(Zq(0), |acc, x| acc + x)
177 }
178}
179impl<const Q: u64> ops::Sub<Zq<Q>> for Zq<Q> {
180 type Output = Self;
181
182 fn sub(self, rhs: Self) -> Zq<Q> {
183 if self.0 >= rhs.0 {
184 Zq(self.0 - rhs.0)
185 } else {
186 Zq((Q + self.0) - rhs.0)
187 }
188 }
189}
190impl<const Q: u64> ops::Sub<&Zq<Q>> for &Zq<Q> {
191 type Output = Zq<Q>;
192
193 fn sub(self, rhs: &Zq<Q>) -> Self::Output {
194 if self.0 >= rhs.0 {
195 Zq(self.0 - rhs.0)
196 } else {
197 Zq((Q + self.0) - rhs.0)
198 }
199 }
200}
201impl<const Q: u64> ops::SubAssign<Zq<Q>> for Zq<Q> {
202 fn sub_assign(&mut self, rhs: Self) {
203 *self = *self - rhs
204 }
205}
206impl<const Q: u64> ops::Neg for Zq<Q> {
207 type Output = Self;
208
209 fn neg(self) -> Self::Output {
210 Zq(Q - self.0)
211 }
212}
213impl<const Q: u64> ops::Mul<Zq<Q>> for Zq<Q> {
214 type Output = Self;
215
216 fn mul(self, rhs: Self) -> Zq<Q> {
217 // TODO non-naive way
218 Zq(((self.0 as u128 * rhs.0 as u128) % Q as u128) as u64)
219 // Zq((self.0 * rhs.0) % Q)
220 }
221}
222impl<const Q: u64> ops::Div<Zq<Q>> for Zq<Q> {
223 type Output = Self;
224
225 fn div(self, rhs: Self) -> Zq<Q> {
226 // TODO non-naive way
227 // Zq((self.0 / rhs.0) % Q)
228 self * rhs.inv()
229 }
230}
231
232impl<const Q: u64> fmt::Display for Zq<Q> {
233 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
234 write!(f, "{}", self.0)
235 }
236}
237impl<const Q: u64> fmt::Debug for Zq<Q> {
238 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
239 write!(f, "{}", self.0)
240 }
241}
242
243#[cfg(test)]
244mod tests {
245 use super::*;
246
247 #[test]
248 fn exp() {
249 const Q: u64 = 1021;
250 let a = Zq::<Q>(3);
251 let b = Zq::<Q>(3);
252 assert_eq!(a.exp(b), Zq(27));
253
254 let a = Zq::<Q>(1000);
255 let b = Zq::<Q>(3);
256 assert_eq!(a.exp(b), Zq(949));
257 }
258 #[test]
259 fn neg() {
260 const Q: u64 = 1021;
261 let a = Zq::<Q>::from_f64(101.0);
262 let b = Zq::<Q>::from_f64(-1.0);
263 assert_eq!(-a, a * b);
264 }
265}