ark_poly/domain/radix2/

mod.rs

//! This module defines `Radix2EvaluationDomain`, an `EvaluationDomain`
//! for performing various kinds of polynomial arithmetic on top of
//! fields that are FFT-friendly.
//!
//! `Radix2EvaluationDomain` supports FFTs of size at most `2^F::TWO_ADICITY`.

pub use crate::domain::utils::{bitreverse_permutation_in_place, Elements};
use crate::domain::{DomainCoeff, EvaluationDomain};
use ark_ff::{FftField, Field};
use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
use ark_std::{fmt, vec::*};

mod fft;

/// Factor that determines if a the degree aware FFT should be called.
const DEGREE_AWARE_FFT_THRESHOLD_FACTOR: usize = 1 << 2;

/// Defines a domain over which finite field (I)FFTs can be performed. Works
/// only for fields that have a large multiplicative subgroup of size that is
/// a power-of-2.
#[derive(Copy, Clone, Hash, Eq, PartialEq, CanonicalSerialize, CanonicalDeserialize)]
pub struct Radix2EvaluationDomain<F: Field> {
    /// The size of the domain.
    pub size: u64,
    /// `log_2(self.size)`.
    pub log_size_of_group: u32,
    /// Size of the domain as a field element.
    pub size_as_field_element: F,
    /// Inverse of the size in the field.
    pub size_inv: F,
    /// A generator of the subgroup.
    pub group_gen: F,
    /// Inverse of the generator of the subgroup.
    pub group_gen_inv: F,
    /// Offset that specifies the coset.
    pub offset: F,
    /// Inverse of the offset that specifies the coset.
    pub offset_inv: F,
    /// Constant coefficient for the vanishing polynomial.
    /// Equals `self.offset^self.size`.
    pub offset_pow_size: F,
}

impl<F: Field> fmt::Debug for Radix2EvaluationDomain<F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Radix-2 multiplicative subgroup of size {}", self.size)
    }
}

impl<F: FftField> EvaluationDomain<F> for Radix2EvaluationDomain<F> {
    type Elements = Elements<F>;

    /// Construct a domain that is large enough for evaluations of a polynomial
    /// having `num_coeffs` coefficients.
    fn new(num_coeffs: usize) -> Option<Self> {
        let size = num_coeffs.next_power_of_two() as u64;

        let log_size_of_group = size.trailing_zeros();

        // libfqfft uses > https://github.com/scipr-lab/libfqfft/blob/e0183b2cef7d4c5deb21a6eaf3fe3b586d738fe0/libfqfft/evaluation_domain/domains/basic_radix2_domain.tcc#L33
        if log_size_of_group > F::TWO_ADICITY {
            return None;
        }

        // Compute the generator for the multiplicative subgroup.
        // It should be the 2^(log_size_of_group) root of unity.
        let group_gen = F::get_root_of_unity(size)?;
        // Check that it is indeed the 2^(log_size_of_group) root of unity.
        debug_assert_eq!(group_gen.pow([size]), F::one());
        let size_as_field_element = F::from(size);

        Some(Self {
            size,
            log_size_of_group,
            size_as_field_element,
            size_inv: size_as_field_element.inverse()?,
            group_gen,
            group_gen_inv: group_gen.inverse()?,
            offset: F::ONE,
            offset_inv: F::ONE,
            offset_pow_size: F::ONE,
        })
    }

    fn get_coset(&self, offset: F) -> Option<Self> {
        Some(Self {
            offset,
            offset_inv: offset.inverse()?,
            offset_pow_size: offset.pow([self.size]),
            ..*self
        })
    }

    fn compute_size_of_domain(num_coeffs: usize) -> Option<usize> {
        let size = num_coeffs.checked_next_power_of_two()?;
        (size.trailing_zeros() <= F::TWO_ADICITY).then_some(size)
    }

    #[inline]
    fn size(&self) -> usize {
        self.size.try_into().unwrap()
    }

    #[inline]
    fn log_size_of_group(&self) -> u64 {
        self.log_size_of_group as u64
    }

    #[inline]
    fn size_inv(&self) -> F {
        self.size_inv
    }

    #[inline]
    fn group_gen(&self) -> F {
        self.group_gen
    }

    #[inline]
    fn group_gen_inv(&self) -> F {
        self.group_gen_inv
    }

    #[inline]
    fn coset_offset(&self) -> F {
        self.offset
    }

    #[inline]
    fn coset_offset_inv(&self) -> F {
        self.offset_inv
    }

    #[inline]
    fn coset_offset_pow_size(&self) -> F {
        self.offset_pow_size
    }

