warriors/trisha/.vendor/triton-vm/src/arithmetic_domain.rs

use std::ops::Mul;
use std::ops::MulAssign;

use num_traits::ConstOne;
use num_traits::Zero;
use rayon::prelude::*;
use twenty_first::math::traits::FiniteField;
use twenty_first::math::traits::PrimitiveRootOfUnity;
use twenty_first::prelude::*;

use crate::error::ArithmeticDomainError;
use crate::error::USIZE_TO_U64_ERR;

type Result<T> = std::result::Result<T, ArithmeticDomainError>;

/// A co-set of a 2-adic subgroup of the multiplicative group of the
/// [base field](BFieldElement).
///
/// ArithmeticDomain’s are structured in a way that makes
/// [evaluation](Self::evaluate) and [interpolation](Self::interpolate) of
/// [polynomials](twenty_first::prelude::Polynomial) on all the domain’s
/// [values](Self::values) efficient.
///
/// Since the multiplicative subgroup of [p](BFieldElement::P) factors into
/// 2^32·other_factors, where the other_factors are not divisible by 2, the
/// longest [ArithmeticDomain] has [length](Self::len) 2^32.
///
/// The exact co-domain can be specified by setting an
/// [offset](Self::with_offset).
//
// Internal invariant: the generator’s order equals the domain’s length.
// Fields are private to uphold the invariant.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct ArithmeticDomain {
    offset: BFieldElement,
    generator: BFieldElement,
    length: usize,
}

impl ArithmeticDomain {
    /// Create a new domain with the given length.
    /// No offset is applied, but can be added through
    /// [`with_offset()`](Self::with_offset).
    ///
    /// # Errors
    ///
    /// Errors if the domain length is
    /// - not a power of 2, or
    /// - exceeds 2^32.
    pub fn of_length(length: usize) -> Result<Self> {
        let domain = Self {
            offset: BFieldElement::ONE,
            generator: Self::generator_for_length(length as u64)?,
            length,
        };
        Ok(domain)
    }

    /// Set the offset of the domain.
    #[must_use]
    pub fn with_offset(mut self, offset: BFieldElement) -> Self {
        self.offset = offset;
        self
    }

    /// Derive a generator for a domain of the given length.
    ///
    /// # Errors
    ///
    /// Errors if the domain length is
    /// - not a power of 2, or
    /// - exceeds 2^32.
    pub fn generator_for_length(domain_length: u64) -> Result<BFieldElement> {
        let error = ArithmeticDomainError::PrimitiveRootNotSupported(domain_length);
        if domain_length == 0 {
            return Err(error);
        }

        BFieldElement::primitive_root_of_unity(domain_length).ok_or(error)
    }

    /// The length of this domain.
    ///
    /// Equivalenty, the size of the group generated by the
    /// [generator](Self::generator).
    //
    // Even the trivial group contains one element:
    #[expect(clippy::len_without_is_empty)]
    pub fn len(&self) -> usize {
        self.length
    }

    /// The field element that generates the subgroup for the co-domain
    /// that this domain represents.
    ///
    /// The order of the generator is the [length of the domain](Self::len).
    /// See also the co-set defining [offset](Self::offset).
    pub fn generator(&self) -> BFieldElement {
        self.generator
    }

    /// The offset that defines the co-domain.
    ///
    /// Can be set through [`Self::with_offset`].
    pub fn offset(&self) -> BFieldElement {
        self.offset
    }

    /// Evaluate a polynomial on every [value](Self::values) of this domain.
    ///
    /// Turns a polynomial from coefficient form into value form
    /// (i.e., into a codeword).
    pub fn evaluate<FF>(&self, polynomial: &Polynomial<FF>) -> Vec<FF>
    where
        FF: FiniteField
            + MulAssign<BFieldElement>
            + Mul<BFieldElement, Output = FF>
            + From<BFieldElement>
            + 'static,
    {
        let (offset, length) = (self.offset, self.length);
        let evaluate_from = |chunk| Polynomial::from(chunk).fast_coset_evaluate(offset, length);

        // avoid `enumerate` to directly get index of the right type
        let mut indexed_chunks = (0..).zip(polynomial.coefficients().chunks(length));

        // only allocate a bunch of zeros if there are no chunks
        let mut values = indexed_chunks.next().map_or_else(
            || vec![FF::ZERO; length],
            |(_, first_chunk)| evaluate_from(first_chunk),
        );
        for (chunk_index, chunk) in indexed_chunks {
            let coefficient_index = chunk_index * u64::try_from(length).unwrap();
            let scaled_offset = offset.mod_pow(coefficient_index);
            values
                .par_iter_mut()
                .zip(evaluate_from(chunk))
                .for_each(|(value, evaluation)| *value += evaluation * scaled_offset);
        }

        values
    }

