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