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