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