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 itertools::zip_eq;
 11use rand::{distributions::Distribution, Rng};
 12use std::array;
 13use std::iter::Sum;
 14use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
 15
 16use crate::{ring::Ring, torus::T64, Rq, Zq};
 17
 18/// 𝕋_<N,Q>[X] = 𝕋<Q>[X]/(X^N +1), polynomials modulo X^N+1 with coefficients in
 19/// 𝕋, where Q=2^64.
 20#[derive(Clone, Debug)]
 21pub struct Tn {
 22    pub n: usize,
 23    pub coeffs: Vec<T64>,
 24}
 25
 26impl Ring for Tn {
 27    type C = T64;
 28    type Param = usize; // n
 29
 30    // const Q: u64 = u64::MAX; // WIP
 31    // const N: usize = N;
 32
 33    fn param(&self) -> Self::Param {
 34        self.n
 35    }
 36    fn coeffs(&self) -> Vec<T64> {
 37        self.coeffs.to_vec()
 38    }
 39
 40    fn zero(n: usize) -> Self {
 41        Self {
 42            n,
 43            coeffs: vec![T64::zero(()); n],
 44        }
 45    }
 46
 47    fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, n: usize) -> Self {
 48        Self {
 49            n,
 50            coeffs: std::iter::repeat_with(|| T64::rand(&mut rng, &dist, ()))
 51                .take(n)
 52                .collect(),
 53        }
 54        // Self(array::from_fn(|_| T64::rand(&mut rng, &dist)))
 55    }
 56
 57    fn from_vec(n: usize, coeffs: Vec<Self::C>) -> Self {
 58        let mut p = coeffs;
 59        modulus(n, &mut p);
 60        Self { n, coeffs: p }
 61    }
 62
 63    fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 64        let elems: Vec<Vec<T64>> = self.coeffs.iter().map(|r| r.decompose(beta, l)).collect();
 65        // transpose it
 66        let r: Vec<Vec<T64>> = (0..elems[0].len())
 67            .map(|i| (0..elems.len()).map(|j| elems[j][i]).collect())
 68            .collect();
 69        // convert it to Tn
 70        r.iter()
 71            .map(|a_i| Self::from_vec(self.n, a_i.clone()))
 72            .collect()
 73    }
 74
 75    fn remodule(&self, p: u64) -> Tn {
 76        todo!()
 77        // Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
 78    }
 79
 80    // fn mod_switch<const P: u64>(&self) -> impl Ring {
 81    fn mod_switch(&self, p: u64) -> Rq {
 82        // unimplemented!()
 83        // TODO WIP
 84        let coeffs = self
 85            .coeffs
 86            .iter()
 87            .map(|c_i| Zq::from_u64(p, c_i.mod_switch(p).0))
 88            .collect();
 89        Rq {
 90            q: p,
 91            n: self.n,
 92            coeffs,
 93            evals: None,
 94        }
 95    }
 96
 97    /// returns [ [(num/den) * self].round() ] mod q
 98    /// ie. performs the multiplication and division over f64, and then it rounds the
 99    /// result, only applying the mod Q at the end
