fhe-study/arith/src/ring.rs
arith/src/ring.rsRaw
1//! Polynomial ring Z[X]/(X^N+1)
2//!
3
4use anyhow::{Result, anyhow};
5use rand::{Rng, distributions::Distribution};
6use std::array;
7use std::fmt;
8use std::ops;
9
10use crate::ntt::NTT;
11use crate::zq::Zq;
12
13// PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1)
14#[derive(Clone, Copy, Debug)]
15pub struct R<const N: usize>([i64; N]);
16
17impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> {
18 fn from(rq: crate::ringq::Rq<Q, N>) -> Self {
19 Self::from_vec_u64(rq.coeffs().to_vec().iter().map(|e| e.0).collect())
20 }
21}
22
23impl<const N: usize> R<N> {
24 pub fn coeffs(&self) -> [i64; N] {
25 self.0
26 }
27 pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
28 crate::Rq::<Q, N>::from(self)
29 }
30
31 pub fn from_vec(coeffs: Vec<i64>) -> Self {
32 let mut p = coeffs;
33 modulus::<N>(&mut p);
34 Self(array::from_fn(|i| p[i]))
35 }
36 // this method is mostly for tests
37 pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
38 let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
39 Self::from_vec(coeffs_i64)
40 }
41 pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
42 let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
43 Self::from_vec(coeffs_i64)
44 }
45 pub fn new(coeffs: [i64; N]) -> Self {
46 Self(coeffs)
47 }
48 pub fn mul_by_i64(&self, s: i64) -> Self {
49 Self(array::from_fn(|i| self.0[i] * s))
50 }
51 // performs the multiplication and division over f64, and then it rounds the
52 // result, only applying the mod Q at the end
53 pub fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> {
54 let r: Vec<f64> = self
55 .coeffs()
56 .iter()
57 .map(|e| ((num as f64 * *e as f64) / den as f64).round())
58 .collect();
59 crate::Rq::<Q, N>::from_vec_f64(r)
60 }
61}
62
63pub fn mul_div_round<const Q: u64, const N: usize>(
64 v: Vec<i64>,
65 num: u64,
66 den: u64,
67) -> crate::Rq<Q, N> {
68 // dbg!(&v);
69 let r: Vec<f64> = v
70 .iter()
71 .map(|e| ((num as f64 * *e as f64) / den as f64).round())
72 .collect();
73 // dbg!(&r);
74 crate::Rq::<Q, N>::from_vec_f64(r)
75}
76
77// TODO rename to make it clear that is not mod q, but mod X^N+1
78// apply mod (X^N+1)
79pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
80 if p.len() < N {
81 return;
82 }
83 for i in N..p.len() {
84 p[i - N] = p[i - N].clone() - p[i].clone();
85 p[i] = 0;
86 }
87 p.truncate(N);
88}
89pub fn modulus_i128<const N: usize>(p: &mut Vec<i128>) {
90 if p.len() < N {
91 return;
92 }
93 for i in N..p.len() {
94 p[i - N] = p[i - N].clone() - p[i].clone();
95 p[i] = 0;
96 }
97 p.truncate(N);
98}
99
100impl<const N: usize> PartialEq for R<N> {
101 fn eq(&self, other: &Self) -> bool {
102 self.0 == other.0
103 }
104}
105impl<const N: usize> ops::Add<R<N>> for R<N> {
106 type Output = Self;
107
108 fn add(self, rhs: Self) -> Self {
109 Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
110 }
111}
112impl<const N: usize> ops::Add<&R<N>> for &R<N> {
113 type Output = R<N>;
114
115 fn add(self, rhs: &R<N>) -> Self::Output {
116 R(array::from_fn(|i| self.0[i] + rhs.0[i]))
117 }
118}
119impl<const N: usize> ops::Sub<R<N>> for R<N> {
120 type Output = Self;
121
122 fn sub(self, rhs: Self) -> Self {
123 Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
124 }
125}
126impl<const N: usize> ops::Sub<&R<N>> for &R<N> {
127 type Output = R<N>;
128
129 fn sub(self, rhs: &R<N>) -> Self::Output {
130 R(array::from_fn(|i| self.0[i] - rhs.0[i]))
131 }
132}
133impl<const N: usize> ops::Mul<R<N>> for R<N> {
134 type Output = Self;
135
136 fn mul(self, rhs: Self) -> Self {
137 naive_poly_mul(&self, &rhs)
138 }
139}
140impl<const N: usize> ops::Mul<&R<N>> for &R<N> {
141 type Output = R<N>;
142
143 fn mul(self, rhs: &R<N>) -> Self::Output {
144 naive_poly_mul(self, rhs)
145 }
146}
147
148// TODO WIP
149pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
150 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
151 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
152 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
153 for i in 0..N {
154 for j in 0..N {
155 result[i + j] = result[i + j] + poly1[i] * poly2[j];
156 }
157 }
158
159 // apply mod (X^N + 1))
160 R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
161}
162
163pub fn naive_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
164 let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
165 let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
166 let mut result: Vec<i128> = vec![0; (N * 2) - 1];
167 for i in 0..N {
168 for j in 0..N {
169 result[i + j] = result[i + j] + poly1[i] * poly2[j];
170 }
171 }
172
173 modulus_i128::<N>(&mut result);
174 // for c_i in result.iter() {
175 // println!("---");
176 // println!("{:?}", &c_i);
177 // println!("{:?}", *c_i as i64);
178 // println!("{:?}", (*c_i as i64) as i128);
179 // assert_eq!(*c_i, (*c_i as i64) as i128, "{:?}", c_i);
180 // }
181
182 // let q: i128 = 65537;
183 // let result: Vec<i64> = result
184 // .iter()
185 // // .map(|c_i| ((c_i % q + q) % q) as i64)
186 // .map(|c_i| (c_i % q) as i64)
187 // // .map(|c_i| *c_i as i64)
188 // .collect();
189 // result
190
191 result.iter().map(|c| *c as i64).collect()
192}
193
194// mul by u64
195impl<const N: usize> ops::Mul<u64> for R<N> {
196 type Output = Self;
197
198 fn mul(self, s: u64) -> Self {
199 self.mul_by_i64(s as i64)
200 }
201}
202impl<const N: usize> ops::Mul<&u64> for &R<N> {
203 type Output = R<N>;
204
205 fn mul(self, s: &u64) -> Self::Output {
206 self.mul_by_i64(*s as i64)
207 }
208}
209
210impl<const N: usize> ops::Neg for R<N> {
211 type Output = Self;
212
213 fn neg(self) -> Self::Output {
214 Self(array::from_fn(|i| -self.0[i]))
215 }
216}