fhe-study/arith/src/ring_n.rs
arith/src/ring_n.rsRaw
1//! Polynomial ring Z[X]/(X^N+1)
2//!
3
4use anyhow::Result;
5use rand::{distributions::Distribution, Rng};
6use std::array;
7use std::fmt;
8use std::iter::Sum;
9use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
10
11use crate::Ring;
12
13// TODO rename to not have name conflicts with the Ring trait (R: Ring)
14// PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1)
15#[derive(Clone, Copy)]
16pub struct R<const N: usize>(pub [i64; N]);
17
18impl<const N: usize> Ring for R<N> {
19 type C = i64;
20 fn coeffs(&self) -> Vec<Self::C> {
21 self.0.to_vec()
22 }
23 fn zero() -> Self {
24 let coeffs: [i64; N] = array::from_fn(|_| 0i64);
25 Self(coeffs)
26 }
27 fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
28 // let coeffs: [i64; N] = array::from_fn(|_| Self::C::rand(&mut rng, &dist));
29 let coeffs: [i64; N] = array::from_fn(|_| dist.sample(&mut rng).round() as i64);
30 Self(coeffs)
31 // let coeffs: [C; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
32 // Self(coeffs)
33 }
34
35 fn from_vec(coeffs: Vec<Self::C>) -> Self {
36 let mut p = coeffs;
37 modulus::<N>(&mut p);
38 Self(array::from_fn(|i| p[i]))
39 }
40
41 // returns the decomposition of each polynomial coefficient
42 fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
43 unimplemented!();
44 // array::from_fn(|i| self.coeffs[i].decompose(beta, l))
45 }
46
47 // performs the multiplication and division over f64, and then it rounds the
48 // result, only applying the mod Q at the end
49 fn mul_div_round(&self, num: u64, den: u64) -> Self {
50 unimplemented!()
51 // fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> {
52 // let r: Vec<f64> = self
53 // .coeffs()
54 // .iter()
55 // .map(|e| ((num as f64 * *e as f64) / den as f64).round())
56 // .collect();
57 // crate::Rq::<Q, N>::from_vec_f64(r)
58 }
59}
60
61impl<const Q: u64, const N: usize> From<crate::ring_nq::Rq<Q, N>> for R<N> {
62 fn from(rq: crate::ring_nq::Rq<Q, N>) -> Self {
63 Self::from_vec_u64(rq.coeffs().to_vec().iter().map(|e| e.0).collect())
64 }
65}
66
67impl<const N: usize> R<N> {
68 pub fn coeffs(&self) -> [i64; N] {
69 self.0
70 }
71 pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
72 crate::Rq::<Q, N>::from(self)
73 }
74
75 // this method is mostly for tests
76 pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
77 let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
78 Self::from_vec(coeffs_i64)
79 }
80 pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
81 let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
82 Self::from_vec(coeffs_i64)
83 }
84 pub fn new(coeffs: [i64; N]) -> Self {
85 Self(coeffs)
86 }
87 pub fn mul_by_i64(&self, s: i64) -> Self {
88 Self(array::from_fn(|i| self.0[i] * s))
89 }
90
91 pub fn infinity_norm(&self) -> u64 {
92 self.coeffs()
93 .iter()
94 // .map(|x| if x.0 > (Q / 2) { Q - x.0 } else { x.0 })
95 .map(|x| x.abs() as u64)
96 .fold(0, |a, b| a.max(b))
97 }
98 pub fn mod_centered_q<const Q: u64>(&self) -> R<N> {
99 let q = Q as i64;
100 let r = self
101 .0
102 .iter()
103 .map(|v| {
104 let mut res = v % q;
105 if res > q / 2 {
106 res = res - q;
107 }
108 res
109 })
110 .collect::<Vec<i64>>();
111 R::<N>::from_vec(r)
112 }
113}
114
115pub fn mul_div_round<const Q: u64, const N: usize>(
116 v: Vec<i64>,
117 num: u64,
118 den: u64,
119) -> crate::Rq<Q, N> {
120 // dbg!(&v);
121 let r: Vec<f64> = v
122 .iter()
123 .map(|e| ((num as f64 * *e as f64) / den as f64).round())
124 .collect();
125 // dbg!(&r);
126 crate::Rq::<Q, N>::from_vec_f64(r)
127}
128
129// TODO rename to make it clear that is not mod q, but mod X^N+1
130// apply mod (X^N+1)
131pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
132 if p.len() < N {
133 return;
134 }
135 for i in N..p.len() {
136 p[i - N] = p[i - N].clone() - p[i].clone();
137 p[i] = 0;
138 }
139 p.truncate(N);
140}
141pub fn modulus_i128<const N: usize>(p: &mut Vec<i128>) {
142 if p.len() < N {
143 return;
144 }
145 for i in N..p.len() {
146 p[i - N] = p[i - N].clone() - p[i].clone();
147 p[i] = 0;
148 }
149 p.truncate(N);
150}
151
152impl<const N: usize> PartialEq for R<N> {
153 fn eq(&self, other: &Self) -> bool {
154 self.0 == other.0
155 }
156}
157impl<const N: usize> Add<R<N>> for R<N> {
158 type Output = Self;
159
160 fn add(self, rhs: Self) -> Self {
161 Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
162 }
163}
164impl<const N: usize> Add<&R<N>> for &R<N> {
165 type Output = R<N>;
166
167 fn add(self, rhs: &R<N>) -> Self::Output {
168 R(array::from_fn(|i| self.0[i] + rhs.0[i]))
169 }
170}
171impl<const N: usize> AddAssign for R<N> {
172 fn add_assign(&mut self, rhs: Self) {
173 for i in 0..N {
174 self.0[i] += rhs.0[i];
175 }
176 }
177}
178
179impl<const N: usize> Sum<R<N>> for R<N> {
180 fn sum<I>(iter: I) -> Self
181 where
182 I: Iterator<Item = Self>,
183 {
184 let mut acc = R::<N>::zero();
185 for e in iter {
186 acc += e;
187 }
188 acc
189 }
190}
191
192impl<const N: usize> Sub<R<N>> for R<N> {
193 type Output = Self;
194
195 fn sub(self, rhs: Self) -> Self {
196 Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
197 }
198}
199impl<const N: usize> Sub<&R<N>> for &R<N> {
200 type Output = R<N>;
201
202 fn sub(self, rhs: &R<N>) -> Self::Output {
203 R(array::from_fn(|i| self.0[i] - rhs.0[i]))
204 }
205}
206
207impl<const N: usize> SubAssign for R<N> {
208 fn sub_assign(&mut self, rhs: Self) {
209 for i in 0..N {
210 self.0[i] -= rhs.0[i];
211 }
212 }
213}
214
215impl<const N: usize> Mul<R<N>> for R<N> {
216 type Output = Self;
217
218 fn mul(self, rhs: Self) -> Self {
219 naive_poly_mul(&self, &rhs)
220 }
221}
222impl<const N: usize> Mul<&R<N>> for &R<N> {
223 type Output = R<N>;
224
225 fn mul(self, rhs: &R<N>) -> Self::Output {
226 naive_poly_mul(self, rhs)
227 }
228}
229
230// TODO WIP
231pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
232 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
233 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
234 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
235 for i in 0..N {
236 for j in 0..N {
237 result[i + j] = result[i + j] + poly1[i] * poly2[j];
238 }
239 }
240
241 // apply mod (X^N + 1))
242 // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
243 modulus_i128::<N>(&mut result);
244 // dbg!(&result);
245 // dbg!(R::<N>(array::from_fn(|i| result[i] as i64)).coeffs());
246
247 // sanity check: check that there are no coeffs > i64_max
248 assert_eq!(
249 result,
250 R::<N>(array::from_fn(|i| result[i] as i64))
251 .coeffs()
252 .iter()
253 .map(|c| *c as i128)
254 .collect::<Vec<_>>()
255 );
256 R(array::from_fn(|i| result[i] as i64))
257}
258pub fn naive_mul_2<const N: usize>(poly1: &Vec<i128>, poly2: &Vec<i128>) -> Vec<i128> {
259 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
260 for i in 0..N {
261 for j in 0..N {
262 result[i + j] = result[i + j] + poly1[i] * poly2[j];
263 }
264 }
265
266 // apply mod (X^N + 1))
267 // R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
268 modulus_i128::<N>(&mut result);
269 result
270}
271
272pub fn naive_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
273 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
274 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
275 let mut result = vec![0; (N * 2) - 1];
276 for i in 0..N {
277 for j in 0..N {
278 result[i + j] = result[i + j] + poly1[i] * poly2[j];
279 }
280 }
281 result.iter().map(|c| *c as i64).collect()
282}
283pub fn naive_mul_TMP<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
284 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
285 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
286 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
287 for i in 0..N {
288 for j in 0..N {
289 result[i + j] = result[i + j] + poly1[i] * poly2[j];
290 }
291 }
292
293 // dbg!(&result);
294 modulus_i128::<N>(&mut result);
295 // for c_i in result.iter() {
296 // println!("---");
297 // println!("{:?}", &c_i);
298 // println!("{:?}", *c_i as i64);
299 // println!("{:?}", (*c_i as i64) as i128);
300 // assert_eq!(*c_i, (*c_i as i64) as i128, "{:?}", c_i);
301 // }
302 result.iter().map(|c| *c as i64).collect()
303}
304
305// wip
306pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
307 let q: i128 = Q as i128;
308 let r = p
309 .iter()
310 .map(|v| {
311 let mut res = v % q;
312 if res > q / 2 {
313 res = res - q;
314 }
315 res
316 })
317 .collect::<Vec<i128>>();
318 R::<N>::from_vec(r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
319}
320
321// mul by u64
322impl<const N: usize> Mul<u64> for R<N> {
323 type Output = Self;
324
325 fn mul(self, s: u64) -> Self {
326 self.mul_by_i64(s as i64)
327 }
328}
329impl<const N: usize> Mul<&u64> for &R<N> {
330 type Output = R<N>;
331
332 fn mul(self, s: &u64) -> Self::Output {
333 self.mul_by_i64(*s as i64)
334 }
335}
336
337impl<const N: usize> Neg for R<N> {
338 type Output = Self;
339
340 fn neg(self) -> Self::Output {
341 Self(array::from_fn(|i| -self.0[i]))
342 }
343}
344
345impl<const N: usize> R<N> {
346 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
347 let mut str = "";
348 let mut zero = true;
349 for (i, coeff) in self.0.iter().enumerate().rev() {
350 if *coeff == 0 {
351 continue;
352 }
353 zero = false;
354 f.write_str(str)?;
355 if *coeff != 1 {
356 f.write_str(coeff.to_string().as_str())?;
357 if i > 0 {
358 f.write_str("*")?;
359 }
360 }
361 if *coeff == 1 && i == 0 {
362 f.write_str(coeff.to_string().as_str())?;
363 }
364 if i == 1 {
365 f.write_str("x")?;
366 } else if i > 1 {
367 f.write_str("x^")?;
368 f.write_str(i.to_string().as_str())?;
369 }
370 str = " + ";
371 }
372 if zero {
373 f.write_str("0")?;
374 }
375
376 f.write_str(" mod Z")?;
377 f.write_str("/(X^")?;
378 f.write_str(N.to_string().as_str())?;
379 f.write_str("+1)")?;
380 Ok(())
381 }
382}
383impl<const N: usize> fmt::Display for R<N> {
384 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
385 self.fmt(f)?;
386 Ok(())
387 }
388}
389impl<const N: usize> fmt::Debug for R<N> {
390 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
391 self.fmt(f)?;
392 Ok(())
393 }
394}
395
396#[cfg(test)]
397mod tests {
398 use super::*;
399 use anyhow::Result;
400
401 #[test]
402 fn test_mul() -> Result<()> {
403 const Q: u64 = 2u64.pow(16) + 1;
404 const N: usize = 2;
405 let q: i64 = Q as i64;
406
407 // test vectors generated with SageMath
408 let a: [i64; N] = [q - 1, q - 1];
409 let b: [i64; N] = [q - 1, q - 1];
410 let c: [i64; N] = [0, 8589934592];
411 test_mul_opt::<Q, N>(a, b, c)?;
412
413 let a: [i64; N] = [1, q - 1];
414 let b: [i64; N] = [1, q - 1];
415 let c: [i64; N] = [-4294967295, 131072];
416 test_mul_opt::<Q, N>(a, b, c)?;
417
418 Ok(())
419 }
420 fn test_mul_opt<const Q: u64, const N: usize>(
421 a: [i64; N],
422 b: [i64; N],
423 expected_c: [i64; N],
424 ) -> Result<()> {
425 let mut a = R::new(a);
426 let mut b = R::new(b);
427 dbg!(&a);
428 dbg!(&b);
429 let expected_c = R::new(expected_c);
430
431 let mut c = naive_mul(&mut a, &mut b);
432 modulus::<N>(&mut c);
433 dbg!(R::<N>::from_vec(c.clone()));
434 assert_eq!(c, expected_c.0.to_vec());
435 Ok(())
436 }
437}