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