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