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