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