p3q/crates/p3q-fri/src/domain.rs

//! Evaluation-domain math shared by every FRI consumer (the generic
//! low-degree verifier in `p3q-verifier` and the AIR DEEP-STARK in
//! `p3q-stark`).
//!
//! This is the single definition of the FRI fold formula and the
//! multiplicative-subgroup LDE domain. Both the prover and verifier —
//! generic and AIR — call these; there is exactly one place the fold
//! arithmetic lives.
//!
//! Conventions:
//! - The LDE domain `D_lde` is the multiplicative subgroup of order
//!   `N = 2^log_lde` generated by `omega_lde`, a primitive `N`-th root
//!   of unity over Goldilocks.
//! - Index `k` refers to the point `omega_lde^k` (no coset shift; v1
//!   uses the pure subgroup).
//! - Binary FRI fold halves the domain each round: layer `i` has size
//!   `2^(log_lde − i)` and generator `omega_lde^(2^i)`.
//! - At layer `i`, index `k` and its sibling `k ^ (1 << (log_lde−i−1))`
//!   correspond to points `x` and `−x` since `omega^{2^i · N/2} = −1`.
//!
//! FRI fold rule (binary, no coset offset): given layer-`i` evals
//! `f_i(x) = a`, `f_i(−x) = b` and challenge `r`,
//!
//! ```text
//! f_{i+1}(x^2) = (a + b) / 2 + r · (a − b) / (2 · x)
//! ```

use p3q_core::{P3qError, Result};
use p3q_field::{Field, Goldilocks};

/// Compute the LDE-domain generator `omega_lde = g^{(p−1)/2^log_lde}`.
pub fn lde_generator(log_lde: u32) -> Result<Goldilocks> {
    Goldilocks::two_adic_generator(log_lde).ok_or(P3qError::FieldOutOfRange)
}

/// Evaluate `omega_lde^k` for a public `k` via square-and-multiply.
#[inline]
pub fn point_at_index(omega_lde: Goldilocks, k: u32) -> Goldilocks {
    omega_lde.pow(k as u64)
}

/// Sibling index for the binary FRI fold at a domain of log-size
/// `log_layer`: the pair of index `k` is `(k, k ^ (N/2))`.
#[inline]
pub fn sibling_index(k: u32, log_layer: u32) -> u32 {
    debug_assert!(log_layer >= 1);
    k ^ (1u32 << (log_layer - 1))
}

/// One binary FRI fold step.
///
/// `a = f_i(x)` (eval at `x = omega^k`), `b = f_i(−x)` (eval at
/// `−x = omega^{k+N/2}`), `x` the layer-`i` point, `r` the round
/// challenge. Returns `f_{i+1}(x^2)`. Rejects when `2x` has no inverse
/// (impossible in a multiplicative subgroup; guarded for defence).
pub fn fri_fold(a: Goldilocks, b: Goldilocks, x: Goldilocks, r: Goldilocks) -> Result<Goldilocks> {
    let sum = a + b;
    let diff = a - b;
    let two_inv = Goldilocks::from_u64(2)
        .inverse()
        .ok_or(P3qError::FieldOutOfRange)?;
    let half_sum = sum * two_inv;
    let two_x = Goldilocks::from_u64(2) * x;
    let two_x_inv = two_x
        .inverse()
        .ok_or(P3qError::FriVerification("FRI fold: 2x is zero"))?;
    let half_diff_over_x = diff * two_x_inv;
    Ok(half_sum + r * half_diff_over_x)
}

/// Evaluate a polynomial in coefficient form at `x` via Horner.
pub fn horner(coeffs: &[Goldilocks], x: Goldilocks) -> Goldilocks {
    let mut acc = Goldilocks::default();
    for &c in coeffs.iter().rev() {
        acc = acc * x + c;
    }
    acc
}

/// Inverse DFT on a multiplicative subgroup of size `N = 2^k` with
/// generator `omega`. Naive `O(N^2)` — fine for the small final
/// polynomials and the trace-interpolation sizes used here. Returns the
/// coefficient vector of the unique degree-`<N` polynomial that
/// interpolates `evals` on `{omega^0, …, omega^{N−1}}`.
pub fn inverse_dft_on_subgroup(evals: &[Goldilocks], omega: Goldilocks) -> Result<Vec<Goldilocks>> {
    let n = evals.len();
    if n == 0 {
        return Err(P3qError::InvalidLength("inverse_dft: empty evals"));
    }
    let n_inv = Goldilocks::from_u64(n as u64)
        .inverse()
        .ok_or(P3qError::FieldOutOfRange)?;
    let omega_inv = omega.inverse().ok_or(P3qError::FieldOutOfRange)?;
    let mut coeffs = Vec::with_capacity(n);
    for k in 0..n {
        let mut acc = Goldilocks::default();
        let mut wk = Goldilocks::one();
        let omega_inv_k = omega_inv.pow(k as u64);
        for v in evals.iter() {
            acc = acc + *v * wk;
            wk = wk * omega_inv_k;
        }
        coeffs.push(acc * n_inv);
    }
    Ok(coeffs)
}

/// Forward DFT on a multiplicative subgroup of size `N = 2^k` with
/// generator `omega`: evaluate the coefficient vector `coeffs` (length
/// `≤ N`) on `{omega^0, …, omega^{N−1}}`. Naive `O(N·deg)` via Horner.
pub fn dft_on_subgroup(coeffs: &[Goldilocks], omega: Goldilocks, n: usize) -> Vec<Goldilocks> {
    (0..n)
        .map(|k| horner(coeffs, omega.pow(k as u64)))
        .collect()
}

