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