1//! 𝕋_<N,q>[X] = ℝ_<N,q>[X] / ℤ_<N,q>[X], polynomials modulo X^N+1 with
  2//! coefficients in 𝕋_Q.
  3//!
  4//! Note: this is not an algebraic ring, since internal-product is not well
  5//! defined. But since we work over the discrete torus 𝕋_q, which we identify as
  6//! 𝕋q = ℤ/qℤ ≈ ℤq, whith q=64. Since we allow product between 𝕋q elements and
  7//! u64, we fit it into the `Ring` trait (from ring.rs) so that we can compose
  8//! the 𝕋_<N,q> implementation with the other objects from the code.
  9
 10use rand::{distributions::Distribution, Rng};
 11use std::array;
 12use std::iter::Sum;
 13use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
 14
 15use crate::{ring::Ring, torus::T64};
 16
 17/// 𝕋_<N,Q>[X] = 𝕋<Q>[X]/(X^N +1), polynomials modulo X^N+1 with coefficients in
 18/// 𝕋, where Q=2^64.
 19#[derive(Clone, Copy, Debug)]
 20pub struct Tn<const N: usize>(pub [T64; N]);
 21
 22impl<const N: usize> Ring for Tn<N> {
 23    type C = T64;
 24
 25    fn coeffs(&self) -> Vec<T64> {
 26        self.0.to_vec()
 27    }
 28
 29    fn zero() -> Self {
 30        Self(array::from_fn(|_| T64::zero()))
 31    }
 32
 33    fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
 34        Self(array::from_fn(|_| T64::rand_f64(&mut rng, &dist)))
 35    }
 36
 37    fn from_vec(coeffs: Vec<Self::C>) -> Self {
 38        let mut p = coeffs;
 39        modulus::<N>(&mut p);
 40        Self(array::from_fn(|i| p[i]))
 41    }
 42
 43    fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 44        let elems: Vec<Vec<T64>> = self.0.iter().map(|r| r.decompose(beta, l)).collect();
 45        // transpose it
 46        let r: Vec<Vec<T64>> = (0..elems[0].len())
 47            .map(|i| (0..elems.len()).map(|j| elems[j][i]).collect())
 48            .collect();
 49        // convert it to Tn<N>
 50        r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
 51    }
 52
 53    /// returns [ [(num/den) * self].round() ] mod q
 54    /// ie. performs the multiplication and division over f64, and then it rounds the
 55    /// result, only applying the mod Q at the end
 56    fn mul_div_round(&self, num: u64, den: u64) -> Self {
 57        let r: Vec<T64> = self
 58            .coeffs()
 59            .iter()
 60            .map(|e| T64(((num as f64 * e.0 as f64) / den as f64).round() as u64))
 61            .collect();
 62        Self::from_vec(r)
 63    }
 64}
 65
 66// apply mod (X^N+1)
 67pub fn modulus<const N: usize>(p: &mut Vec<T64>) {
 68    if p.len() < N {
 69        return;
 70    }
 71    for i in N..p.len() {
 72        p[i - N] = p[i - N].clone() - p[i].clone();
 73        p[i] = T64::zero();
 74    }
 75    p.truncate(N);
 76}
 77
 78impl<const N: usize> Add<Tn<N>> for Tn<N> {
 79    type Output = Self;
 80
 81    fn add(self, rhs: Self) -> Self {
 82        Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
 83    }
 84}
 85impl<const N: usize> Add<&Tn<N>> for &Tn<N> {
 86    type Output = Tn<N>;
 87
 88    fn add(self, rhs: &Tn<N>) -> Self::Output {
 89        Tn(array::from_fn(|i| self.0[i] + rhs.0[i]))
 90    }
 91}
 92impl<const N: usize> AddAssign for Tn<N> {
 93    fn add_assign(&mut self, rhs: Self) {
 94        for i in 0..N {
 95            self.0[i] += rhs.0[i];
 96        }
 97    }
 98}
 99
100impl<const N: usize> Sum<Tn<N>> for Tn<N> {
101    fn sum<I>(iter: I) -> Self
102    where
103        I: Iterator<Item = Self>,
104    {
105        let mut acc = Tn::<N>::zero();
106        for e in iter {
107            acc += e;
108        }
109        acc
110    }
111}
112
113impl<const N: usize> Sub<Tn<N>> for Tn<N> {
114    type Output = Self;
115
116    fn sub(self, rhs: Self) -> Self {
117        Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
118    }
119}
120impl<const N: usize> Sub<&Tn<N>> for &Tn<N> {
121    type Output = Tn<N>;
122
123    fn sub(self, rhs: &Tn<N>) -> Self::Output {
124        Tn(array::from_fn(|i| self.0[i] - rhs.0[i]))
125    }
126}
127
128impl<const N: usize> SubAssign for Tn<N> {
129    fn sub_assign(&mut self, rhs: Self) {
130        for i in 0..N {
131            self.0[i] -= rhs.0[i];
132        }
133    }
134}
135
136impl<const N: usize> PartialEq for Tn<N> {
137    fn eq(&self, other: &Self) -> bool {
138        self.0 == other.0
139    }
140}
141
142impl<const N: usize> Mul<Tn<N>> for Tn<N> {
143    type Output = Self;
144
145    fn mul(self, rhs: Self) -> Self {
146        naive_poly_mul(&self, &rhs)
147    }
148}
149impl<const N: usize> Mul<&Tn<N>> for &Tn<N> {
150    type Output = Tn<N>;
151
152    fn mul(self, rhs: &Tn<N>) -> Self::Output {
153        naive_poly_mul(self, rhs)
154    }
155}
156
157fn naive_poly_mul<const N: usize>(poly1: &Tn<N>, poly2: &Tn<N>) -> Tn<N> {
158    let poly1: Vec<u128> = poly1.0.iter().map(|c| c.0 as u128).collect();
159    let poly2: Vec<u128> = poly2.0.iter().map(|c| c.0 as u128).collect();
160    let mut result: Vec<u128> = vec![0; (N * 2) - 1];
161    for i in 0..N {
162        for j in 0..N {
163            result[i + j] = result[i + j] + poly1[i] * poly2[j];
164        }
165    }
166
167    // apply mod (X^N + 1))
168    modulus_u128::<N>(&mut result);
169
170    // sanity check: check that there are no coeffs > i64_max
171    assert_eq!(
172        result,
173        Tn::<N>(array::from_fn(|i| T64(result[i] as u64)))
174            .coeffs()
175            .iter()
176            .map(|c| c.0 as u128)
177            .collect::<Vec<_>>()
178    );
179    Tn(array::from_fn(|i| T64(result[i] as u64)))
180}
181fn modulus_u128<const N: usize>(p: &mut Vec<u128>) {
182    if p.len() < N {
183        return;
184    }
185    for i in N..p.len() {
186        p[i - N] = p[i - N].clone() - p[i].clone();
187        p[i] = 0;
188    }
189    p.truncate(N);
190}
191
192// mul by u64
193impl<const N: usize> Mul<u64> for Tn<N> {
194    type Output = Self;
195    fn mul(self, s: u64) -> Self {
196        Self(array::from_fn(|i| self.0[i] * s))
197    }
198}
199impl<const N: usize> Mul<&u64> for &Tn<N> {
200    type Output = Tn<N>;
201    fn mul(self, s: &u64) -> Self::Output {
202        Tn::<N>(array::from_fn(|i| self.0[i] * *s))
203    }
204}