fhe-study/arith/src/ring.rs
arith/src/ring.rsRaw
1//! Polynomial ring Z[X]/(X^N+1)
2//!
3
4use anyhow::Result;
5use rand::{distributions::Distribution, Rng};
6use std::array;
7use std::fmt;
8use std::iter::Sum;
9use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
10
11use crate::Ring;
12
13// TODO rename to not have name conflicts with the Ring trait (R: Ring)
14// PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1)
15#[derive(Clone, Copy)]
16pub struct R<const N: usize>(pub [i64; N]);
17
18impl<const N: usize> Ring for R<N> {
19 type C = i64;
20 fn coeffs(&self) -> Vec<Self::C> {
21 self.0.to_vec()
22 }
23 fn zero() -> Self {
24 let coeffs: [i64; N] = array::from_fn(|_| 0i64);
25 Self(coeffs)
26 }
27 fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
28 // let coeffs: [i64; N] = array::from_fn(|_| Self::C::rand(&mut rng, &dist));
29 let coeffs: [i64; N] = array::from_fn(|_| dist.sample(&mut rng).round() as i64);
30 Self(coeffs)
31 // let coeffs: [C; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
32 // Self(coeffs)
33 }
34
35 // returns the decomposition of each polynomial coefficient
36 fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
37 unimplemented!();
38 // array::from_fn(|i| self.coeffs[i].decompose(beta, l))
39 }
40}
41
42impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> {
43 fn from(rq: crate::ringq::Rq<Q, N>) -> Self {
44 Self::from_vec_u64(rq.coeffs().to_vec().iter().map(|e| e.0).collect())
45 }
46}
47
48impl<const N: usize> R<N> {
49 pub fn coeffs(&self) -> [i64; N] {
50 self.0
51 }
52 pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
53 crate::Rq::<Q, N>::from(self)
54 }
55
56 pub fn from_vec(coeffs: Vec<i64>) -> Self {
57 let mut p = coeffs;
58 modulus::<N>(&mut p);
59 Self(array::from_fn(|i| p[i]))
60 }
61 // this method is mostly for tests
62 pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
63 let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
64 Self::from_vec(coeffs_i64)
65 }
66 pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
67 let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
68 Self::from_vec(coeffs_i64)
69 }
70 pub fn new(coeffs: [i64; N]) -> Self {
71 Self(coeffs)
72 }
73 pub fn mul_by_i64(&self, s: i64) -> Self {
74 Self(array::from_fn(|i| self.0[i] * s))
75 }
76 // performs the multiplication and division over f64, and then it rounds the
77 // result, only applying the mod Q at the end
78 pub fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> {
79 let r: Vec<f64> = self
80 .coeffs()
81 .iter()
82 .map(|e| ((num as f64 * *e as f64) / den as f64).round())
83 .collect();
84 crate::Rq::<Q, N>::from_vec_f64(r)
85 }
86
87 pub fn infinity_norm(&self) -> u64 {
88 self.coeffs()
89 .iter()
90 // .map(|x| if x.0 > (Q / 2) { Q - x.0 } else { x.0 })
91 .map(|x| x.abs() as u64)
92 .fold(0, |a, b| a.max(b))
93 }
94 pub fn mod_centered_q<const Q: u64>(&self) -> R<N> {
95 let q = Q as i64;
96 let r = self
97 .0
98 .iter()
99 .map(|v| {
100 let mut res = v % q;
101 if res > q / 2 {
102 res = res - q;
103 }
104 res
105 })
106 .collect::<Vec<i64>>();
107 R::<N>::from_vec(r)
108 }
109}
110
111pub fn mul_div_round<const Q: u64, const N: usize>(
112 v: Vec<i64>,
113 num: u64,
114 den: u64,
115) -> crate::Rq<Q, N> {
116 // dbg!(&v);
117 let r: Vec<f64> = v
118 .iter()
119 .map(|e| ((num as f64 * *e as f64) / den as f64).round())
120 .collect();
121 // dbg!(&r);
122 crate::Rq::<Q, N>::from_vec_f64(r)
123}
124
125// TODO rename to make it clear that is not mod q, but mod X^N+1
126// apply mod (X^N+1)
127pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
128 if p.len() < N {
129 return;
130 }
131 for i in N..p.len() {
132 p[i - N] = p[i - N].clone() - p[i].clone();
133 p[i] = 0;
134 }
135 p.truncate(N);
136}
137pub fn modulus_i128<const N: usize>(p: &mut Vec<i128>) {
138 if p.len() < N {
139 return;
140 }
141 for i in N..p.len() {
142 p[i - N] = p[i - N].clone() - p[i].clone();
143 p[i] = 0;
144 }
145 p.truncate(N);
146}
147
148impl<const N: usize> PartialEq for R<N> {
149 fn eq(&self, other: &Self) -> bool {
150 self.0 == other.0
151 }
152}
153impl<const N: usize> Add<R<N>> for R<N> {
154 type Output = Self;
155
156 fn add(self, rhs: Self) -> Self {
157 Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
158 }
159}
160impl<const N: usize> Add<&R<N>> for &R<N> {
161 type Output = R<N>;
162
163 fn add(self, rhs: &R<N>) -> Self::Output {
164 R(array::from_fn(|i| self.0[i] + rhs.0[i]))
165 }
166}
167impl<const N: usize> AddAssign for R<N> {
168 fn add_assign(&mut self, rhs: Self) {
169 for i in 0..N {
170 self.0[i] += rhs.0[i];
171 }
172 }
173}
174
175impl<const N: usize> Sum<R<N>> for R<N> {
176 fn sum<I>(iter: I) -> Self
177 where
178 I: Iterator<Item = Self>,
179 {
180 let mut acc = R::<N>::zero();
181 for e in iter {
182 acc += e;
183 }
184 acc
185 }
186}
187
188impl<const N: usize> Sub<R<N>> for R<N> {
189 type Output = Self;
190
191 fn sub(self, rhs: Self) -> Self {
192 Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
193 }
194}
195impl<const N: usize> Sub<&R<N>> for &R<N> {
196 type Output = R<N>;
197
198 fn sub(self, rhs: &R<N>) -> Self::Output {
199 R(array::from_fn(|i| self.0[i] - rhs.0[i]))
200 }
201}
202
203impl<const N: usize> SubAssign for R<N> {
204 fn sub_assign(&mut self, rhs: Self) {
205 for i in 0..N {
206 self.0[i] -= rhs.0[i];
207 }
208 }
209}
210
211impl<const N: usize> Mul<R<N>> for R<N> {
212 type Output = Self;
213
214 fn mul(self, rhs: Self) -> Self {
215 naive_poly_mul(&self, &rhs)
216 }
217}
218impl<const N: usize> Mul<&R<N>> for &R<N> {
219 type Output = R<N>;
220
221 fn mul(self, rhs: &R<N>) -> Self::Output {
222 naive_poly_mul(self, rhs)
223 }
224}
225
226// TODO WIP
227pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
228 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
229 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
230 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
231 for i in 0..N {
232 for j in 0..N {
233 result[i + j] = result[i + j] + poly1[i] * poly2[j];
234 }
235 }
236
237 // apply mod (X^N + 1))
238 // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
239 modulus_i128::<N>(&mut result);
240 // dbg!(&result);
241 // dbg!(R::<N>(array::from_fn(|i| result[i] as i64)).coeffs());
242
243 // sanity check: check that there are no coeffs > i64_max
244 assert_eq!(
245 result,
246 R::<N>(array::from_fn(|i| result[i] as i64))
247 .coeffs()
248 .iter()
249 .map(|c| *c as i128)
250 .collect::<Vec<_>>()
251 );
252 R(array::from_fn(|i| result[i] as i64))
253}
254pub fn naive_mul_2<const N: usize>(poly1: &Vec<i128>, poly2: &Vec<i128>) -> Vec<i128> {
255 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
256 for i in 0..N {
257 for j in 0..N {
258 result[i + j] = result[i + j] + poly1[i] * poly2[j];
259 }
260 }
261
262 // apply mod (X^N + 1))
263 // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
264 modulus_i128::<N>(&mut result);
265 result
266}
267
268pub fn naive_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
269 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
270 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
271 let mut result = vec![0; (N * 2) - 1];
272 for i in 0..N {
273 for j in 0..N {
274 result[i + j] = result[i + j] + poly1[i] * poly2[j];
275 }
276 }
277 result.iter().map(|c| *c as i64).collect()
278}
279pub fn naive_mul_TMP<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
280 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
281 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
282 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
283 for i in 0..N {
284 for j in 0..N {
285 result[i + j] = result[i + j] + poly1[i] * poly2[j];
286 }
287 }
288
289 // dbg!(&result);
290 modulus_i128::<N>(&mut result);
291 // for c_i in result.iter() {
292 // println!("---");
293 // println!("{:?}", &c_i);
294 // println!("{:?}", *c_i as i64);
295 // println!("{:?}", (*c_i as i64) as i128);
296 // assert_eq!(*c_i, (*c_i as i64) as i128, "{:?}", c_i);
297 // }
298 result.iter().map(|c| *c as i64).collect()
299}
300
301// wip
302pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
303 let q: i128 = Q as i128;
304 let r = p
305 .iter()
306 .map(|v| {
307 let mut res = v % q;
308 if res > q / 2 {
309 res = res - q;
310 }
311 res
312 })
313 .collect::<Vec<i128>>();
314 R::<N>::from_vec(r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
315}
316
317// mul by u64
318impl<const N: usize> Mul<u64> for R<N> {
319 type Output = Self;
320
321 fn mul(self, s: u64) -> Self {
322 self.mul_by_i64(s as i64)
323 }
324}
325impl<const N: usize> Mul<&u64> for &R<N> {
326 type Output = R<N>;
327
328 fn mul(self, s: &u64) -> Self::Output {
329 self.mul_by_i64(*s as i64)
330 }
331}
332
333impl<const N: usize> Neg for R<N> {
334 type Output = Self;
335
336 fn neg(self) -> Self::Output {
337 Self(array::from_fn(|i| -self.0[i]))
338 }
339}
340
341impl<const N: usize> R<N> {
342 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
343 let mut str = "";
344 let mut zero = true;
345 for (i, coeff) in self.0.iter().enumerate().rev() {
346 if *coeff == 0 {
347 continue;
348 }
349 zero = false;
350 f.write_str(str)?;
351 if *coeff != 1 {
352 f.write_str(coeff.to_string().as_str())?;
353 if i > 0 {
354 f.write_str("*")?;
355 }
356 }
357 if *coeff == 1 && i == 0 {
358 f.write_str(coeff.to_string().as_str())?;
359 }
360 if i == 1 {
361 f.write_str("x")?;
362 } else if i > 1 {
363 f.write_str("x^")?;
364 f.write_str(i.to_string().as_str())?;
365 }
366 str = " + ";
367 }
368 if zero {
369 f.write_str("0")?;
370 }
371
372 f.write_str(" mod Z")?;
373 f.write_str("/(X^")?;
374 f.write_str(N.to_string().as_str())?;
375 f.write_str("+1)")?;
376 Ok(())
377 }
378}
379impl<const N: usize> fmt::Display for R<N> {
380 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
381 self.fmt(f)?;
382 Ok(())
383 }
384}
385impl<const N: usize> fmt::Debug for R<N> {
386 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
387 self.fmt(f)?;
388 Ok(())
389 }
390}
391
392#[cfg(test)]
393mod tests {
394 use super::*;
395 use anyhow::Result;
396
397 #[test]
398 fn test_mul() -> Result<()> {
399 const Q: u64 = 2u64.pow(16) + 1;
400 const N: usize = 2;
401 let q: i64 = Q as i64;
402
403 // test vectors generated with SageMath
404 let a: [i64; N] = [q - 1, q - 1];
405 let b: [i64; N] = [q - 1, q - 1];
406 let c: [i64; N] = [0, 8589934592];
407 test_mul_opt::<Q, N>(a, b, c)?;
408
409 let a: [i64; N] = [1, q - 1];
410 let b: [i64; N] = [1, q - 1];
411 let c: [i64; N] = [-4294967295, 131072];
412 test_mul_opt::<Q, N>(a, b, c)?;
413
414 Ok(())
415 }
416 fn test_mul_opt<const Q: u64, const N: usize>(
417 a: [i64; N],
418 b: [i64; N],
419 expected_c: [i64; N],
420 ) -> Result<()> {
421 let mut a = R::new(a);
422 let mut b = R::new(b);
423 dbg!(&a);
424 dbg!(&b);
425 let expected_c = R::new(expected_c);
426
427 let mut c = naive_mul(&mut a, &mut b);
428 modulus::<N>(&mut c);
429 dbg!(R::<N>::from_vec(c.clone()));
430 assert_eq!(c, expected_c.0.to_vec());
431 Ok(())
432 }
433}