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