1//! Polynomial ring Z[X]/(X^N+1)
  2//!
  3
  4use anyhow::{anyhow, Result};
  5use rand::{distributions::Distribution, Rng};
  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}
 62pub fn mul_div_round<const Q: u64, const N: usize>(
 63    v: Vec<i64>,
 64    num: u64,
 65    den: u64,
 66) -> crate::Rq<Q, N> {
 67    // dbg!(&v);
 68    let r: Vec<f64> = v
 69        .iter()
 70        .map(|e| ((num as f64 * *e as f64) / den as f64).round())
 71        .collect();
 72    // dbg!(&r);
 73    crate::Rq::<Q, N>::from_vec_f64(r)
 74}
 75
 76// apply mod (X^N+1)
 77pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
 78    if p.len() < N {
 79        return;
 80    }
 81    for i in N..p.len() {
 82        p[i - N] = p[i - N].clone() - p[i].clone();
 83        p[i] = 0;
 84    }
 85    p.truncate(N);
 86}
 87
 88impl<const N: usize> PartialEq for R<N> {
 89    fn eq(&self, other: &Self) -> bool {
 90        self.0 == other.0
 91    }
 92}
 93impl<const N: usize> ops::Add<R<N>> for R<N> {
 94    type Output = Self;
 95
 96    fn add(self, rhs: Self) -> Self {
 97        Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
 98    }
 99}
100impl<const N: usize> ops::Add<&R<N>> for &R<N> {
101    type Output = R<N>;
102
103    fn add(self, rhs: &R<N>) -> Self::Output {
104        R(array::from_fn(|i| self.0[i] + rhs.0[i]))
105    }
106}
107impl<const N: usize> ops::Sub<R<N>> for R<N> {
108    type Output = Self;
109
110    fn sub(self, rhs: Self) -> Self {
111        Self(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 = R<N>;
116
117    fn sub(self, rhs: &R<N>) -> Self::Output {
118        R(array::from_fn(|i| self.0[i] - rhs.0[i]))
119    }
120}
121impl<const N: usize> ops::Mul<R<N>> for R<N> {
122    type Output = Self;
123
124    fn mul(self, rhs: Self) -> Self {
125        naive_poly_mul(&self, &rhs)
126    }
127}
128impl<const N: usize> ops::Mul<&R<N>> for &R<N> {
129    type Output = R<N>;
130
131    fn mul(self, rhs: &R<N>) -> Self::Output {
132        naive_poly_mul(self, rhs)
133    }
134}
135
136// TODO with NTT(?)
137pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
138    let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
139    let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
140    let mut result: Vec<i128> = vec![0; (N * 2) - 1];
141    for i in 0..N {
142        for j in 0..N {
143            result[i + j] = result[i + j] + poly1[i] * poly2[j];
144        }
145    }
146
147    // apply mod (X^N + 1))
148    R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
149}
150pub fn naive_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
151    let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
152    let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
153    let mut result = vec![0; (N * 2) - 1];
154    for i in 0..N {
155        for j in 0..N {
156            result[i + j] = result[i + j] + poly1[i] * poly2[j];
157        }
158    }
159    result.iter().map(|c| *c as i64).collect()
160}
161
162// mul by u64
163impl<const N: usize> ops::Mul<u64> for R<N> {
164    type Output = Self;
165
166    fn mul(self, s: u64) -> Self {
167        self.mul_by_i64(s as i64)
168    }
169}
170impl<const N: usize> ops::Mul<&u64> for &R<N> {
171    type Output = R<N>;
172
173    fn mul(self, s: &u64) -> Self::Output {
174        self.mul_by_i64(*s as i64)
175    }
176}
177
178impl<const N: usize> ops::Neg for R<N> {
179    type Output = Self;
180
181    fn neg(self) -> Self::Output {
182        Self(array::from_fn(|i| -self.0[i]))
183    }
184}