fhe-study/arith/src/zq.rs
arith/src/zq.rsRaw
1use std::fmt;
2use std::ops;
3
4// Z_q, integers modulus q, not necessarily prime
5#[derive(Clone, Copy, PartialEq)]
6pub struct Zq<const Q: u64>(pub u64);
7
8// WIP
9// impl<const Q: u64> From<Vec<u64>> for Vec<Zq<Q>> {
10// fn from(v: Vec<u64>) -> Self {
11// v.into_iter().map(Zq::new).collect()
12// }
13// }
14
15pub(crate) fn modulus_u64<const Q: u64>(e: u64) -> u64 {
16 (e % Q + Q) % Q
17}
18impl<const Q: u64> Zq<Q> {
19 pub fn from_u64(e: u64) -> Self {
20 if e >= Q {
21 // (e % Q + Q) % Q
22 return Zq(modulus_u64::<Q>(e));
23 // return Zq(e % Q);
24 }
25 Zq(e)
26 }
27 pub fn from_f64(e: f64) -> Self {
28 // WIP method
29 let e: i64 = e.round() as i64;
30 let q = Q as i64;
31 if e < 0 || e >= q {
32 return Zq(((e % q + q) % q) as u64);
33 }
34 Zq(e as u64)
35
36 // if e < 0 {
37 // // dbg!(&e);
38 // // dbg!(Zq::<Q>(((Q as i64 + e) % Q as i64) as u64));
39 // // return Zq(((Q as i64 + e) % Q as i64) as u64);
40 // return Zq(e as u64 % Q);
41 // } else if e >= Q as i64 {
42 // return Zq((e % Q as i64) as u64);
43 // }
44 // Zq(e as u64)
45 }
46 pub fn from_bool(b: bool) -> Self {
47 if b {
48 Zq(1)
49 } else {
50 Zq(0)
51 }
52 }
53 pub fn zero() -> Self {
54 Zq(0u64)
55 }
56 pub fn square(self) -> Self {
57 self * self
58 }
59 // modular exponentiation
60 pub fn exp(self, e: Self) -> Self {
61 // mul-square approach
62 let mut res = Self(1);
63 let mut rem = e.clone();
64 let mut exp = self;
65 // for rem != Self(0) {
66 while rem != Self(0) {
67 // if odd
68 // TODO use a more readible expression
69 if 1 - ((rem.0 & 1) << 1) as i64 == -1 {
70 res = res * exp;
71 }
72 exp = exp.square();
73 rem = Self(rem.0 >> 1);
74 }
75 res
76 }
77 // multiplicative inverse
78 // WARNING: if this is needed, it means that 'Zq' is a Finite Field. For the moment we assume
79 // we work in a Finite Field
80 pub fn inv_OLD(self) -> Self {
81 // TODO
82 // let a = self.0;
83 // let q = Q;
84 let mut t = 0;
85 let mut r = Q;
86 let mut new_t = 0;
87 let mut new_r = self.0.clone();
88 while new_r != 0 {
89 let q = r / new_r;
90
91 t = new_t.clone();
92 new_t = t - q;
93
94 r = new_r.clone();
95 new_r = r - (q * new_r);
96 }
97 // if t < 0 {
98 // t = t + q;
99 // }
100 return Zq::from_u64(t);
101 }
102 pub fn inv(self) -> Zq<Q> {
103 let (g, x, _) = Self::egcd(self.0 as i128, Q as i128);
104 if g != 1 {
105 // None
106 panic!("E");
107 } else {
108 let q = Q as i128;
109 Zq(((x % q + q) % q) as u64) // TODO maybe just Zq::new(x)
110 }
111 }
112 fn egcd(a: i128, b: i128) -> (i128, i128, i128) {
113 if a == 0 {
114 (b, 0, 1)
115 } else {
116 let (g, x, y) = Self::egcd(b % a, a);
117 (g, y - (b / a) * x, x)
118 }
119 }
120
121 /// perform the mod switch operation from Q to Q', where Q2=Q'
122 pub fn mod_switch<const Q2: u64>(&self) -> Zq<Q2> {
123 Zq::<Q2>::from_u64(((self.0 as f64 * Q2 as f64) / Q as f64).round() as u64)
124 }
125}
126
127impl<const Q: u64> Zq<Q> {
128 fn r#mod(self) -> Self {
129 if self.0 >= Q {
130 return Zq(self.0 % Q);
131 }
132 self
133 }
134}
135
136impl<const Q: u64> ops::Add<Zq<Q>> for Zq<Q> {
137 type Output = Self;
138
139 fn add(self, rhs: Self) -> Self::Output {
140 let mut r = self.0 + rhs.0;
141 if r >= Q {
142 r -= Q;
143 }
144 Zq(r)
145 }
146}
147impl<const Q: u64> ops::Add<&Zq<Q>> for &Zq<Q> {
148 type Output = Zq<Q>;
149
150 fn add(self, rhs: &Zq<Q>) -> Self::Output {
151 let mut r = self.0 + rhs.0;
152 if r >= Q {
153 r -= Q;
154 }
155 Zq(r)
156 }
157}
158impl<const Q: u64> ops::AddAssign<Zq<Q>> for Zq<Q> {
159 fn add_assign(&mut self, rhs: Self) {
160 *self = *self + rhs
161 }
162}
163impl<const Q: u64> std::iter::Sum for Zq<Q> {
164 fn sum<I>(iter: I) -> Self
165 where
166 I: Iterator<Item = Self>,
167 {
168 iter.fold(Zq(0), |acc, x| acc + x)
169 }
170}
171impl<const Q: u64> ops::Sub<Zq<Q>> for Zq<Q> {
172 type Output = Self;
173
174 fn sub(self, rhs: Self) -> Zq<Q> {
175 if self.0 >= rhs.0 {
176 Zq(self.0 - rhs.0)
177 } else {
178 Zq((Q + self.0) - rhs.0)
179 }
180 }
181}
182impl<const Q: u64> ops::Sub<&Zq<Q>> for &Zq<Q> {
183 type Output = Zq<Q>;
184
185 fn sub(self, rhs: &Zq<Q>) -> Self::Output {
186 if self.0 >= rhs.0 {
187 Zq(self.0 - rhs.0)
188 } else {
189 Zq((Q + self.0) - rhs.0)
190 }
191 }
192}
193impl<const Q: u64> ops::SubAssign<Zq<Q>> for Zq<Q> {
194 fn sub_assign(&mut self, rhs: Self) {
195 *self = *self - rhs
196 }
197}
198impl<const Q: u64> ops::Neg for Zq<Q> {
199 type Output = Self;
200
201 fn neg(self) -> Self::Output {
202 Zq(Q - self.0)
203 }
204}
205impl<const Q: u64> ops::Mul<Zq<Q>> for Zq<Q> {
206 type Output = Self;
207
208 fn mul(self, rhs: Self) -> Zq<Q> {
209 // TODO non-naive way
210 Zq(((self.0 as u128 * rhs.0 as u128) % Q as u128) as u64)
211 // Zq((self.0 * rhs.0) % Q)
212 }
213}
214impl<const Q: u64> ops::Div<Zq<Q>> for Zq<Q> {
215 type Output = Self;
216
217 fn div(self, rhs: Self) -> Zq<Q> {
218 // TODO non-naive way
219 // Zq((self.0 / rhs.0) % Q)
220 self * rhs.inv()
221 }
222}
223
224impl<const Q: u64> fmt::Display for Zq<Q> {
225 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
226 write!(f, "{}", self.0)
227 }
228}
229impl<const Q: u64> fmt::Debug for Zq<Q> {
230 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
231 write!(f, "{}", self.0)
232 }
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238
239 #[test]
240 fn exp() {
241 const Q: u64 = 1021;
242 let a = Zq::<Q>(3);
243 let b = Zq::<Q>(3);
244 assert_eq!(a.exp(b), Zq(27));
245
246 let a = Zq::<Q>(1000);
247 let b = Zq::<Q>(3);
248 assert_eq!(a.exp(b), Zq(949));
249 }
250 #[test]
251 fn neg() {
252 const Q: u64 = 1021;
253 let a = Zq::<Q>::from_f64(101.0);
254 let b = Zq::<Q>::from_f64(-1.0);
255 assert_eq!(-a, a * b);
256 }
257}