1//! This file implements the wrapper on top of the ntt.rs to be able to compute
  2//! the NTT for non-prime modulus, specifically for modulus 2^64 (for u64).
  3
  4use crate::ntt::NTT as NTT_p;
  5
  6const P0: u64 = 8070449433331580929; // max use: 1<<62
  7const P1: u64 = 8070450532384645121;
  8
  9// const P0: u64 = 0x80000000080001; // max use: 1<<55
 10// const P1: u64 = 0x80000000130001;
 11
 12#[derive(Debug)]
 13pub struct NTT {}
 14
 15impl NTT {
 16    pub fn ntt(n: usize, a: &Vec<u64>) -> (Vec<u64>, Vec<u64>) {
 17        // apply modulus p_i
 18        let a_0: Vec<u64> = a.iter().map(|a_i| (a_i % P0 + P0) % P0).collect();
 19        let a_1: Vec<u64> = a.iter().map(|a_i| (a_i % P1 + P1) % P1).collect();
 20
 21        let r_0 = NTT_p::ntt(P0, n, &a_0);
 22        let r_1 = NTT_p::ntt(P1, n, &a_1);
 23
 24        (r_0, r_1)
 25    }
 26
 27    pub fn intt(n: usize, r: &(Vec<u64>, Vec<u64>)) -> Vec<u64> {
 28        let a_0 = NTT_p::intt(P0, n, &r.0);
 29        let a_1 = NTT_p::intt(P1, n, &r.1);
 30
 31        reconstruct(a_0, a_1)
 32    }
 33}
 34
 35/// applies CRT to reconstruct the composite original value. Uses Garner's CRT algorithm for two
 36/// moduli: combine (r1 mod p1, r2 mod p2) -> Z/(p1*p2)
 37fn reconstruct(a0: Vec<u64>, a1: Vec<u64>) -> Vec<u64> {
 38    // y_i = q/q_i
 39    // let y0: u128 = P1 as u128;
 40    let y1: u128 = P0 as u128; // Q/P1 = P0*P1/P1 = P0
 41
 42    // y_i^-1 mod q_i = z_i
 43    // let z0: u128 = inv_mod(P0 as u128, y0); // M_i = N_i^-1 mod q_i
 44    let z1: u128 = inv_mod(P1 as u128, y1);
 45
 46    // m1 = q1^-1 mod q2
 47    // aux = (a2 - a1) * m1 mod q2
 48    // a = a1 + (q1 * m1) * aux
 49    let p0: u128 = P0 as u128;
 50    let p1: u128 = P1 as u128;
 51    let a: Vec<u64> = itertools::zip_eq(a0, a1)
 52        .map(|(a0_i, a1_i)| a0_i as u128 + ((p0 * z1) % p1) * (((a1_i - a0_i) as u128 * z1) % p1))
 53        .map(|v| v as u64)
 54        .collect();
 55    a
 56}
 57
 58fn exp_mod(q: u128, x: u128, k: u128) -> u128 {
 59    // work on u128 to avoid overflow
 60    let mut r = 1u128;
 61    let mut x = x.clone();
 62    let mut k = k.clone();
 63    x = x % q;
 64    // exponentiation by square strategy
 65    while k > 0 {
 66        if k % 2 == 1 {
 67            r = (r * x) % q;
 68        }
 69        x = (x * x) % q;
 70        k /= 2;
 71    }
 72    r
 73}
 74/// returns x^-1 mod Q
 75fn inv_mod(q: u128, x: u128) -> u128 {
 76    // by Fermat's Little Theorem, x^-1 mod q \equiv  x^{q-2} mod q
 77    exp_mod(q, x, q - 2)
 78}
 79
 80#[cfg(test)]
 81mod tests {
 82    use super::*;
 83
 84    use anyhow::Result;
 85
 86    #[test]
 87    fn test_dbg() -> Result<()> {
 88        let n: usize = 16;
 89
 90        // let q = 1u128 << 64;
 91        // assert!(P0 as u128 * P1 as u128 > (n as u128 * (q * q)) / 2);
 92
 93        use rand::Rng;
 94        let mut rng = rand::thread_rng();
 95        let a: Vec<u64> = (0..n).map(|_| rng.gen_range(0..(1 << 62))).collect();
 96
 97        dbg!(a.len());
 98
 99        let a_ntt = NTT::ntt(n, &a);
100        dbg!(&a_ntt);
101
102        let a_intt = NTT::intt(n, &a_ntt);
103
104        dbg!(&a_intt);
105        assert_eq!(a_intt, a);
106
107        Ok(())
108    }
109}