fhe-study/arithmetic/src/ring.rs
arithmetic/src/ring.rsRaw
1use rand::{distributions::Distribution, Rng};
2use std::array;
3use std::fmt;
4use std::ops;
5
6use crate::ntt::NTT;
7use crate::zq::Zq;
8use anyhow::{anyhow, Result};
9
10// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
11#[derive(Clone, Copy)]
12pub struct PR<const Q: u64, const N: usize> {
13 pub(crate) coeffs: [Zq<Q>; N],
14
15 // evals are set when doig a PRxPR multiplication, so it can be reused in future
16 // multiplications avoiding recomputing it
17 pub(crate) evals: Option<[Zq<Q>; N]>,
18}
19
20// TODO define a trait "PolynomialRingTrait" or similar, so that when other structs use it can just
21// use the trait and not need to add '<Q, N>' to their params
22
23// apply mod (X^N+1)
24pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
25 if p.len() < N {
26 return;
27 }
28 for i in N..p.len() {
29 p[i - N] = p[i - N].clone() - p[i].clone();
30 p[i] = Zq(0);
31 }
32 p.truncate(N);
33}
34
35// PR stands for PolynomialRing
36impl<const Q: u64, const N: usize> PR<Q, N> {
37 pub fn coeffs(&self) -> [Zq<Q>; N] {
38 self.coeffs
39 }
40
41 pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
42 let mut p = coeffs;
43 modulus::<Q, N>(&mut p);
44 let coeffs = array::from_fn(|i| p[i]);
45 Self {
46 coeffs,
47 evals: None,
48 }
49 }
50 // this method is mostly for tests
51 pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
52 let coeffs_mod_q = coeffs.iter().map(|c| Zq::new(*c)).collect();
53 Self::from_vec(coeffs_mod_q)
54 }
55 pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
56 Self { coeffs, evals }
57 }
58
59 pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
60 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
61 Ok(Self {
62 coeffs,
63 evals: None,
64 })
65 }
66 pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
67 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
68 Ok(Self {
69 coeffs,
70 evals: None,
71 })
72 }
73 // WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
74 pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
75 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
76 Ok(PR {
77 coeffs,
78 evals: None,
79 })
80 }
81
82 // TODO review if needed, or if with this interface
83 pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq<Q>>>) -> Result<Vec<Zq<Q>>> {
84 matrix_vec_product(m, &self.coeffs.to_vec())
85 }
86 pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
87 Self {
88 coeffs: array::from_fn(|i| self.coeffs[i] * *s),
89 evals: None,
90 }
91 }
92 pub fn mul_by_u64(&self, s: u64) -> Self {
93 let s = Zq::new(s);
94 Self {
95 coeffs: array::from_fn(|i| self.coeffs[i] * s),
96 // coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
97 evals: None,
98 }
99 }
100 pub fn mul_by_f64(&self, s: f64) -> Self {
101 Self {
102 coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
103 evals: None,
104 }
105 }
106
107 pub fn mul(&mut self, rhs: &mut Self) -> Self {
108 mul_mut(self, rhs)
109 }
110
111 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
112 // TODO simplify
113 let mut str = "";
114 let mut zero = true;
115 for (i, coeff) in self.coeffs.iter().enumerate().rev() {
116 if coeff.0 == 0 {
117 continue;
118 }
119 zero = false;
120 f.write_str(str)?;
121 if coeff.0 != 1 {
122 f.write_str(coeff.0.to_string().as_str())?;
123 if i > 0 {
124 f.write_str("*")?;
125 }
126 }
127 if coeff.0 == 1 && i == 0 {
128 f.write_str(coeff.0.to_string().as_str())?;
129 }
130 if i == 1 {
131 f.write_str("x")?;
132 } else if i > 1 {
133 f.write_str("x^")?;
134 f.write_str(i.to_string().as_str())?;
135 }
136 str = " + ";
137 }
138 if zero {
139 f.write_str("0")?;
140 }
141
142 f.write_str(" mod Z_")?;
143 f.write_str(Q.to_string().as_str())?;
144 f.write_str("/(X^")?;
145 f.write_str(N.to_string().as_str())?;
146 f.write_str("+1)")?;
147 Ok(())
148 }
149}
150pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
151 // assert_eq!(m.len(), m[0].len()); // TODO change to returning err
152 // assert_eq!(m.len(), v.len());
153 if m.len() != m[0].len() {
154 return Err(anyhow!("expected 'm' to be a square matrix"));
155 }
156 if m.len() != v.len() {
157 return Err(anyhow!(
158 "m.len: {} should be equal to v.len(): {}",
159 m.len(),
160 v.len(),
161 ));
162 }
163
164 Ok(m.iter()
165 .map(|row| {
166 row.iter()
167 .zip(v.iter())
168 .map(|(&row_i, &v_i)| row_i * v_i)
169 .sum()
170 })
171 .collect::<Vec<Zq<Q>>>())
172}
173pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
174 // TODO case when m[0].len()=0
175 // TODO non square matrix
176 let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
177 for (i, m_row) in m.iter().enumerate() {
178 for (j, m_ij) in m_row.iter().enumerate() {
179 r[j][i] = *m_ij;
180 }
181 }
182 r
183}
184
185impl<const Q: u64, const N: usize> PartialEq for PR<Q, N> {
186 fn eq(&self, other: &Self) -> bool {
187 self.coeffs == other.coeffs
188 }
189}
190impl<const Q: u64, const N: usize> ops::Add<PR<Q, N>> for PR<Q, N> {
191 type Output = Self;
192
193 fn add(self, rhs: Self) -> Self {
194 Self {
195 coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
196 evals: None,
197 }
198 // Self {
199 // coeffs: self
200 // .coeffs
201 // .iter()
202 // .zip(rhs.coeffs)
203 // .map(|(a, b)| *a + b)
204 // .collect(),
205 // evals: None,
206 // }
207 // Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in +
208 }
209}
210impl<const Q: u64, const N: usize> ops::Add<&PR<Q, N>> for &PR<Q, N> {
211 type Output = PR<Q, N>;
212
213 fn add(self, rhs: &PR<Q, N>) -> Self::Output {
214 PR {
215 coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
216 evals: None,
217 }
218 }
219}
220impl<const Q: u64, const N: usize> ops::Sub<PR<Q, N>> for PR<Q, N> {
221 type Output = Self;
222
223 fn sub(self, rhs: Self) -> Self {
224 Self {
225 coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
226 evals: None,
227 }
228 }
229}
230impl<const Q: u64, const N: usize> ops::Sub<&PR<Q, N>> for &PR<Q, N> {
231 type Output = PR<Q, N>;
232
233 fn sub(self, rhs: &PR<Q, N>) -> Self::Output {
234 PR {
235 coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
236 evals: None,
237 }
238 }
239}
240impl<const Q: u64, const N: usize> ops::Mul<PR<Q, N>> for PR<Q, N> {
241 type Output = Self;
242
243 fn mul(self, rhs: Self) -> Self {
244 mul(&self, &rhs)
245 }
246}
247impl<const Q: u64, const N: usize> ops::Mul<&PR<Q, N>> for &PR<Q, N> {
248 type Output = PR<Q, N>;
249
250 fn mul(self, rhs: &PR<Q, N>) -> Self::Output {
251 mul(self, rhs)
252 }
253}
254
255// mul by Zq element
256impl<const Q: u64, const N: usize> ops::Mul<Zq<Q>> for PR<Q, N> {
257 type Output = Self;
258
259 fn mul(self, s: Zq<Q>) -> Self {
260 self.mul_by_zq(&s)
261 }
262}
263impl<const Q: u64, const N: usize> ops::Mul<&Zq<Q>> for &PR<Q, N> {
264 type Output = PR<Q, N>;
265
266 fn mul(self, s: &Zq<Q>) -> Self::Output {
267 self.mul_by_zq(s)
268 }
269}
270// mul by u64
271impl<const Q: u64, const N: usize> ops::Mul<u64> for PR<Q, N> {
272 type Output = Self;
273
274 fn mul(self, s: u64) -> Self {
275 self.mul_by_u64(s)
276 }
277}
278impl<const Q: u64, const N: usize> ops::Mul<&u64> for &PR<Q, N> {
279 type Output = PR<Q, N>;
280
281 fn mul(self, s: &u64) -> Self::Output {
282 self.mul_by_u64(*s)
283 }
284}
285
286impl<const Q: u64, const N: usize> ops::Neg for PR<Q, N> {
287 type Output = Self;
288
289 fn neg(self) -> Self::Output {
290 Self {
291 coeffs: array::from_fn(|i| -self.coeffs[i]),
292 evals: None,
293 }
294 }
295}
296
297fn mul_mut<const Q: u64, const N: usize>(lhs: &mut PR<Q, N>, rhs: &mut PR<Q, N>) -> PR<Q, N> {
298 // reuse evaluations if already computed
299 if !lhs.evals.is_some() {
300 lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
301 };
302 if !rhs.evals.is_some() {
303 rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
304 };
305 let lhs_evals = lhs.evals.unwrap();
306 let rhs_evals = rhs.evals.unwrap();
307
308 let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
309 let c = NTT::<Q, { N }>::intt(c_ntt);
310 PR::new(c, Some(c_ntt))
311}
312fn mul<const Q: u64, const N: usize>(lhs: &PR<Q, N>, rhs: &PR<Q, N>) -> PR<Q, N> {
313 // reuse evaluations if already computed
314 let lhs_evals = if lhs.evals.is_some() {
315 lhs.evals.unwrap()
316 } else {
317 NTT::<Q, N>::ntt(lhs.coeffs)
318 };
319 let rhs_evals = if rhs.evals.is_some() {
320 rhs.evals.unwrap()
321 } else {
322 NTT::<Q, N>::ntt(rhs.coeffs)
323 };
324
325 let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
326 let c = NTT::<Q, { N }>::intt(c_ntt);
327 PR::new(c, Some(c_ntt))
328}
329
330impl<const Q: u64, const N: usize> fmt::Display for PR<Q, N> {
331 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
332 self.fmt(f)?;
333 Ok(())
334 }
335}
336impl<const Q: u64, const N: usize> fmt::Debug for PR<Q, N> {
337 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
338 self.fmt(f)?;
339 Ok(())
340 }
341}
342
343#[cfg(test)]
344mod tests {
345 use super::*;
346
347 #[test]
348 fn poly_ring() {
349 // the test values used are generated with SageMath
350 const Q: u64 = 7;
351 const N: usize = 3;
352
353 // p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
354 let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
355 assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
356
357 // try with coefficients bigger than Q
358 let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
359 assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
360
361 // try with other ring
362 let p = PR::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
363 assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
364
365 let p = PR::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
366 assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
367
368 let p = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
369 assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
370
371 let a = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
372 assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
373
374 let b = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
375 assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
376
377 // add
378 assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
379 assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
380 // assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
381
382 // sub
383 assert_eq!(
384 (a.clone() - b.clone()).to_string(),
385 "x^2 + x + 1 mod Z_7/(X^3+1)"
386 );
387 }
388
389 fn test_mul_opt<const Q: u64, const N: usize>(
390 a: [u64; N],
391 b: [u64; N],
392 expected_c: [u64; N],
393 ) -> Result<()> {
394 let a: [Zq<Q>; N] = array::from_fn(|i| Zq::new(a[i]));
395 let mut a = PR::new(a, None);
396 let b: [Zq<Q>; N] = array::from_fn(|i| Zq::new(b[i]));
397 let mut b = PR::new(b, None);
398 let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::new(expected_c[i]));
399 let expected_c = PR::new(expected_c, None);
400
401 let c = mul_mut(&mut a, &mut b);
402 assert_eq!(c, expected_c);
403 Ok(())
404 }
405 #[test]
406 fn test_mul() -> Result<()> {
407 const Q: u64 = 2u64.pow(16) + 1;
408 const N: usize = 4;
409
410 let a: [u64; N] = [1u64, 2, 3, 4];
411 let b: [u64; N] = [1u64, 2, 3, 4];
412 let c: [u64; N] = [65513, 65517, 65531, 20];
413 test_mul_opt::<Q, N>(a, b, c)?;
414
415 let a: [u64; N] = [0u64, 0, 0, 2];
416 let b: [u64; N] = [0u64, 0, 0, 2];
417 let c: [u64; N] = [0u64, 0, 65533, 0];
418 test_mul_opt::<Q, N>(a, b, c)?;
419
420 // TODO more testvectors
421
422 Ok(())
423 }
424}