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