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