    #[inline]
    fn fft_in_place<T: DomainCoeff<F>>(&self, coeffs: &mut Vec<T>) {
        if coeffs.len() * DEGREE_AWARE_FFT_THRESHOLD_FACTOR <= self.size() {
            self.degree_aware_fft_in_place(coeffs);
        } else {
            coeffs.resize(self.size(), T::zero());
            self.in_order_fft_in_place(coeffs);
        }
    }

    #[inline]
    fn ifft_in_place<T: DomainCoeff<F>>(&self, evals: &mut Vec<T>) {
        evals.resize(self.size(), T::zero());
        self.in_order_ifft_in_place(&mut *evals);
    }

    /// Return an iterator over the elements of the domain.
    fn elements(&self) -> Elements<F> {
        Elements {
            cur_elem: self.offset,
            cur_pow: 0,
            size: self.size,
            group_gen: self.group_gen,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::DEGREE_AWARE_FFT_THRESHOLD_FACTOR;
    use crate::{
        polynomial::{univariate::*, DenseUVPolynomial, Polynomial},
        EvaluationDomain, Radix2EvaluationDomain,
    };
    use ark_ff::{FftField, Field, One, UniformRand, Zero};
    use ark_std::{collections::BTreeSet, rand::Rng, test_rng, vec};
    use ark_test_curves::bls12_381::Fr;

    #[test]
    fn vanishing_polynomial_evaluation() {
        let rng = &mut test_rng();
        for coeffs in 0..10 {
            let domain = Radix2EvaluationDomain::<Fr>::new(coeffs).unwrap();
            let coset_domain = domain.get_coset(Fr::GENERATOR).unwrap();
            let z = domain.vanishing_polynomial();
            let z_coset = coset_domain.vanishing_polynomial();
            for _ in 0..100 {
                let point: Fr = rng.gen();
                assert_eq!(
                    z.evaluate(&point),
                    domain.evaluate_vanishing_polynomial(point)
                );
                assert_eq!(
                    z_coset.evaluate(&point),
                    coset_domain.evaluate_vanishing_polynomial(point)
                );
            }
        }
    }

    #[test]
    fn vanishing_polynomial_vanishes_on_domain() {
        for coeffs in 0..1000 {
            let domain = Radix2EvaluationDomain::<Fr>::new(coeffs).unwrap();
            let z = domain.vanishing_polynomial();
            for point in domain.elements() {
                assert!(z.evaluate(&point).is_zero())
            }

            let coset_domain = domain.get_coset(Fr::GENERATOR).unwrap();
            let z = coset_domain.vanishing_polynomial();
            for point in coset_domain.elements() {
                assert!(z.evaluate(&point).is_zero())
            }
        }
    }

    #[test]
    fn filter_polynomial_test() {
        for log_domain_size in 1..=4 {
            let domain_size = 1 << log_domain_size;
            let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
            for log_subdomain_size in 1..=log_domain_size {
                let subdomain_size = 1 << log_subdomain_size;
                let subdomain = Radix2EvaluationDomain::<Fr>::new(subdomain_size).unwrap();

                // Obtain all possible offsets of `subdomain` within `domain`.
                let mut possible_offsets = vec![Fr::one()];
                let domain_generator = domain.group_gen();

                let mut offset = domain_generator;
                let subdomain_generator = subdomain.group_gen();
                while offset != subdomain_generator {
                    possible_offsets.push(offset);
                    offset *= domain_generator;
                }

                assert_eq!(possible_offsets.len(), domain_size / subdomain_size);

                // Get all possible cosets of `subdomain` within `domain`.
                let cosets = possible_offsets
                    .iter()
                    .map(|offset| subdomain.get_coset(*offset).unwrap());

                for coset in cosets {
                    let coset_elements = coset.elements().collect::<BTreeSet<_>>();
                    let filter_poly = domain.filter_polynomial(&coset);
                    assert_eq!(filter_poly.degree(), domain_size - subdomain_size);
                    for element in domain.elements() {
                        let evaluation = domain.evaluate_filter_polynomial(&coset, element);
                        assert_eq!(evaluation, filter_poly.evaluate(&element));
                        if coset_elements.contains(&element) {
                            assert_eq!(evaluation, Fr::one())
                        } else {
                            assert_eq!(evaluation, Fr::zero())
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn size_of_elements() {
        for coeffs in 1..10 {
            let size = 1 << coeffs;
            let domain = Radix2EvaluationDomain::<Fr>::new(size).unwrap();
            let domain_size = domain.size();
            assert_eq!(domain_size, domain.elements().count());
        }
    }

    #[test]
    fn elements_contents() {
        for coeffs in 1..10 {
            let size = 1 << coeffs;
            let domain = Radix2EvaluationDomain::<Fr>::new(size).unwrap();
            let offset = Fr::GENERATOR;
            let coset_domain = domain.get_coset(offset).unwrap();
            for (i, (x, coset_x)) in domain.elements().zip(coset_domain.elements()).enumerate() {
                assert_eq!(x, domain.group_gen.pow([i as u64]));
                assert_eq!(x, domain.element(i));
                assert_eq!(coset_x, offset * coset_domain.group_gen.pow([i as u64]));
                assert_eq!(coset_x, coset_domain.element(i));
            }
        }
    }

    /// Test that lagrange interpolation for a random polynomial at a random
    /// point works.
    #[test]
    fn non_systematic_lagrange_coefficients_test() {
        for domain_dim in 1..10 {
            let domain_size = 1 << domain_dim;
            let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
            let coset_domain = domain.get_coset(Fr::GENERATOR).unwrap();
            // Get random pt + lagrange coefficients
            let rand_pt = Fr::rand(&mut test_rng());
            let lagrange_coeffs = domain.evaluate_all_lagrange_coefficients(rand_pt);
            let coset_lagrange_coeffs = coset_domain.evaluate_all_lagrange_coefficients(rand_pt);

            // Sample the random polynomial, evaluate it over the domain and the random
            // point.
            let rand_poly = DensePolynomial::<Fr>::rand(domain_size - 1, &mut test_rng());
            let poly_evals = domain.fft(rand_poly.coeffs());
            let coset_poly_evals = coset_domain.fft(rand_poly.coeffs());
            let actual_eval = rand_poly.evaluate(&rand_pt);

            // Do lagrange interpolation, and compare against the actual evaluation
            let mut interpolated_eval = Fr::zero();
            let mut coset_interpolated_eval = Fr::zero();
            for i in 0..domain_size {
                interpolated_eval += lagrange_coeffs[i] * poly_evals[i];
                coset_interpolated_eval += coset_lagrange_coeffs[i] * coset_poly_evals[i];
            }
            assert_eq!(actual_eval, interpolated_eval);
            assert_eq!(actual_eval, coset_interpolated_eval);
        }
    }

    /// Test that lagrange coefficients for a point in the domain is correct
    #[test]
    fn systematic_lagrange_coefficients_test() {
        // This runs in time O(N^2) in the domain size, so keep the domain dimension
        // low. We generate lagrange coefficients for each element in the domain.
        for domain_dim in 1..5 {
            let domain_size = 1 << domain_dim;
            let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
            let coset_domain = domain.get_coset(Fr::GENERATOR).unwrap();
            for (i, (x, coset_x)) in domain.elements().zip(coset_domain.elements()).enumerate() {
                let lagrange_coeffs = domain.evaluate_all_lagrange_coefficients(x);
                let coset_lagrange_coeffs =
                    coset_domain.evaluate_all_lagrange_coefficients(coset_x);
                for (j, (y, coset_y)) in lagrange_coeffs
                    .into_iter()
                    .zip(coset_lagrange_coeffs)
                    .enumerate()
                {
                    // Lagrange coefficient for the evaluation point, which should be 1
                    if i == j {
                        assert_eq!(y, Fr::one());
                        assert_eq!(coset_y, Fr::one());
                    } else {
                        assert_eq!(y, Fr::zero());
                        assert_eq!(coset_y, Fr::zero());
                    }
                }
            }
        }
    }

    #[test]
    fn test_fft_correctness() {
        // Tests that the ffts output the correct result.
        // This assumes a correct polynomial evaluation at point procedure.
        // It tests consistency of FFT/IFFT, and coset_fft/coset_ifft,
        // along with testing that each individual evaluation is correct.

        // Runs in time O(degree^2)
        let log_degree = 5;
        let degree = 1 << log_degree;
        let rand_poly = DensePolynomial::<Fr>::rand(degree - 1, &mut test_rng());

        for log_domain_size in log_degree..(log_degree + 2) {
            let domain_size = 1 << log_domain_size;
            let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
            let coset_domain =
                Radix2EvaluationDomain::<Fr>::new_coset(domain_size, Fr::GENERATOR).unwrap();
            let poly_evals = domain.fft(&rand_poly.coeffs);
            let poly_coset_evals = coset_domain.fft(&rand_poly.coeffs);

            for (i, (x, coset_x)) in domain.elements().zip(coset_domain.elements()).enumerate() {
                assert_eq!(poly_evals[i], rand_poly.evaluate(&x));
                assert_eq!(poly_coset_evals[i], rand_poly.evaluate(&coset_x));
            }

            let rand_poly_from_subgroup =
                DensePolynomial::from_coefficients_vec(domain.ifft(&poly_evals));
            let rand_poly_from_coset =
                DensePolynomial::from_coefficients_vec(coset_domain.ifft(&poly_coset_evals));

            assert_eq!(
                rand_poly, rand_poly_from_subgroup,
                "degree = {}, domain size = {}",
                degree, domain_size
            );
            assert_eq!(
                rand_poly, rand_poly_from_coset,
                "degree = {}, domain size = {}",
                degree, domain_size
            );
        }
    }

    #[test]
    fn degree_aware_fft_correctness() {
        // Test that the degree aware FFT (O(n log d)) matches the regular FFT (O(n log n)).
        let num_coeffs = 1 << 5;
        let rand_poly = DensePolynomial::<Fr>::rand(num_coeffs - 1, &mut test_rng());
        let domain_size = num_coeffs * DEGREE_AWARE_FFT_THRESHOLD_FACTOR;
        let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
        let coset_domain = domain.get_coset(Fr::GENERATOR).unwrap();

        let deg_aware_fft_evals = domain.fft(&rand_poly);
        let coset_deg_aware_fft_evals = coset_domain.fft(&rand_poly);

        for (i, (x, coset_x)) in domain.elements().zip(coset_domain.elements()).enumerate() {
            assert_eq!(deg_aware_fft_evals[i], rand_poly.evaluate(&x));
            assert_eq!(coset_deg_aware_fft_evals[i], rand_poly.evaluate(&coset_x));
        }
    }

    #[test]
    fn test_roots_of_unity() {
        // Tests that the roots of unity result is the same as domain.elements()
        let max_degree = 10;
        for log_domain_size in 0..max_degree {
            let domain_size = 1 << log_domain_size;
            let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
            let actual_roots = domain.roots_of_unity(domain.group_gen);
            for &value in &actual_roots {
                assert!(domain.evaluate_vanishing_polynomial(value).is_zero());
            }
            let expected_roots_elements = domain.elements();
            for (expected, &actual) in expected_roots_elements.zip(&actual_roots) {
                assert_eq!(expected, actual);
            }
            assert_eq!(actual_roots.len(), domain_size / 2);
        }
    }

    #[test]
    #[cfg(feature = "parallel")]
    fn parallel_fft_consistency() {
        use ark_std::{test_rng, vec::*};
        use ark_test_curves::bls12_381::Fr;

        // This implements the Cooley-Turkey FFT, derived from libfqfft
        // The libfqfft implementation uses pseudocode from [CLRS 2n Ed, pp. 864].
        fn serial_radix2_fft(a: &mut [Fr], omega: Fr, log_n: u32) {
            use ark_std::convert::TryFrom;
            let n = u32::try_from(a.len())
                .expect("cannot perform FFTs larger on vectors of len > (1 << 32)");
            assert_eq!(n, 1 << log_n);

            // swap coefficients in place
            crate::domain::utils::bitreverse_permutation_in_place(a, log_n);

            let mut m = 1;
            for _i in 1..=log_n {
                // w_m is 2^i-th root of unity
                let w_m = omega.pow([(n / (2 * m)) as u64]);

                let mut k = 0;
                while k < n {
                    // w = w_m^j at the start of every loop iteration
                    let mut w = Fr::one();
                    for j in 0..m {
                        let mut t = a[(k + j + m) as usize];
                        t *= w;
                        let mut tmp = a[(k + j) as usize];
                        tmp -= t;
                        a[(k + j + m) as usize] = tmp;
                        a[(k + j) as usize] += t;
                        w *= &w_m;
                    }

                    k += 2 * m;
                }

                m *= 2;
            }
        }

        fn serial_radix2_ifft(a: &mut [Fr], omega: Fr, log_n: u32) {
            serial_radix2_fft(a, omega.inverse().unwrap(), log_n);
            let domain_size_inv = Fr::from(a.len() as u64).inverse().unwrap();
            for coeff in a.iter_mut() {
                *coeff *= Fr::from(domain_size_inv);
            }
        }

        fn serial_radix2_coset_fft(a: &mut [Fr], omega: Fr, log_n: u32) {
            let coset_shift = Fr::GENERATOR;
            let mut cur_pow = Fr::one();
            for coeff in a.iter_mut() {
                *coeff *= cur_pow;
                cur_pow *= coset_shift;
            }
            serial_radix2_fft(a, omega, log_n);
        }

        fn serial_radix2_coset_ifft(a: &mut [Fr], omega: Fr, log_n: u32) {
            serial_radix2_ifft(a, omega, log_n);
            let coset_shift = Fr::GENERATOR.inverse().unwrap();
            let mut cur_pow = Fr::one();
            for coeff in a.iter_mut() {
                *coeff *= cur_pow;
                cur_pow *= coset_shift;
            }
        }

        fn test_consistency<R: Rng>(rng: &mut R, max_coeffs: u32) {
            for _ in 0..5 {
                for log_d in 0..max_coeffs {
                    let d = 1 << log_d;

                    let expected_poly = (0..d).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
                    let mut expected_vec = expected_poly.clone();
                    let mut actual_vec = expected_vec.clone();

                    let domain = Radix2EvaluationDomain::new(d).unwrap();
                    let coset_domain = domain.get_coset(Fr::GENERATOR).unwrap();

                    serial_radix2_fft(&mut expected_vec, domain.group_gen, log_d);
                    domain.fft_in_place(&mut actual_vec);
                    assert_eq!(expected_vec, actual_vec);

                    serial_radix2_ifft(&mut expected_vec, domain.group_gen, log_d);
                    domain.ifft_in_place(&mut actual_vec);
                    assert_eq!(expected_vec, actual_vec);
                    assert_eq!(expected_vec, expected_poly);

                    serial_radix2_coset_fft(&mut expected_vec, domain.group_gen, log_d);
                    coset_domain.fft_in_place(&mut actual_vec);
                    assert_eq!(expected_vec, actual_vec);

                    serial_radix2_coset_ifft(&mut expected_vec, domain.group_gen, log_d);
                    coset_domain.ifft_in_place(&mut actual_vec);
                    assert_eq!(expected_vec, actual_vec);
                }
            }
        }

        let rng = &mut test_rng();

        test_consistency(rng, 10);
    }

    #[test]
    fn test_domain_creation() {
        // Test powers of 2
        let domain_8 = Radix2EvaluationDomain::<Fr>::new(8).unwrap();
        assert_eq!(domain_8.size(), 8);
        assert_eq!(domain_8.log_size_of_group(), 3);

        // Test rounding to next power of 2
        let domain_7 = Radix2EvaluationDomain::<Fr>::new(7).unwrap();
        assert_eq!(domain_7.size(), 8);
    }

    #[test]
    fn test_root_of_unity() {
        let domain = Radix2EvaluationDomain::<Fr>::new(8).unwrap();
        let root = domain.group_gen();

        // Check order: root^8 should be 1
        let expected = root.pow([8]);
        assert_eq!(expected, Fr::one());
    }

    #[test]
    fn test_inverse_root_of_unity() {
        let domain = Radix2EvaluationDomain::<Fr>::new(8).unwrap();
        let root = domain.group_gen();
        let root_inv = domain.group_gen_inv();

        // root * root_inv should be 1
        let expected = root * root_inv;
        assert_eq!(expected, Fr::one());
    }

    #[test]
    fn test_size_inverse() {
        let domain = Radix2EvaluationDomain::<Fr>::new(8).unwrap();
        let size_inv = domain.size_inv();
        let expected = Fr::from(8).inverse().unwrap();

        assert_eq!(size_inv, expected);
    }

    #[test]
    fn test_fft_ifft_identity() {
        let domain = Radix2EvaluationDomain::<Fr>::new(8).unwrap();
        let mut coeffs = vec![
            Fr::from(1),
            Fr::from(2),
            Fr::from(3),
            Fr::from(4),
            Fr::from(5),
            Fr::from(6),
            Fr::from(7),
            Fr::from(8),
        ];

        let original = coeffs.clone();
        domain.fft_in_place(&mut coeffs);
        domain.ifft_in_place(&mut coeffs);

        // Check if IFFT(FFT(coeffs)) == coeffs
        assert_eq!(coeffs, original);
    }

    #[test]
    fn test_vanishing_polynomial() {
        let domain = Radix2EvaluationDomain::<Fr>::new(4).unwrap();
        let z = domain.vanishing_polynomial();

        for elem in domain.elements() {
            assert_eq!(z.evaluate(&elem), Fr::zero());
        }
    }

    #[test]
    fn test_compute_size_of_domain() {
        assert_eq!(
            Radix2EvaluationDomain::<Fr>::compute_size_of_domain(7),
            Some(8)
        );
        assert_eq!(
            Radix2EvaluationDomain::<Fr>::compute_size_of_domain(8),
            Some(8)
        );
        assert_eq!(
            Radix2EvaluationDomain::<Fr>::compute_size_of_domain(15),
            Some(16)
        );
    }
}