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, Rq, Zq};
 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    const Q: u64 = u64::MAX; // WIP
 26    const N: usize = N;
 27
 28    fn coeffs(&self) -> Vec<T64> {
 29        self.0.to_vec()
 30    }
 31
 32    fn zero() -> Self {
 33        Self(array::from_fn(|_| T64::zero()))
 34    }
 35
 36    fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
 37        Self(array::from_fn(|_| T64::rand_f64(&mut rng, &dist)))
 38    }
 39
 40    fn from_vec(coeffs: Vec<Self::C>) -> Self {
 41        let mut p = coeffs;
 42        modulus::<N>(&mut p);
 43        Self(array::from_fn(|i| p[i]))
 44    }
 45
 46    fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 47        let elems: Vec<Vec<T64>> = self.0.iter().map(|r| r.decompose(beta, l)).collect();
 48        // transpose it
 49        let r: Vec<Vec<T64>> = (0..elems[0].len())
 50            .map(|i| (0..elems.len()).map(|j| elems[j][i]).collect())
 51            .collect();
 52        // convert it to Tn<N>
 53        r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
 54    }
 55
 56    fn remodule<const P: u64>(&self) -> Tn<N> {
 57        todo!()
 58        // Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
 59    }
 60
 61    // fn mod_switch<const P: u64>(&self) -> impl Ring {
 62    fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
 63        // unimplemented!()
 64        // TODO WIP
 65        let coeffs = array::from_fn(|i| Zq::<P>::from_u64(self.0[i].mod_switch::<P>().0));
 66        Rq::<P, N> {
 67            coeffs,
 68            evals: None,
 69        }
 70    }
 71
 72    /// returns [ [(num/den) * self].round() ] mod q
 73    /// ie. performs the multiplication and division over f64, and then it rounds the
 74    /// result, only applying the mod Q at the end
 75    fn mul_div_round(&self, num: u64, den: u64) -> Self {
 76        let r: Vec<T64> = self
 77            .coeffs()
 78            .iter()
 79            .map(|e| T64(((num as f64 * e.0 as f64) / den as f64).round() as u64))
 80            .collect();
 81        Self::from_vec(r)
 82    }
 83}
 84
 85// apply mod (X^N+1)
 86pub fn modulus<const N: usize>(p: &mut Vec<T64>) {
 87    if p.len() < N {
 88        return;
 89    }
 90    for i in N..p.len() {
 91        p[i - N] = p[i - N].clone() - p[i].clone();
 92        p[i] = T64::zero();
 93    }
 94    p.truncate(N);
 95}
 96
 97impl<const N: usize> Add<Tn<N>> for Tn<N> {
 98    type Output = Self;
 99
100    fn add(self, rhs: Self) -> Self {
101        Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
102    }
103}
104impl<const N: usize> Add<&Tn<N>> for &Tn<N> {
105    type Output = Tn<N>;
106
107    fn add(self, rhs: &Tn<N>) -> Self::Output {
108        Tn(array::from_fn(|i| self.0[i] + rhs.0[i]))
109    }
110}
111impl<const N: usize> AddAssign for Tn<N> {
112    fn add_assign(&mut self, rhs: Self) {
113        for i in 0..N {
114            self.0[i] += rhs.0[i];
115        }
116    }
117}
118
119impl<const N: usize> Sum<Tn<N>> for Tn<N> {
120    fn sum<I>(iter: I) -> Self
121    where
122        I: Iterator<Item = Self>,
123    {
124        let mut acc = Tn::<N>::zero();
125        for e in iter {
126            acc += e;
127        }
128        acc
129    }
130}
131
132impl<const N: usize> Sub<Tn<N>> for Tn<N> {
133    type Output = Self;
134
135    fn sub(self, rhs: Self) -> Self {
136        Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
137    }
138}
139impl<const N: usize> Sub<&Tn<N>> for &Tn<N> {
140    type Output = Tn<N>;
141
142    fn sub(self, rhs: &Tn<N>) -> Self::Output {
143        Tn(array::from_fn(|i| self.0[i] - rhs.0[i]))
144    }
145}
146
147impl<const N: usize> SubAssign for Tn<N> {
148    fn sub_assign(&mut self, rhs: Self) {
149        for i in 0..N {
150            self.0[i] -= rhs.0[i];
151        }
152    }
153}
154
155impl<const N: usize> PartialEq for Tn<N> {
156    fn eq(&self, other: &Self) -> bool {
157        self.0 == other.0
158    }
159}
160
161impl<const N: usize> Mul<Tn<N>> for Tn<N> {
162    type Output = Self;
163
164    fn mul(self, rhs: Self) -> Self {
165        naive_poly_mul(&self, &rhs)
166    }
167}
168impl<const N: usize> Mul<&Tn<N>> for &Tn<N> {
169    type Output = Tn<N>;
170
171    fn mul(self, rhs: &Tn<N>) -> Self::Output {
172        naive_poly_mul(self, rhs)
173    }
174}
175
176fn naive_poly_mul<const N: usize>(poly1: &Tn<N>, poly2: &Tn<N>) -> Tn<N> {
177    let poly1: Vec<u128> = poly1.0.iter().map(|c| c.0 as u128).collect();
178    let poly2: Vec<u128> = poly2.0.iter().map(|c| c.0 as u128).collect();
179    let mut result: Vec<u128> = vec![0; (N * 2) - 1];
180    for i in 0..N {
181        for j in 0..N {
182            result[i + j] = result[i + j] + poly1[i] * poly2[j];
183        }
184    }
185
186    // apply mod (X^N + 1))
187    modulus_u128::<N>(&mut result);
188
189    Tn(array::from_fn(|i| T64(result[i] as u64)))
190}
191fn modulus_u128<const N: usize>(p: &mut Vec<u128>) {
192    if p.len() < N {
193        return;
194    }
195    for i in N..p.len() {
196        // p[i - N] = p[i - N].clone() - p[i].clone();
197        p[i - N] = p[i - N].wrapping_sub(p[i]);
198        p[i] = 0;
199    }
200    p.truncate(N);
201}
202
203impl<const N: usize> Mul<T64> for Tn<N> {
204    type Output = Self;
205    fn mul(self, s: T64) -> Self {
206        Self(array::from_fn(|i| self.0[i] * s))
207    }
208}
209// mul by u64
210impl<const N: usize> Mul<u64> for Tn<N> {
211    type Output = Self;
212    fn mul(self, s: u64) -> Self {
213        Self(array::from_fn(|i| self.0[i] * s))
214    }
215}
216impl<const N: usize> Mul<&u64> for &Tn<N> {
217    type Output = Tn<N>;
218    fn mul(self, s: &u64) -> Self::Output {
219        Tn::<N>(array::from_fn(|i| self.0[i] * *s))
220    }
221}