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