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