number theoretic transform

ntt_notes.rs

pub trait FFT: Sized + Copy {
    type F: Sized
        + Copy
        + From<Self>
        + Neg
        + Add<Output = Self::F>
        + Div<Output = Self::F>
        + Mul<Output = Self::F>
        + Sub<Output = Self::F>;

    const ZERO: Self;

    /// A primitive nth root of one raised to the powers 0, 1, 2, ..., n/2 - 1
    fn get_roots(n: usize, inverse: bool) -> Vec<Self::F>;
    /// 1 for forward transform, 1/n for inverse transform
    fn get_factor(n: usize, inverse: bool) -> Self::F;
    /// The inverse of Self::F::from()
    fn extract(f: Self::F) -> Self;
}

// NTT notes: see problem 30-6 in CLRS for details, keeping in mind that
//      2187 and  410692747 are inverses and 2^26th roots of 1 mod (7<<26)+1
//  15311432 and  469870224 are inverses and 2^23rd roots of 1 mod (119<<23)+1
// 440564289 and 1713844692 are inverses and 2^27th roots of 1 mod (15<<27)+1
//       125 and 2267742733 are inverses and 2^30th roots of 1 mod (3<<30)+1
impl FFT for i64 {
    type F = Field;

    const ZERO: Self = 0;

    fn get_roots(n: usize, inverse: bool) -> Vec<Self::F> {
        assert!(n <= 1 << 23);
        let mut prim_root = Self::F::from(15_311_432);
        if inverse {
            prim_root = prim_root.recip();
        }
        for _ in (0..).take_while(|&i| n < 1 << (23 - i)) {
            prim_root = prim_root * prim_root;
        }

        let mut roots = Vec::with_capacity(n / 2);
        let mut root = Self::F::from(1);
        for _ in 0..roots.capacity() {
            roots.push(root);
            root = root * prim_root;
        }
        roots
    }

    fn get_factor(n: usize, inverse: bool) -> Self::F {
        Self::F::from(if inverse { n as Self } else { 1 }).recip()
    }

    fn extract(f: Self::F) -> Self {
        f.val
    }
}