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