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