fhe-study/arith/src/ringq.rs
arith/src/ringq.rsRaw
1//! Polynomial ring Z_q[X]/(X^N+1)
2//!
3
4use rand::{distributions::Distribution, Rng};
5use std::array;
6use std::fmt;
7use std::iter::Sum;
8use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
9
10use crate::ntt::NTT;
11use crate::zq::{modulus_u64, Zq};
12use anyhow::{anyhow, Result};
13
14use crate::Ring;
15
16// NOTE: currently using fixed-size arrays, but pending to see if with
17// real-world parameters the stack can keep up; if not will move everything to
18// use Vec.
19/// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
20/// The implementation assumes that q is prime.
21#[derive(Clone, Copy)]
22pub struct Rq<const Q: u64, const N: usize> {
23 pub(crate) coeffs: [Zq<Q>; N],
24
25 // evals are set when doig a PRxPR multiplication, so it can be reused in future
26 // multiplications avoiding recomputing it
27 pub(crate) evals: Option<[Zq<Q>; N]>,
28}
29
30impl<const Q: u64, const N: usize> Ring for Rq<Q, N> {
31 type C = Zq<Q>;
32 fn coeffs(&self) -> Vec<Self::C> {
33 self.coeffs.to_vec()
34 }
35 fn zero() -> Self {
36 let coeffs = array::from_fn(|_| Zq::zero());
37 Self {
38 coeffs,
39 evals: None,
40 }
41 }
42 fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
43 // let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
44 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Self::C::rand(&mut rng, &dist));
45 Self {
46 coeffs,
47 evals: None,
48 }
49 }
50}
51
52// TODO define a trait "PolynomialRingTrait" or similar, so that when other structs use it can just
53// use the trait and not need to add '<Q, N>' to their params
54
55impl<const Q: u64, const N: usize> From<crate::ring::R<N>> for Rq<Q, N> {
56 fn from(r: crate::ring::R<N>) -> Self {
57 Self::from_vec(
58 r.coeffs()
59 .iter()
60 .map(|e| Zq::<Q>::from_f64(*e as f64))
61 .collect(),
62 )
63 }
64}
65
66// apply mod (X^N+1)
67pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
68 if p.len() < N {
69 return;
70 }
71 for i in N..p.len() {
72 p[i - N] = p[i - N].clone() - p[i].clone();
73 p[i] = Zq(0);
74 }
75 p.truncate(N);
76}
77
78// PR stands for PolynomialRing
79impl<const Q: u64, const N: usize> Rq<Q, N> {
80 pub fn coeffs(&self) -> [Zq<Q>; N] {
81 self.coeffs
82 }
83 pub fn compute_evals(&mut self) {
84 self.evals = Some(NTT::<Q, N>::ntt(self.coeffs));
85 }
86 pub fn to_r(self) -> crate::R<N> {
87 crate::R::<N>::from(self)
88 }
89
90 // TODO rm since it is implemented in Ring trait impl
91 // pub fn zero() -> Self {
92 // let coeffs = array::from_fn(|_| Zq::zero());
93 // Self {
94 // coeffs,
95 // evals: None,
96 // }
97 // }
98 pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
99 let mut p = coeffs;
100 modulus::<Q, N>(&mut p);
101 let coeffs = array::from_fn(|i| p[i]);
102 Self {
103 coeffs,
104 evals: None,
105 }
106 }
107 // this method is mostly for tests
108 pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
109 let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_u64(*c)).collect();
110 Self::from_vec(coeffs_mod_q)
111 }
112 pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
113 let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_f64(*c)).collect();
114 Self::from_vec(coeffs_mod_q)
115 }
116 pub fn from_vec_i64(coeffs: Vec<i64>) -> Self {
117 let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_f64(*c as f64)).collect();
118 Self::from_vec(coeffs_mod_q)
119 }
120 pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
121 Self { coeffs, evals }
122 }
123
124 pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
125 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
126 Ok(Self {
127 coeffs,
128 evals: None,
129 })
130 }
131 pub fn rand_f64_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
132 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
133 Ok(Self {
134 coeffs,
135 evals: None,
136 })
137 }
138 pub fn rand_f64(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
139 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
140 Ok(Self {
141 coeffs,
142 evals: None,
143 })
144 }
145 pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Result<Self> {
146 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
147 Ok(Self {
148 coeffs,
149 evals: None,
150 })
151 }
152 // WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
153 pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
154 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
155 Ok(Rq {
156 coeffs,
157 evals: None,
158 })
159 }
160 // Warning: this method will behave differently depending on the values P and Q:
161 // if Q<P, it just 'renames' the modulus parameter to P
162 // if Q>=P, it crops to mod P
163 pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
164 Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
165 }
166
167 /// perform the mod switch operation from Q to Q', where Q2=Q'
168 fn mod_switch<const Q2: u64>(&self) -> Rq<Q2, N> {
169 Rq::<Q2, N> {
170 coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<Q2>()),
171 evals: None,
172 }
173 }
174
175 // applies mod(T) to all coefficients of self
176 pub fn coeffs_mod<const T: u64>(&self) -> Self {
177 Rq::<Q, N>::from_vec_u64(
178 self.coeffs()
179 .iter()
180 .map(|m_i| modulus_u64::<T>(m_i.0))
181 .collect(),
182 )
183 }
184
185 // TODO review if needed, or if with this interface
186 pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq<Q>>>) -> Result<Vec<Zq<Q>>> {
187 matrix_vec_product(m, &self.coeffs.to_vec())
188 }
189 pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
190 Self {
191 coeffs: array::from_fn(|i| self.coeffs[i] * *s),
192 evals: None,
193 }
194 }
195 pub fn mul_by_u64(&self, s: u64) -> Self {
196 let s = Zq::from_u64(s);
197 Self {
198 coeffs: array::from_fn(|i| self.coeffs[i] * s),
199 // coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
200 evals: None,
201 }
202 }
203 pub fn mul_by_f64(&self, s: f64) -> Self {
204 Self {
205 coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
206 evals: None,
207 }
208 }
209
210 pub fn mul(&mut self, rhs: &mut Self) -> Self {
211 mul_mut(self, rhs)
212 }
213 // divides by the given scalar 's' and rounds, returning a Rq<Q,N>
214 // TODO rm
215 pub fn div_round(&self, s: u64) -> Self {
216 let r: Vec<f64> = self
217 .coeffs()
218 .iter()
219 .map(|e| (e.0 as f64 / s as f64).round())
220 .collect();
221 Rq::<Q, N>::from_vec_f64(r)
222 }
223 // returns [ [(num/den) * self].round() ] mod q
224 // ie. performs the multiplication and division over f64, and then it rounds the
225 // result, only applying the mod Q at the end
226 pub fn mul_div_round(&self, num: u64, den: u64) -> Self {
227 let r: Vec<f64> = self
228 .coeffs()
229 .iter()
230 .map(|e| ((num as f64 * e.0 as f64) / den as f64).round())
231 .collect();
232 Rq::<Q, N>::from_vec_f64(r)
233 }
234
235 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
236 // TODO simplify
237 let mut str = "";
238 let mut zero = true;
239 for (i, coeff) in self.coeffs.iter().enumerate().rev() {
240 if coeff.0 == 0 {
241 continue;
242 }
243 zero = false;
244 f.write_str(str)?;
245 if coeff.0 != 1 {
246 f.write_str(coeff.0.to_string().as_str())?;
247 if i > 0 {
248 f.write_str("*")?;
249 }
250 }
251 if coeff.0 == 1 && i == 0 {
252 f.write_str(coeff.0.to_string().as_str())?;
253 }
254 if i == 1 {
255 f.write_str("x")?;
256 } else if i > 1 {
257 f.write_str("x^")?;
258 f.write_str(i.to_string().as_str())?;
259 }
260 str = " + ";
261 }
262 if zero {
263 f.write_str("0")?;
264 }
265
266 f.write_str(" mod Z_")?;
267 f.write_str(Q.to_string().as_str())?;
268 f.write_str("/(X^")?;
269 f.write_str(N.to_string().as_str())?;
270 f.write_str("+1)")?;
271 Ok(())
272 }
273
274 pub fn infinity_norm(&self) -> u64 {
275 self.coeffs()
276 .iter()
277 .map(|x| if x.0 > (Q / 2) { Q - x.0 } else { x.0 })
278 .fold(0, |a, b| a.max(b))
279 }
280 pub fn mod_centered_q(&self) -> crate::ring::R<N> {
281 self.to_r().mod_centered_q::<Q>()
282 }
283}
284pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
285 // assert_eq!(m.len(), m[0].len()); // TODO change to returning err
286 // assert_eq!(m.len(), v.len());
287 if m.len() != m[0].len() {
288 return Err(anyhow!("expected 'm' to be a square matrix"));
289 }
290 if m.len() != v.len() {
291 return Err(anyhow!(
292 "m.len: {} should be equal to v.len(): {}",
293 m.len(),
294 v.len(),
295 ));
296 }
297
298 Ok(m.iter()
299 .map(|row| {
300 row.iter()
301 .zip(v.iter())
302 .map(|(&row_i, &v_i)| row_i * v_i)
303 .sum()
304 })
305 .collect::<Vec<Zq<Q>>>())
306}
307pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
308 // TODO case when m[0].len()=0
309 // TODO non square matrix
310 let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
311 for (i, m_row) in m.iter().enumerate() {
312 for (j, m_ij) in m_row.iter().enumerate() {
313 r[j][i] = *m_ij;
314 }
315 }
316 r
317}
318
319impl<const Q: u64, const N: usize> PartialEq for Rq<Q, N> {
320 fn eq(&self, other: &Self) -> bool {
321 self.coeffs == other.coeffs
322 }
323}
324impl<const Q: u64, const N: usize> Add<Rq<Q, N>> for Rq<Q, N> {
325 type Output = Self;
326
327 fn add(self, rhs: Self) -> Self {
328 Self {
329 coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
330 evals: None,
331 }
332 // Self {
333 // coeffs: self
334 // .coeffs
335 // .iter()
336 // .zip(rhs.coeffs)
337 // .map(|(a, b)| *a + b)
338 // .collect(),
339 // evals: None,
340 // }
341 // Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in +
342 }
343}
344impl<const Q: u64, const N: usize> Add<&Rq<Q, N>> for &Rq<Q, N> {
345 type Output = Rq<Q, N>;
346
347 fn add(self, rhs: &Rq<Q, N>) -> Self::Output {
348 Rq {
349 coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
350 evals: None,
351 }
352 }
353}
354impl<const Q: u64, const N: usize> AddAssign for Rq<Q, N> {
355 fn add_assign(&mut self, rhs: Self) {
356 for i in 0..N {
357 self.coeffs[i] += rhs.coeffs[i];
358 }
359 }
360}
361
362impl<const Q: u64, const N: usize> Sum<Rq<Q, N>> for Rq<Q, N> {
363 fn sum<I>(iter: I) -> Self
364 where
365 I: Iterator<Item = Self>,
366 {
367 let mut acc = Rq::<Q, N>::zero();
368 for e in iter {
369 acc += e;
370 }
371 acc
372 }
373}
374
375impl<const Q: u64, const N: usize> Sub<Rq<Q, N>> for Rq<Q, N> {
376 type Output = Self;
377
378 fn sub(self, rhs: Self) -> Self {
379 Self {
380 coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
381 evals: None,
382 }
383 }
384}
385impl<const Q: u64, const N: usize> Sub<&Rq<Q, N>> for &Rq<Q, N> {
386 type Output = Rq<Q, N>;
387
388 fn sub(self, rhs: &Rq<Q, N>) -> Self::Output {
389 Rq {
390 coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
391 evals: None,
392 }
393 }
394}
395impl<const Q: u64, const N: usize> SubAssign for Rq<Q, N> {
396 fn sub_assign(&mut self, rhs: Self) {
397 for i in 0..N {
398 self.coeffs[i] -= rhs.coeffs[i];
399 }
400 }
401}
402
403impl<const Q: u64, const N: usize> Mul<Rq<Q, N>> for Rq<Q, N> {
404 type Output = Self;
405
406 fn mul(self, rhs: Self) -> Self {
407 mul(&self, &rhs)
408 }
409}
410impl<const Q: u64, const N: usize> Mul<&Rq<Q, N>> for &Rq<Q, N> {
411 type Output = Rq<Q, N>;
412
413 fn mul(self, rhs: &Rq<Q, N>) -> Self::Output {
414 mul(self, rhs)
415 }
416}
417
418// mul by Zq element
419impl<const Q: u64, const N: usize> Mul<Zq<Q>> for Rq<Q, N> {
420 type Output = Self;
421
422 fn mul(self, s: Zq<Q>) -> Self {
423 self.mul_by_zq(&s)
424 }
425}
426impl<const Q: u64, const N: usize> Mul<&Zq<Q>> for &Rq<Q, N> {
427 type Output = Rq<Q, N>;
428
429 fn mul(self, s: &Zq<Q>) -> Self::Output {
430 self.mul_by_zq(s)
431 }
432}
433// mul by u64
434impl<const Q: u64, const N: usize> Mul<u64> for Rq<Q, N> {
435 type Output = Self;
436
437 fn mul(self, s: u64) -> Self {
438 self.mul_by_u64(s)
439 }
440}
441impl<const Q: u64, const N: usize> Mul<&u64> for &Rq<Q, N> {
442 type Output = Rq<Q, N>;
443
444 fn mul(self, s: &u64) -> Self::Output {
445 self.mul_by_u64(*s)
446 }
447}
448// mul by f64
449impl<const Q: u64, const N: usize> Mul<f64> for Rq<Q, N> {
450 type Output = Self;
451
452 fn mul(self, s: f64) -> Self {
453 self.mul_by_f64(s)
454 }
455}
456impl<const Q: u64, const N: usize> Mul<&f64> for &Rq<Q, N> {
457 type Output = Rq<Q, N>;
458
459 fn mul(self, s: &f64) -> Self::Output {
460 self.mul_by_f64(*s)
461 }
462}
463
464impl<const Q: u64, const N: usize> Neg for Rq<Q, N> {
465 type Output = Self;
466
467 fn neg(self) -> Self::Output {
468 Self {
469 coeffs: array::from_fn(|i| -self.coeffs[i]),
470 evals: None,
471 }
472 }
473}
474
475// note: this assumes that Q is prime
476fn mul_mut<const Q: u64, const N: usize>(lhs: &mut Rq<Q, N>, rhs: &mut Rq<Q, N>) -> Rq<Q, N> {
477 // reuse evaluations if already computed
478 if !lhs.evals.is_some() {
479 lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
480 };
481 if !rhs.evals.is_some() {
482 rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
483 };
484 let lhs_evals = lhs.evals.unwrap();
485 let rhs_evals = rhs.evals.unwrap();
486
487 let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
488 let c = NTT::<Q, { N }>::intt(c_ntt);
489 Rq::new(c, Some(c_ntt))
490}
491// note: this assumes that Q is prime
492// TODO impl karatsuba for non-prime Q
493fn mul<const Q: u64, const N: usize>(lhs: &Rq<Q, N>, rhs: &Rq<Q, N>) -> Rq<Q, N> {
494 // reuse evaluations if already computed
495 let lhs_evals = if lhs.evals.is_some() {
496 lhs.evals.unwrap()
497 } else {
498 NTT::<Q, N>::ntt(lhs.coeffs)
499 };
500 let rhs_evals = if rhs.evals.is_some() {
501 rhs.evals.unwrap()
502 } else {
503 NTT::<Q, N>::ntt(rhs.coeffs)
504 };
505
506 let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
507 let c = NTT::<Q, { N }>::intt(c_ntt);
508 Rq::new(c, Some(c_ntt))
509}
510
511impl<const Q: u64, const N: usize> fmt::Display for Rq<Q, N> {
512 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
513 self.fmt(f)?;
514 Ok(())
515 }
516}
517impl<const Q: u64, const N: usize> fmt::Debug for Rq<Q, N> {
518 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
519 self.fmt(f)?;
520 Ok(())
521 }
522}
523
524#[cfg(test)]
525mod tests {
526 use super::*;
527
528 #[test]
529 fn poly_ring() {
530 // the test values used are generated with SageMath
531 const Q: u64 = 7;
532 const N: usize = 3;
533
534 // p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
535 let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
536 assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
537
538 // try with coefficients bigger than Q
539 let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
540 assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
541
542 // try with other ring
543 let p = Rq::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
544 assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
545
546 let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
547 assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
548
549 let p = Rq::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
550 assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
551
552 let a = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
553 assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
554
555 let b = Rq::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
556 assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
557
558 // add
559 assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
560 assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
561 // assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
562
563 // sub
564 assert_eq!(
565 (a.clone() - b.clone()).to_string(),
566 "x^2 + x + 1 mod Z_7/(X^3+1)"
567 );
568 }
569
570 #[test]
571 fn test_mul() -> Result<()> {
572 const Q: u64 = 2u64.pow(16) + 1;
573 const N: usize = 4;
574
575 let a: [u64; N] = [1u64, 2, 3, 4];
576 let b: [u64; N] = [1u64, 2, 3, 4];
577 let c: [u64; N] = [65513, 65517, 65531, 20];
578 test_mul_opt::<Q, N>(a, b, c)?;
579
580 let a: [u64; N] = [0u64, 0, 0, 2];
581 let b: [u64; N] = [0u64, 0, 0, 2];
582 let c: [u64; N] = [0u64, 0, 65533, 0];
583 test_mul_opt::<Q, N>(a, b, c)?;
584
585 // TODO more testvectors
586
587 Ok(())
588 }
589 fn test_mul_opt<const Q: u64, const N: usize>(
590 a: [u64; N],
591 b: [u64; N],
592 expected_c: [u64; N],
593 ) -> Result<()> {
594 let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
595 let mut a = Rq::new(a, None);
596 let b: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(b[i]));
597 let mut b = Rq::new(b, None);
598 let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(expected_c[i]));
599 let expected_c = Rq::new(expected_c, None);
600
601 let c = mul_mut(&mut a, &mut b);
602 assert_eq!(c, expected_c);
603 Ok(())
604 }
605}