1use crate::FieldExt;
  2
  3pub fn fft<F>(coeffs: &[F], domain: &[F]) -> Vec<F>
  4where
  5    F: FieldExt,
  6{
  7    debug_assert_eq!(coeffs.len(), domain.len());
  8    if coeffs.len() == 1 {
  9        return coeffs.to_vec();
 10    }
 11
 12    // Split into evens and odds
 13    let L = coeffs
 14        .iter()
 15        .enumerate()
 16        .filter(|(i, _)| i % 2 == 0)
 17        .map(|(_, x)| *x)
 18        .collect::<Vec<F>>();
 19
 20    let R = coeffs
 21        .iter()
 22        .enumerate()
 23        .filter(|(i, _)| i % 2 == 1)
 24        .map(|(_, x)| *x)
 25        .collect::<Vec<F>>();
 26
 27    // Square the domain values
 28    let domain_squared: Vec<F> = (0..(domain.len() / 2)).map(|i| domain[i * 2]).collect();
 29
 30    let fft_e = fft(&L, &domain_squared);
 31    let fft_o = fft(&R, &domain_squared);
 32
 33    let mut evals_L = Vec::with_capacity(coeffs.len() / 2);
 34    let mut evals_R = Vec::with_capacity(coeffs.len() / 2);
 35    for i in 0..(coeffs.len() / 2) {
 36        // We can use the previous evaluations to create a list of evaluations
 37        // of the domain
 38        evals_L.push(fft_e[i] + fft_o[i] * domain[i]);
 39        evals_R.push(fft_e[i] - fft_o[i] * domain[i]);
 40    }
 41
 42    evals_L.extend(evals_R);
 43    return evals_L;
 44}
 45
 46#[cfg(test)]
 47mod tests {
 48    use halo2curves::ff::Field;
 49    use halo2curves::ff::PrimeField;
 50    use halo2curves::pasta::Fp;
 51
 52    // Test the fft function by running the inverse fft
 53    fn ifft<F: FieldExt + Field>(domain: &[F], evals: &[F]) -> Vec<F> {
 54        let mut coeffs = vec![];
 55        let len_mod_inv = F::from(domain.len() as u64).invert().unwrap();
 56        let vals = fft(&evals, &domain);
 57
 58        coeffs.push(vals[0] * len_mod_inv);
 59        for val in vals[1..].iter().rev() {
 60            coeffs.push(*val * len_mod_inv);
 61        }
 62
 63        coeffs
 64    }
 65
 66    use super::*;
 67    #[test]
 68    fn test_fft_ifft() {
 69        // f(x) = 1 + 2x + 3x^2 + 4x^3
 70        let mut coeffs = vec![
 71            Fp::from(1),
 72            Fp::from(2),
 73            Fp::from(3),
 74            Fp::from(4),
 75            Fp::from(5),
 76            Fp::from(6),
 77            Fp::from(7),
 78            Fp::from(81),
 79        ];
 80
 81        let mut domain = vec![];
 82
 83        let root_of_unity = Fp::ROOT_OF_UNITY;
 84
 85        let subgroup_order = (coeffs.len() * 2).next_power_of_two();
 86
 87        coeffs.resize(subgroup_order, Fp::ZERO);
 88
 89        // Generator for the subgroup with order _subgroup_order_ in the field
 90        let generator = root_of_unity.pow(&[
 91            2u32.pow(32 - ((subgroup_order as f64).log2() as u32)) as u64,
 92            0,
 93            0,
 94            0,
 95        ]);
 96
 97        for i in 0..(subgroup_order) {
 98            domain.push(generator.pow(&[i as u64, 0, 0, 0]));
 99        }
100
101        let mut expected_evals = vec![];
102
103        for w in &domain {
104            let mut eval = Fp::ZERO;
105            for (i, coeff) in (&coeffs).iter().enumerate() {
106                eval += *coeff * w.pow(&[i as u64, 0, 0, 0]);
107            }
108            expected_evals.push(eval);
109        }
110
111        let evals = fft(&coeffs, &domain);
112        debug_assert!(evals == expected_evals);
113
114        let recovered_coeffs = ifft(&domain, &evals);
115        debug_assert!(recovered_coeffs == coeffs);
116    }
117}