100    fn mul_div_round(&self, num: u64, den: u64) -> Self {
101        let r: Vec<T64> = self
102            .coeffs()
103            .iter()
104            .map(|e| T64(((num as f64 * e.0 as f64) / den as f64).round() as u64))
105            .collect();
106        Self::from_vec(self.n, r)
107    }
108}
109
110impl Tn {
111    // multiply self by X^-h
112    pub fn left_rotate(&self, h: usize) -> Self {
113        let n = self.n;
114
115        let h = h % n;
116        assert!(h < n);
117        let c = &self.coeffs;
118        // c[h], c[h+1], c[h+2], ..., c[n-1],  -c[0], -c[1], ..., -c[h-1]
119        // let r: Vec<T64> = vec![c[h..N], c[0..h].iter().map(|&c_i| -c_i).collect()].concat();
120        let r: Vec<T64> = c[h..n]
121            .iter()
122            .copied()
123            .chain(c[0..h].iter().map(|&x| -x))
124            .collect();
125        Self::from_vec(self.n, r)
126    }
127
128    pub fn from_vec_u64(n: usize, v: Vec<u64>) -> Self {
129        let coeffs = v.iter().map(|c| T64(*c)).collect();
130        Self::from_vec(n, coeffs)
131    }
132}
133
134// apply mod (X^N+1)
135pub fn modulus(n: usize, p: &mut Vec<T64>) {
136    if p.len() < n {
137        return;
138    }
139    for i in n..p.len() {
140        p[i - n] = p[i - n].clone() - p[i].clone();
141        p[i] = T64::zero(());
142    }
143    p.truncate(n);
144}
145
146impl Add<Tn> for Tn {
147    type Output = Self;
148
149    fn add(self, rhs: Self) -> Self {
150        // Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
151        assert_eq!(self.n, rhs.n);
152        Self {
153            n: self.n,
154            coeffs: zip_eq(self.coeffs, rhs.coeffs)
155                .map(|(l, r)| l + r)
156                .collect(),
157        }
158    }
159}
160impl Add<&Tn> for &Tn {
161    type Output = Tn;
162
163    fn add(self, rhs: &Tn) -> Self::Output {
164        // Tn(array::from_fn(|i| self.0[i] + rhs.0[i]))
165        assert_eq!(self.n, rhs.n);
166        Tn {
167            n: self.n,
168            coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
169                .map(|(l, r)| l + r)
170                .collect(),
171        }
172    }
173}
174impl AddAssign for Tn {
175    fn add_assign(&mut self, rhs: Self) {
176        assert_eq!(self.n, rhs.n);
177        for i in 0..self.n {
178            self.coeffs[i] += rhs.coeffs[i];
179        }
180    }
181}
182
183impl Sum<Tn> for Tn {
184    fn sum<I>(iter: I) -> Self
185    where
186        I: Iterator<Item = Self>,
187    {
188        // let mut acc = Tn::<N>::zero();
189        // for e in iter {
190        //     acc += e;
191        // }
192        // acc
193        let first = *iter.next().unwrap().borrow();
194        iter.fold(first, |acc, x| acc + x)
195    }
196}
197
198impl Sub<Tn> for Tn {
199    type Output = Self;
200
201    fn sub(self, rhs: Self) -> Self {
202        assert_eq!(self.n, rhs.n);
203        Self {
204            n: self.n,
205            coeffs: zip_eq(self.coeffs, rhs.coeffs)
206                .map(|(l, r)| l - r)
207                .collect(),
208        }
209    }
210}
211impl Sub<&Tn> for &Tn {
212    type Output = Tn;
213
214    fn sub(self, rhs: &Tn) -> Self::Output {
215        // Tn(array::from_fn(|i| self.0[i] - rhs.0[i]))
216        assert_eq!(self.n, rhs.n);
217        Tn {
218            n: self.n,
219            coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
220                .map(|(l, r)| l - r)
221                .collect(),
222        }
223    }
224}
225
226impl SubAssign for Tn {
227    fn sub_assign(&mut self, rhs: Self) {
228        // for i in 0..N {
229        //     self.0[i] -= rhs.0[i];
230        // }
231        assert_eq!(self.n, rhs.n);
232        for i in 0..self.n {
233            self.coeffs[i] -= rhs.coeffs[i];
234        }
235    }
236}
237
238impl Neg for Tn {
239    type Output = Self;
240
241    fn neg(self) -> Self::Output {
242        // Tn(array::from_fn(|i| -self.0[i]))
243        Self {
244            n: self.n,
245            coeffs: self.coeffs.iter().map(|c_i| -*c_i).collect(),
246        }
247    }
248}
249
250impl PartialEq for Tn {
251    fn eq(&self, other: &Self) -> bool {
252        self.coeffs == other.coeffs && self.n == other.n
253    }
254}
255
256impl Mul<Tn> for Tn {
257    type Output = Self;
258
259    fn mul(self, rhs: Self) -> Self {
260        naive_poly_mul(&self, &rhs)
261    }
262}
263impl Mul<&Tn> for &Tn {
264    type Output = Tn;
265
266    fn mul(self, rhs: &Tn) -> Self::Output {
267        naive_poly_mul(self, rhs)
268    }
269}
270
271fn naive_poly_mul(poly1: &Tn, poly2: &Tn) -> Tn {
272    assert_eq!(poly1.n, poly2.n);
273    let n = poly1.n;
274
275    let poly1: Vec<u128> = poly1.coeffs.iter().map(|c| c.0 as u128).collect();
276    let poly2: Vec<u128> = poly2.coeffs.iter().map(|c| c.0 as u128).collect();
277    let mut result: Vec<u128> = vec![0; (n * 2) - 1];
278    for i in 0..n {
279        for j in 0..n {
280            result[i + j] = result[i + j] + poly1[i] * poly2[j];
281        }
282    }
283
284    // apply mod (X^n + 1))
285    modulus_u128(n, &mut result);
286
287    Tn {
288        n,
289        // coeffs: array::from_fn(|i| T64(result[i] as u64)),
290        coeffs: result.iter().map(|r_i| T64(*r_i as u64)).collect(),
291    }
292}
293fn modulus_u128(n: usize, p: &mut Vec<u128>) {
294    if p.len() < n {
295        return;
296    }
297    for i in n..p.len() {
298        // p[i - n] = p[i - n].clone() - p[i].clone();
299        p[i - n] = p[i - n].wrapping_sub(p[i]);
300        p[i] = 0;
301    }
302    p.truncate(n);
303}
304
305impl Mul<T64> for Tn {
306    type Output = Self;
307    fn mul(self, s: T64) -> Self {
308        Self {
309            n: self.n,
310            // coeffs: array::from_fn(|i| self.coeffs[i] * s),
311            coeffs: self.coeffs.iter().map(|c_i| *c_i * s).collect(),
312        }
313    }
314}
315// mul by u64
316impl Mul<u64> for Tn {
317    type Output = Self;
318    fn mul(self, s: u64) -> Self {
319        // Self(array::from_fn(|i| self.0[i] * s))
320        Tn {
321            n: self.n,
322            coeffs: self.coeffs.iter().map(|c_i| *c_i * s).collect(),
323        }
324    }
325}
326impl Mul<&u64> for &Tn {
327    type Output = Tn;
328    fn mul(self, s: &u64) -> Self::Output {
329        // Tn::<N>(array::from_fn(|i| self.0[i] * *s))
330        Self {
331            n: self.n,
332            coeffs: self.coeffs.iter().map(|c_i| c_i * s).collect(),
333        }
334    }
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340
341    #[test]
342    fn test_left_rotate() {
343        let n: usize = 4;
344        let f = Tn::from_vec(
345            n,
346            vec![2i64, 3, -4, -1]
347                .iter()
348                .map(|c| T64(*c as u64))
349                .collect(),
350        );
351
352        // expect f*x^-3 == -1 -2x -3x^2 +4x^3
353        assert_eq!(
354            f.left_rotate(3),
355            Tn::from_vec(
356                n,
357                vec![-1i64, -2, -3, 4]
358                    .iter()
359                    .map(|c| T64(*c as u64))
360                    .collect(),
361            )
362        );
363        // expect f*x^-1 == 3 -4x -1x^2 -2x^3
364        assert_eq!(
365            f.left_rotate(1),
366            Tn::from_vec(
367                n,
368                vec![3i64, -4, -1, -2]
369                    .iter()
370                    .map(|c| T64(*c as u64))
371                    .collect(),
372            )
373        );
374    }
375}