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