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