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