fhe-study/arithmetic/src/ring.rs
arithmetic/src/ring.rsRaw
1use rand::{distributions::Distribution, Rng};
2use std::array;
3use std::fmt;
4use std::ops;
5
6use crate::zq::Zq;
7use anyhow::{anyhow, Result};
8
9// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
10#[derive(Clone, Copy)]
11pub struct PR<const Q: u64, const N: usize> {
12 pub(crate) coeffs: [Zq<Q>; N],
13
14 // evals are set when doig a PRxPR multiplication, so it can be reused in future
15 // multiplications avoiding recomputing it
16 pub(crate) evals: Option<[Zq<Q>; N]>,
17}
18
19// TODO define a trait "PolynomialRingTrait" or similar, so that when other structs use it can just
20// use the trait and not need to add '<Q, N>' to their params
21
22// apply mod (X^N+1)
23pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
24 if p.len() < N {
25 return;
26 }
27 for i in N..p.len() {
28 p[i - N] = p[i - N].clone() - p[i].clone();
29 p[i] = Zq(0);
30 }
31 p.truncate(N);
32}
33
34// PR stands for PolynomialRing
35impl<const Q: u64, const N: usize> PR<Q, N> {
36 pub fn coeffs(&self) -> [Zq<Q>; N] {
37 self.coeffs
38 }
39
40 pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
41 let mut p = coeffs;
42 modulus::<Q, N>(&mut p);
43 let coeffs = array::from_fn(|i| p[i]);
44 Self {
45 coeffs,
46 evals: None,
47 }
48 }
49 // this method is mostly for tests
50 pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
51 let coeffs_mod_q = coeffs.iter().map(|c| Zq::new(*c)).collect();
52 Self::from_vec(coeffs_mod_q)
53 }
54 pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
55 Self { coeffs, evals }
56 }
57
58 pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
59 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
60 Ok(Self {
61 coeffs,
62 evals: None,
63 })
64 }
65 pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
66 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
67 Ok(Self {
68 coeffs,
69 evals: None,
70 })
71 }
72 // WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
73 pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
74 let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
75 Ok(PR {
76 coeffs,
77 evals: None,
78 })
79 }
80
81 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
82 // TODO simplify
83 let mut str = "";
84 let mut zero = true;
85 for (i, coeff) in self.coeffs.iter().enumerate().rev() {
86 if coeff.0 == 0 {
87 continue;
88 }
89 zero = false;
90 f.write_str(str)?;
91 if coeff.0 != 1 {
92 f.write_str(coeff.0.to_string().as_str())?;
93 if i > 0 {
94 f.write_str("*")?;
95 }
96 }
97 if coeff.0 == 1 && i == 0 {
98 f.write_str(coeff.0.to_string().as_str())?;
99 }
100 if i == 1 {
101 f.write_str("x")?;
102 } else if i > 1 {
103 f.write_str("x^")?;
104 f.write_str(i.to_string().as_str())?;
105 }
106 str = " + ";
107 }
108 if zero {
109 f.write_str("0")?;
110 }
111
112 f.write_str(" mod Z_")?;
113 f.write_str(Q.to_string().as_str())?;
114 f.write_str("/(X^")?;
115 f.write_str(N.to_string().as_str())?;
116 f.write_str("+1)")?;
117 Ok(())
118 }
119}
120pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
121 // assert_eq!(m.len(), m[0].len()); // TODO change to returning err
122 // assert_eq!(m.len(), v.len());
123 if m.len() != m[0].len() {
124 return Err(anyhow!("expected 'm' to be a square matrix"));
125 }
126 if m.len() != v.len() {
127 return Err(anyhow!(
128 "m.len: {} should be equal to v.len(): {}",
129 m.len(),
130 v.len(),
131 ));
132 }
133
134 Ok(m.iter()
135 .map(|row| {
136 row.iter()
137 .zip(v.iter())
138 .map(|(&row_i, &v_i)| row_i * v_i)
139 .sum()
140 })
141 .collect::<Vec<Zq<Q>>>())
142}
143pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
144 // TODO case when m[0].len()=0
145 // TODO non square matrix
146 let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
147 for (i, m_row) in m.iter().enumerate() {
148 for (j, m_ij) in m_row.iter().enumerate() {
149 r[j][i] = *m_ij;
150 }
151 }
152 r
153}
154
155impl<const Q: u64, const N: usize> PartialEq for PR<Q, N> {
156 fn eq(&self, other: &Self) -> bool {
157 self.coeffs == other.coeffs
158 }
159}
160impl<const Q: u64, const N: usize> ops::Add<PR<Q, N>> for PR<Q, N> {
161 type Output = Self;
162
163 fn add(self, rhs: Self) -> Self {
164 Self {
165 coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
166 evals: None,
167 }
168 // Self {
169 // coeffs: self
170 // .coeffs
171 // .iter()
172 // .zip(rhs.coeffs)
173 // .map(|(a, b)| *a + b)
174 // .collect(),
175 // evals: None,
176 // }
177 // Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in +
178 }
179}
180impl<const Q: u64, const N: usize> ops::Add<&PR<Q, N>> for &PR<Q, N> {
181 type Output = PR<Q, N>;
182
183 fn add(self, rhs: &PR<Q, N>) -> Self::Output {
184 PR {
185 coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
186 evals: None,
187 }
188 }
189}
190impl<const Q: u64, const N: usize> ops::Sub<PR<Q, N>> for PR<Q, N> {
191 type Output = Self;
192
193 fn sub(self, rhs: Self) -> Self {
194 Self {
195 coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
196 evals: None,
197 }
198 }
199}
200impl<const Q: u64, const N: usize> ops::Sub<&PR<Q, N>> for &PR<Q, N> {
201 type Output = PR<Q, N>;
202
203 fn sub(self, rhs: &PR<Q, N>) -> Self::Output {
204 PR {
205 coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
206 evals: None,
207 }
208 }
209}
210
211impl<const Q: u64, const N: usize> ops::Neg for PR<Q, N> {
212 type Output = Self;
213
214 fn neg(self) -> Self::Output {
215 Self {
216 coeffs: array::from_fn(|i| -self.coeffs[i]),
217 evals: None,
218 }
219 }
220}
221
222impl<const Q: u64, const N: usize> fmt::Display for PR<Q, N> {
223 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
224 self.fmt(f)?;
225 Ok(())
226 }
227}
228impl<const Q: u64, const N: usize> fmt::Debug for PR<Q, N> {
229 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
230 self.fmt(f)?;
231 Ok(())
232 }
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238
239 #[test]
240 fn poly_ring() {
241 // the test values used are generated with SageMath
242 const Q: u64 = 7;
243 const N: usize = 3;
244
245 // p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
246 let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
247 assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
248
249 // try with coefficients bigger than Q
250 let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
251 assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
252
253 // try with other ring
254 let p = PR::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
255 assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
256
257 let p = PR::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
258 assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
259
260 let p = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
261 assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
262
263 let a = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
264 assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
265
266 let b = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
267 assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
268
269 // add
270 assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
271 assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
272 // assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
273
274 // sub
275 assert_eq!(
276 (a.clone() - b.clone()).to_string(),
277 "x^2 + x + 1 mod Z_7/(X^3+1)"
278 );
279 }
280}