    /// Interpolate a polynomial with respect to the [values](Self::values) of
    /// this domain.
    ///
    /// Turns a polynomial from value form (i.e., a codeword) into coefficient
    /// form.
    ///
    /// # Panics
    ///
    /// Panics if the length of the argument does not match the
    /// [length](Self::len) of this domain.
    pub fn interpolate<FF>(&self, values: &[FF]) -> Polynomial<'static, FF>
    where
        FF: FiniteField + MulAssign<BFieldElement> + Mul<BFieldElement, Output = FF>,
    {
        debug_assert_eq!(self.length, values.len()); // required by `fast_coset_interpolate`

        Polynomial::fast_coset_interpolate(self.offset, values)
    }

    /// Move a codeword across domains.
    ///
    /// Given a [polynomial](twenty_first::prelude::Polynomial) in evaluation
    /// form (i.e., a codeword) on _this_ domain, produce that polynomial’s
    /// codeword for the given target domain.
    ///
    /// This is usually done to increase the amount of “redundancy” contained
    /// in the resulting codeword. This, in turn, enables modern low-degree
    /// tests like [FRI](crate::fri::Fri).
    /// However, the target domain may (in principle) be shorter than the
    /// source domain. If the target domain is shorter than or equal to the
    /// polynomial’s degree, information will be lost.
    pub fn low_degree_extension<FF>(&self, codeword: &[FF], target_domain: Self) -> Vec<FF>
    where
        FF: FiniteField
            + MulAssign<BFieldElement>
            + Mul<BFieldElement, Output = FF>
            + From<BFieldElement>
            + 'static,
    {
        target_domain.evaluate(&self.interpolate(codeword))
    }

    #[doc(hidden)]
    #[deprecated(since = "1.0.1", note = "use `domain.value()` instead")]
    pub fn domain_value(&self, n: u32) -> BFieldElement {
        self.value(n)
    }

    #[doc(hidden)]
    #[deprecated(since = "1.0.1", note = "use `domain.values()` instead")]
    pub fn domain_values(&self) -> Vec<BFieldElement> {
        self.values()
    }

    /// The `n`th element of the domain.
    pub fn value(&self, n: u32) -> BFieldElement {
        self.generator.mod_pow_u32(n) * self.offset
    }

    /// All the values that make up this domain.
    pub fn values(&self) -> Vec<BFieldElement> {
        let mut accumulator = BFieldElement::ONE;
        let mut domain_values = Vec::with_capacity(self.length);
        for _ in 0..self.length {
            domain_values.push(accumulator * self.offset);
            accumulator *= self.generator;
        }
        assert_eq!(
            BFieldElement::ONE,
            accumulator,
            "internal error: domain length must equal the order of the generator"
        );

        domain_values
    }

    /// A polynomial that evaluates to 0 on (and only on)
    /// the [domain’s values][Self::values].
    pub fn zerofier(&self) -> Polynomial<'_, BFieldElement> {
        if self.offset.is_zero() {
            return Polynomial::x_to_the(1);
        }