/// Forward DFT on the coset `shift · <omega>` of size `N = 2^k`:
/// evaluate `coeffs` at `{shift·omega^0, …, shift·omega^{N−1}}`. Used by
/// the AIR prover to evaluate trace / constraint polynomials on a coset
/// of the LDE domain where the trace-domain vanishing polynomial `Z_H`
/// is invertible (it vanishes on the base subgroup, never on a proper
/// coset). Naive `O(N·deg)`.
pub fn coset_dft(
    coeffs: &[Goldilocks],
    shift: Goldilocks,
    omega: Goldilocks,
    n: usize,
) -> Vec<Goldilocks> {
    (0..n)
        .map(|k| horner(coeffs, shift * omega.pow(k as u64)))
        .collect()
}

/// Inverse DFT on the coset `shift · <omega>` of size `N = 2^k`: recover
/// the coefficient vector of the degree-`<N` polynomial whose values on
/// `{shift·omega^k}` are `evals`. This is the inverse of [`coset_dft`].
///
/// Derivation: if `p(x) = Σ c_m x^m`, then `p(shift·omega^k) = Σ c_m
/// shift^m omega^{mk}`. Let `d_m = c_m shift^m`. Then `evals` is the
/// plain subgroup DFT of `d`, so `d = idft(evals)` and `c_m = d_m /
/// shift^m`.
pub fn coset_idft(
    evals: &[Goldilocks],
    shift: Goldilocks,
    omega: Goldilocks,
) -> Result<Vec<Goldilocks>> {
    let d = inverse_dft_on_subgroup(evals, omega)?;
    let shift_inv = shift.inverse().ok_or(P3qError::FieldOutOfRange)?;
    let mut acc = Goldilocks::one();
    let mut coeffs = Vec::with_capacity(d.len());
    for dm in d {
        coeffs.push(dm * acc);
        acc = acc * shift_inv;
    }
    Ok(coeffs)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn lde_generator_matches_two_adic() {
        for k in 1..=10 {
            assert_eq!(
                lde_generator(k).unwrap(),
                Goldilocks::two_adic_generator(k).unwrap()
            );
        }
    }

    #[test]
    fn sibling_index_is_involutive() {
        for log in 1..=10u32 {
            for k in 0u32..(1 << log) {
                let s = sibling_index(k, log);
                assert_eq!(sibling_index(s, log), k);
                assert_eq!(k ^ s, 1u32 << (log - 1));
            }
        }
    }

    #[test]
    fn fri_fold_is_polynomial_decomposition() {
        // f(x) = c0 + c1 x + c2 x^2 + c3 x^3; one fold gives a degree-<2
        // poly in y = x^2: (c0 + r c1) + (c2 + r c3) y.
        let c = [
            Goldilocks::from_u64(7),
            Goldilocks::from_u64(11),
            Goldilocks::from_u64(13),
            Goldilocks::from_u64(17),
        ];
        let r = Goldilocks::from_u64(19);
        let omega = Goldilocks::two_adic_generator(2).unwrap();
        let x = omega;
        let a = horner(&c, x);
        let b = horner(&c, -x);
        let folded = fri_fold(a, b, x, r).unwrap();
        let y = x * x;
        let expected = horner(&[c[0] + r * c[1], c[2] + r * c[3]], y);
        assert_eq!(folded, expected);
    }

    #[test]
    fn dft_inverse_roundtrip() {
        let coeffs: Vec<_> = (1u64..=8).map(Goldilocks::from_u64).collect();
        let omega = Goldilocks::two_adic_generator(3).unwrap();
        let evals = dft_on_subgroup(&coeffs, omega, 8);
        let recovered = inverse_dft_on_subgroup(&evals, omega).unwrap();
        assert_eq!(recovered, coeffs);
    }

    #[test]
    fn horner_matches_naive() {
        let coeffs: Vec<_> = (1u64..=5).map(Goldilocks::from_u64).collect();
        let x = Goldilocks::from_u64(3);
        let want = (1..=5u64).fold(Goldilocks::default(), |acc, i| {
            acc + Goldilocks::from_u64(i) * x.pow(i - 1)
        });
        assert_eq!(horner(&coeffs, x), want);
    }

    #[test]
    fn coset_dft_inverse_roundtrip() {
        let coeffs: Vec<_> = (3u64..11).map(Goldilocks::from_u64).collect();
        let omega = Goldilocks::two_adic_generator(3).unwrap();
        let shift = Goldilocks::from_u64(7); // generator coset shift
        let evals = coset_dft(&coeffs, shift, omega, 8);
        let recovered = coset_idft(&evals, shift, omega).unwrap();
        assert_eq!(recovered, coeffs);
    }

    #[test]
    fn coset_dft_matches_pointwise() {
        let coeffs: Vec<_> = (1u64..=4).map(Goldilocks::from_u64).collect();
        let omega = Goldilocks::two_adic_generator(2).unwrap();
        let shift = Goldilocks::from_u64(7);
        let evals = coset_dft(&coeffs, shift, omega, 4);
        for (k, e) in evals.iter().enumerate() {
            let x = shift * omega.pow(k as u64);
            assert_eq!(*e, horner(&coeffs, x));
        }
    }
}