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 Param = 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(
91 rq.param.n,
92 rq.coeffs().to_vec().iter().map(|e| e.v).collect(),
93 )
94 }
95}
96
97impl R {
98 // pub fn coeffs(&self) -> [i64; N] {
99 // self.0
100 // }
101 pub fn to_rq(self, q: u64) -> crate::Rq {
102 crate::Rq::from((q, self))
103 }
104
105 // this method is mostly for tests
106 pub fn from_vec_u64(n: usize, coeffs: Vec<u64>) -> Self {
107 let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
108 Self::from_vec(n, coeffs_i64)
109 }
110 pub fn from_vec_f64(n: usize, coeffs: Vec<f64>) -> Self {
111 let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
112 Self::from_vec(n, coeffs_i64)
113 }
114 pub fn new(n: usize, coeffs: Vec<i64>) -> Self {
115 assert_eq!(n, coeffs.len());
116 Self { n, coeffs }
117 }
118 pub fn mul_by_i64(&self, s: i64) -> Self {
119 Self {
120 n: self.n,
121 coeffs: self.coeffs.iter().map(|c_i| c_i * s).collect(),
122 }
123 }
124
125 pub fn infinity_norm(&self) -> u64 {
126 self.coeffs()
127 .iter()
128 // .map(|x| if x.0 > (Q / 2) { Q - x.0 } else { x.0 })
129 .map(|x| x.abs() as u64)
130 .fold(0, |a, b| a.max(b))
131 }
132 pub fn mod_centered_q(&self, q: u64) -> R {
133 let q = q as i64;
134 let r = self
135 .coeffs
136 .iter()
137 .map(|v| {
138 let mut res = v % q;
139 if res > q / 2 {
140 res = res - q;
141 }
142 res
143 })
144 .collect::<Vec<i64>>();
145 R::from_vec(self.n, r)
146 }
147}
148
149pub fn mul_div_round(q: u64, n: usize, v: Vec<i64>, num: u64, den: u64) -> crate::Rq {
150 // dbg!(&v);
151 let r: Vec<f64> = v
152 .iter()
153 .map(|e| ((num as f64 * *e as f64) / den as f64).round())
154 .collect();
155 // dbg!(&r);
156 crate::Rq::from_vec_f64(&crate::ring::RingParam { q, n }, r)
157}
158
159// TODO rename to make it clear that is not mod q, but mod X^N+1
160// apply mod (X^N+1)
161pub fn modulus(n: usize, p: &mut Vec<i64>) {
162 if p.len() < n {
163 return;
164 }
165 for i in n..p.len() {
166 p[i - n] = p[i - n].clone() - p[i].clone();
167 p[i] = 0;
168 }
169 p.truncate(n);
170}
171pub fn modulus_i128(n: usize, p: &mut Vec<i128>) {
172 if p.len() < n {
173 return;
174 }
175 for i in n..p.len() {
176 p[i - n] = p[i - n].clone() - p[i].clone();
177 p[i] = 0;
178 }
179 p.truncate(n);
180}
181
182impl PartialEq for R {
183 fn eq(&self, other: &Self) -> bool {
184 self.coeffs == other.coeffs && self.n == other.n
185 }
186}
187impl Add<R> for R {
188 type Output = Self;
189
190 fn add(self, rhs: Self) -> Self {
191 assert_eq!(self.n, rhs.n);
192 Self {
193 n: self.n,
194 coeffs: zip_eq(self.coeffs, rhs.coeffs)
195 .map(|(l, r)| l + r)
196 .collect(),
197 }
198 }
199}
200impl Add<&R> for &R {
201 type Output = R;
202
203 fn add(self, rhs: &R) -> Self::Output {
204 // R(array::from_fn(|i| self.0[i] + rhs.0[i]))
205 assert_eq!(self.n, rhs.n);
206 R {
207 n: self.n,
208 coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
209 .map(|(l, r)| l + r)
210 .collect(),
211 }
212 }
213}
214impl AddAssign for R {
215 fn add_assign(&mut self, rhs: Self) {
216 assert_eq!(self.n, rhs.n);
217 for i in 0..self.n {
218 self.coeffs[i] += rhs.coeffs[i];
219 }
220 }
221}
222
223impl Sum<R> for R {
224 fn sum<I>(mut iter: I) -> Self
225 where
226 I: Iterator<Item = Self>,
227 {
228 // let mut acc = R::zero();
229 // for e in iter {
230 // acc += e;
231 // }
232 // acc
233 let first = iter.next().unwrap();
234 iter.fold(first, |acc, x| acc + x)
235 }
236}
237
238impl Sub<R> for R {
239 type Output = Self;
240
241 fn sub(self, rhs: Self) -> Self {
242 // Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
243 assert_eq!(self.n, rhs.n);
244 Self {
245 n: self.n,
246 coeffs: zip_eq(self.coeffs, rhs.coeffs)
247 .map(|(l, r)| l - r)
248 .collect(),
249 }
250 }
251}
252impl Sub<&R> for &R {
253 type Output = R;
254
255 fn sub(self, rhs: &R) -> Self::Output {
256 // R(array::from_fn(|i| self.0[i] - rhs.0[i]))
257 assert_eq!(self.n, rhs.n);
258 R {
259 n: self.n,
260 coeffs: zip_eq(&self.coeffs, &rhs.coeffs)
261 .map(|(l, r)| l - r)
262 .collect(),
263 }
264 }
265}
266
267impl SubAssign for R {
268 fn sub_assign(&mut self, rhs: Self) {
269 assert_eq!(self.n, rhs.n);
270 for i in 0..self.n {
271 self.coeffs[i] -= rhs.coeffs[i];
272 }
273 }
274}
275
276impl Mul<R> for R {
277 type Output = Self;
278
279 fn mul(self, rhs: Self) -> Self {
280 naive_poly_mul(&self, &rhs)
281 }
282}
283impl Mul<&R> for &R {
284 type Output = R;
285
286 fn mul(self, rhs: &R) -> Self::Output {
287 naive_poly_mul(self, rhs)
288 }
289}
290
291// TODO WIP
292pub fn naive_poly_mul(poly1: &R, poly2: &R) -> R {
293 assert_eq!(poly1.n, poly2.n);
294 let n = poly1.n;
295
296 let poly1: Vec<i128> = poly1.coeffs.iter().map(|c| *c as i128).collect();
297 let poly2: Vec<i128> = poly2.coeffs.iter().map(|c| *c as i128).collect();
298 let mut result: Vec<i128> = vec![0; (n * 2) - 1];
299 for i in 0..n {
300 for j in 0..n {
301 result[i + j] = result[i + j] + poly1[i] * poly2[j];
302 }
303 }
304
305 // apply mod (X^N + 1))
306 // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
307 modulus_i128(n, &mut result);
308 // dbg!(&result);
309 // dbg!(R::<N>(array::from_fn(|i| result[i] as i64)).coeffs());
310
311 let result_i64: Vec<i64> = result.iter().map(|c_i| *c_i as i64).collect();
312 let r = R::from_vec(n, result_i64);
313 // sanity check: check that there are no coeffs > i64_max
314 assert_eq!(
315 result,
316 r.coeffs.iter().map(|c| *c as i128).collect::<Vec<_>>()
317 );
318 r
319}
320pub fn naive_mul_2(n: usize, poly1: &Vec<i128>, poly2: &Vec<i128>) -> Vec<i128> {
321 let mut result: Vec<i128> = vec![0; (n * 2) - 1];
322 for i in 0..n {
323 for j in 0..n {
324 result[i + j] = result[i + j] + poly1[i] * poly2[j];
325 }
326 }
327
328 // apply mod (X^N + 1))
329 // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
330 modulus_i128(n, &mut result);
331 result
332}
333
334pub fn naive_mul(poly1: &R, poly2: &R) -> Vec<i64> {
335 assert_eq!(poly1.n, poly2.n);
336 let n = poly1.n;
337
338 let poly1: Vec<i128> = poly1.coeffs.iter().map(|c| *c as i128).collect();
339 let poly2: Vec<i128> = poly2.coeffs.iter().map(|c| *c as i128).collect();
340 let mut result = vec![0; (n * 2) - 1];
341 for i in 0..n {
342 for j in 0..n {
343 result[i + j] = result[i + j] + poly1[i] * poly2[j];
344 }
345 }
346 result.iter().map(|c| *c as i64).collect()
347}
348pub fn naive_mul_TMP(poly1: &R, poly2: &R) -> Vec<i64> {
349 assert_eq!(poly1.n, poly2.n);
350 let n = poly1.n;
351
352 let poly1: Vec<i128> = poly1.coeffs.iter().map(|c| *c as i128).collect();
353 let poly2: Vec<i128> = poly2.coeffs.iter().map(|c| *c as i128).collect();
354 let mut result: Vec<i128> = vec![0; (n * 2) - 1];
355 for i in 0..n {
356 for j in 0..n {
357 result[i + j] = result[i + j] + poly1[i] * poly2[j];
358 }
359 }
360
361 // dbg!(&result);
362 modulus_i128(n, &mut result);
363 // for c_i in result.iter() {
364 // println!("---");
365 // println!("{:?}", &c_i);
366 // println!("{:?}", *c_i as i64);
367 // println!("{:?}", (*c_i as i64) as i128);
368 // assert_eq!(*c_i, (*c_i as i64) as i128, "{:?}", c_i);
369 // }
370 result.iter().map(|c| *c as i64).collect()
371}
372
373// wip
374pub fn mod_centered_q(q: u64, n: usize, p: Vec<i128>) -> R {
375 let q: i128 = q as i128;
376 let r = p
377 .iter()
378 .map(|v| {
379 let mut res = v % q;
380 if res > q / 2 {
381 res = res - q;
382 }
383 res
384 })
385 .collect::<Vec<i128>>();
386 R::from_vec(n, r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
387}
388
389impl Mul<i64> for R {
390 type Output = Self;
391
392 fn mul(self, s: i64) -> Self {
393 self.mul_by_i64(s)
394 }
395}
396// mul by u64
397impl Mul<u64> for R {
398 type Output = Self;
399
400 fn mul(self, s: u64) -> Self {
401 self.mul_by_i64(s as i64)
402 }
403}
404impl Mul<&u64> for &R {
405 type Output = R;
406
407 fn mul(self, s: &u64) -> Self::Output {
408 self.mul_by_i64(*s as i64)
409 }
410}
411
412impl Neg for R {
413 type Output = Self;
414
415 fn neg(self) -> Self::Output {
416 // Self(array::from_fn(|i| -self.0[i]))
417 Self {
418 n: self.n,
419 coeffs: self.coeffs.iter().map(|c_i| -c_i).collect(),
420 }
421 }
422}
423
424impl R {
425 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
426 let mut str = "";
427 let mut zero = true;
428 for (i, coeff) in self.coeffs.iter().enumerate().rev() {
429 if *coeff == 0 {
430 continue;
431 }
432 zero = false;
433 f.write_str(str)?;
434 if *coeff != 1 {
435 f.write_str(coeff.to_string().as_str())?;
436 if i > 0 {
437 f.write_str("*")?;
438 }
439 }
440 if *coeff == 1 && i == 0 {
441 f.write_str(coeff.to_string().as_str())?;
442 }
443 if i == 1 {
444 f.write_str("x")?;
445 } else if i > 1 {
446 f.write_str("x^")?;
447 f.write_str(i.to_string().as_str())?;
448 }
449 str = " + ";
450 }
451 if zero {
452 f.write_str("0")?;
453 }
454
455 f.write_str(" mod Z")?;
456 f.write_str("/(X^")?;
457 f.write_str(self.n.to_string().as_str())?;
458 f.write_str("+1)")?;
459 Ok(())
460 }
461}
462impl fmt::Display for R {
463 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
464 self.fmt(f)?;
465 Ok(())
466 }
467}
468impl fmt::Debug for R {
469 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
470 self.fmt(f)?;
471 Ok(())
472 }
473}
474
475#[cfg(test)]
476mod tests {
477 use super::*;
478 use anyhow::Result;
479
480 #[test]
481 fn test_mul() -> Result<()> {
482 let n: usize = 2;
483 let q: i64 = (2u64.pow(16) + 1) as i64;
484
485 // test vectors generated with SageMath
486 let a: Vec<i64> = vec![q - 1, q - 1];
487 let b: Vec<i64> = vec![q - 1, q - 1];
488 let c: Vec<i64> = vec![0, 8589934592];
489 test_mul_opt(n, a, b, c)?;
490
491 let a: Vec<i64> = vec![1, q - 1];
492 let b: Vec<i64> = vec![1, q - 1];
493 let c: Vec<i64> = vec![-4294967295, 131072];
494 test_mul_opt(n, a, b, c)?;
495
496 Ok(())
497 }
498 fn test_mul_opt(n: usize, a: Vec<i64>, b: Vec<i64>, expected_c: Vec<i64>) -> Result<()> {
499 let mut a = R::new(n, a);
500 let mut b = R::new(n, b);
501 dbg!(&a);
502 dbg!(&b);
503 let expected_c = R::new(n, expected_c);
504
505 let mut c = naive_mul(&mut a, &mut b);
506 modulus(n, &mut c);
507 dbg!(R::from_vec(n, c.clone()));
508 assert_eq!(c, expected_c.coeffs);
509 Ok(())
510 }
511}