        Polynomial::x_to_the(self.length)
            - Polynomial::from_constant(self.offset.mod_pow(self.length as u64))
    }

    /// [`Self::zerofier`] times the argument.
    /// More performant than polynomial multiplication.
    /// See [`Self::zerofier`] for details.
    pub fn mul_zerofier_with<FF>(&self, polynomial: Polynomial<FF>) -> Polynomial<'static, FF>
    where
        FF: FiniteField + Mul<BFieldElement, Output = FF>,
    {
        // use knowledge of zerofier's shape for faster multiplication
        polynomial.clone().shift_coefficients(self.length)
            - polynomial.scalar_mul(self.offset.mod_pow(self.length as u64))
    }

    /// Raise this domain to the given exponent.
    ///
    /// If `self` corresponds to the domain ο·<ω>, then raising it to the k-th
    /// power results in a new domain ο^k·<ω^k>. This new domain is shorter
    /// by a factor of k, but contains at minimum one element, namely ο^k.
    ///
    /// # Errors
    ///
    /// Returns an error if the exponent is not a power of two.
    pub fn pow(&self, exponent: usize) -> Result<Self> {
        if !exponent.is_power_of_two() {
            let err = ArithmeticDomainError::IllegalExponent(exponent);
            return Err(err);
        }

        // Since 0 is not a power of 2, this division is safe.
        let length = (self.length / exponent).max(1);
        let exponent = u64::try_from(exponent).expect(USIZE_TO_U64_ERR);
        let domain = Self {
            offset: self.offset.mod_pow(exponent),
            generator: self.generator.mod_pow(exponent),
            length,
        };

        Ok(domain)
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
pub(crate) mod tests {
    use itertools::Itertools;
    use proptest::collection::vec;
    use proptest::prelude::*;
    use proptest_arbitrary_interop::arb;
    use test_strategy::proptest;

    use crate::shared_tests::arbitrary_polynomial;
    use crate::shared_tests::arbitrary_polynomial_of_degree;

    use super::*;

    prop_compose! {
        pub fn arbitrary_domain()(
            length in (0_usize..17).prop_map(|x| 1 << x),
            offset in arb(),
        ) -> ArithmeticDomain {
            ArithmeticDomain::of_length(length).unwrap().with_offset(offset)
        }
    }

    #[proptest]
    fn evaluate_empty_polynomial(
        #[strategy(arbitrary_domain())] domain: ArithmeticDomain,
        #[strategy(arbitrary_polynomial_of_degree(-1))] poly: Polynomial<'static, XFieldElement>,
    ) {
        domain.evaluate(&poly);
    }

    #[proptest]
    fn evaluate_constant_polynomial(
        #[strategy(arbitrary_domain())] domain: ArithmeticDomain,
        #[strategy(arbitrary_polynomial_of_degree(0))] poly: Polynomial<'static, XFieldElement>,
    ) {
        domain.evaluate(&poly);
    }

    #[proptest]
    fn evaluate_linear_polynomial(
        #[strategy(arbitrary_domain())] domain: ArithmeticDomain,
        #[strategy(arbitrary_polynomial_of_degree(1))] poly: Polynomial<'static, XFieldElement>,
    ) {
        domain.evaluate(&poly);
    }

    #[proptest]
    fn evaluate_polynomial(
        #[strategy(arbitrary_domain())] domain: ArithmeticDomain,
        #[strategy(arbitrary_polynomial())] polynomial: Polynomial<'static, XFieldElement>,
    ) {
        domain.evaluate(&polynomial);
    }

    #[test]
    fn domain_values() {
        let poly = Polynomial::<BFieldElement>::x_to_the(3);
        let x_cubed_coefficients = poly.clone().into_coefficients();

        for order in [4, 8, 32] {
            let generator = BFieldElement::primitive_root_of_unity(order).unwrap();
            let offset = BFieldElement::generator();
            let b_domain = ArithmeticDomain::of_length(order as usize)
                .unwrap()
                .with_offset(offset);

            let expected_b_values = (0..order)
                .map(|i| offset * generator.mod_pow(i))
                .collect_vec();
            let actual_b_values_1 = b_domain.values();
            let actual_b_values_2 = (0..order as u32).map(|i| b_domain.value(i)).collect_vec();
            assert_eq!(expected_b_values, actual_b_values_1);
            assert_eq!(expected_b_values, actual_b_values_2);

            let values = b_domain.evaluate(&poly);
            assert_ne!(values, x_cubed_coefficients);

            let interpolant = b_domain.interpolate(&values);
            assert_eq!(poly, interpolant);

            // Verify that batch-evaluated values match a manual evaluation
            for i in 0..order {
                let indeterminate = b_domain.value(i as u32);
                let evaluation: BFieldElement = poly.evaluate(indeterminate);
                assert_eq!(evaluation, values[i as usize]);
            }
        }
    }

    #[test]
    fn low_degree_extension() {
        let short_domain_len = 32;
        let long_domain_len = 128;
        let unit_distance = long_domain_len / short_domain_len;

        let short_domain = ArithmeticDomain::of_length(short_domain_len).unwrap();
        let long_domain = ArithmeticDomain::of_length(long_domain_len).unwrap();

        let polynomial = Polynomial::new(bfe_vec![1, 2, 3, 4]);
        let short_codeword = short_domain.evaluate(&polynomial);
        let long_codeword = short_domain.low_degree_extension(&short_codeword, long_domain);

        assert_eq!(long_codeword.len(), long_domain_len);

        let long_codeword_sub_view = long_codeword
            .into_iter()
            .step_by(unit_distance)
            .collect_vec();
        assert_eq!(short_codeword, long_codeword_sub_view);
    }

    #[proptest]
    fn domain_to_the_pow_contains_points_to_the_pow(
        #[strategy(arbitrary_domain())] big_domain: ArithmeticDomain,
        #[strategy(0..=4)]
        #[map(|i| 1_usize << i)]
        exponent: usize,
    ) {
        let small_domain = big_domain.pow(exponent)?;
        if big_domain.length >= exponent {
            prop_assert_eq!(big_domain.length / exponent, small_domain.length);
        } else {
            prop_assert_eq!(1, small_domain.length);
        }

        let exponent = exponent.try_into()?;
        for (big_domain_point, small_domain_point) in big_domain
            .values()
            .into_iter()
            .zip(small_domain.values().into_iter().cycle())
        {
            prop_assert_eq!(big_domain_point.mod_pow(exponent), small_domain_point);
        }
    }

    #[proptest]
    fn pow_converges_on_domain_of_len_1(
        #[strategy(arbitrary_domain())] domain: ArithmeticDomain,
        #[strategy(0..=4)]
        #[map(|i| 1_usize << i)]
        spurious_exponent: usize,
    ) {
        let domain = domain.pow(domain.length)?;
        prop_assert_eq!(BFieldElement::ONE, domain.generator);
        prop_assert_eq!(1, domain.length);

        let domain = domain.pow(spurious_exponent)?;
        prop_assert_eq!(BFieldElement::ONE, domain.generator);
        prop_assert_eq!(1, domain.length);
    }

    #[proptest]
    fn can_evaluate_polynomial_larger_than_domain(
        #[strategy(1_usize..10)] _log_domain_length: usize,
        #[strategy(1_usize..5)] _expansion_factor: usize,
        #[strategy(Just(1 << #_log_domain_length))] domain_length: usize,
        #[strategy(vec(arb(),#domain_length*#_expansion_factor))] coefficients: Vec<BFieldElement>,
        #[strategy(arb())] offset: BFieldElement,
    ) {
        let domain = ArithmeticDomain::of_length(domain_length)
            .unwrap()
            .with_offset(offset);
        let polynomial = Polynomial::new(coefficients);

        let values0 = domain.evaluate(&polynomial);
        let values1 = polynomial.batch_evaluate(&domain.values());
        assert_eq!(values0, values1);
    }

    #[proptest]
    fn zerofier_is_actually_zerofier(#[strategy(arbitrary_domain())] domain: ArithmeticDomain) {
        let actual_zerofier = Polynomial::zerofier(&domain.values());
        prop_assert_eq!(actual_zerofier, domain.zerofier());
    }

    #[proptest]
    fn multiplication_with_zerofier_is_identical_to_method_mul_with_zerofier(
        #[strategy(arbitrary_domain())] domain: ArithmeticDomain,
        #[strategy(arbitrary_polynomial())] polynomial: Polynomial<'static, XFieldElement>,
    ) {
        let mul = domain.zerofier() * polynomial.clone();
        let mul_with = domain.mul_zerofier_with(polynomial);
        prop_assert_eq!(mul, mul_with);
    }
}