p3-sumcheck 0.7.0

//! Suffix-mode stacked-sumcheck prover.

use alloc::vec;
use alloc::vec::Vec;

use p3_challenger::{FieldChallenger, GrindingChallenger};
use p3_field::{ExtensionField, Field, dot_product};
use p3_maybe_rayon::prelude::*;
use p3_multilinear_util::point::Point;
use p3_multilinear_util::poly::Poly;
use p3_multilinear_util::split_eq::SplitEq;

use crate::lagrange::lagrange_weights_01inf_multi;
use crate::layout::opening::{Opening, ProverMultiClaim};
use crate::layout::prover::{Layout, StackedClaims};
use crate::layout::witness::Table;
use crate::layout::{LayoutStrategy, Witness};
use crate::product_polynomial::ProductPolynomial;
use crate::strategy::{SumcheckProver, VariableOrder};
use crate::svo::{SvoPoint, calculate_accumulators_batch};
use crate::table::{OpeningBatch, OpeningEvals, OpeningRequest};
use crate::{Claim, SumcheckData, extrapolate_01inf};

/// Stacked-sumcheck prover with suffix-first variable binding.
///
/// # Flow
///
/// - SVO accumulators are precomputed at claim-recording time.
/// - Each preprocessing round reads its slice of those accumulators.
/// - The residual product polynomial is built once, after all rounds.
#[derive(Debug, Clone)]
pub struct SuffixProver<F: Field, EF: ExtensionField<F>> {
    /// Recorded opening claims and the layout context that batches them.
    ///
    /// - Suffix binding walks the per-table data directly.
    /// - No separate copy of the stacked polynomial is kept.
    pub(crate) claims: StackedClaims<F, EF>,
}

impl<F: Field, EF: ExtensionField<F>> Layout<F, EF> for SuffixProver<F, EF> {
    fn from_witness(witness: Witness<F>) -> Self {
        // Move the witness fields out so the prover owns them outright.
        // The stacked polynomial is intentionally discarded: every suffix-mode
        // primitive walks the per-table data instead.
        let parts = witness.into_parts();
        Self {
            claims: StackedClaims::new(
                parts.tables,
                parts.placements,
                parts.num_variables,
                parts.folding,
            ),
        }
    }

    fn new_witness(tables: Vec<Table<F>>, folding: usize) -> Witness<F> {
        Witness::new(tables, folding)
    }

    fn claims(&self) -> &StackedClaims<F, EF> {
        &self.claims
    }

    /// Records opening claims for the selected columns of one table.
    ///
    /// All requested columns share one sampled local opening point.
    ///
    /// - Current openings evaluate a column at that point.
    /// - Next openings evaluate the repeat-last successor view at that point.
    /// - Returned evaluations list all current openings first, then all next openings.
    ///
    /// # Arguments
    ///
    /// - `table_idx`  — source table index.
    /// - `batch`      — current and next columns opened at this point.
    /// - `challenger` — Fiat-Shamir transcript.
    ///
    /// # Fiat-Shamir
    ///
    /// - Samples the opening point internally from the transcript.
    /// - Absorbs the evaluations before returning.
    /// - The verifier performs the symmetric absorption.
    ///
    /// # Panics
    ///
    /// - At least one current or next column must be requested.
    #[tracing::instrument(skip_all)]
    fn eval_at<Ch>(
        &mut self,
        table_idx: usize,
        batch: &OpeningRequest,
        point: &Point<EF>,
        challenger: &mut Ch,
    ) -> OpeningEvals<EF>
    where
        Ch: FieldChallenger<F> + GrindingChallenger<Witness = F>,
    {
        // Split the request into its two column groups.
        let current = batch.current();
        let next = batch.next();
        // Precondition: opening nothing would silently push an empty claim.
        assert!(
            !batch.is_empty(),
            "opening schedule must name at least one column"
        );

        let table = &self.claims.tables[table_idx];
        // The opening point lives in the table's local frame, one coordinate per variable.
        debug_assert_eq!(point.num_variables(), table.num_variables());

        // Factorise the point with the suffix split; every selected column reuses it.
        let point = SvoPoint::new_unpacked(self.claims.folding, point, VariableOrder::Suffix);

        // Current group: evaluate each column at the point.
        // Each entry yields an opening (carrying preprocessing residuals) plus the bare eval.
        // Each column is an independent O(2^n) pass, so columns evaluate in parallel.
        let (current_openings, current_evals): (Vec<_>, Vec<EF>) = current
            .into_par_iter()
            .copied()
            .map(|poly_idx| {
                let (eval, partial_evals) = point.eval(table.poly(poly_idx));
                (Opening::new_with_data(poly_idx, eval, partial_evals), eval)
            })
            .unzip();

        // Next group: evaluate the repeat-last successor view at the same point.
        // `current_openings` is fully built above and only read here, so this is safe to
        // run in parallel alongside it.
        let (next_openings, next_evals): (Vec<_>, Vec<EF>) = next
            .into_par_iter()
            .copied()
            .map(|poly_idx| {
                // Reuse: if this column was already opened in the current group,
                // its last-round equality residual feeds the successor computation
                // for free, avoiding a redundant pass over the column.
                //
                // The per-round list is empty when there are no SVO rounds (folding
                // is zero), so fall back to recomputing the residual rather than
                // unwrapping an absent last round.
                let d_eq = current_openings
                    .iter()
                    .find(|opening| opening.poly_idx() == Some(poly_idx))
                    .and_then(|opening| opening.data().rounds().last())
                    .map(|round| round.poly());
                let (eval, partial_evals) = point.eval_next_suffix(table.poly(poly_idx), d_eq);
                (Opening::new_with_data(poly_idx, eval, partial_evals), eval)
            })
            .unzip();

        // Bind the evaluations into the transcript, current group first then next.
        // The verifier absorbs the same bytes in the same order.
        challenger.observe_algebra_slice(&current_evals);
        challenger.observe_algebra_slice(&next_evals);

        // Store the batch with its shared SVO point.
        self.claims.claim_map[table_idx].push(ProverMultiClaim::new(
            point,
            current_openings,
            next_openings,
        ));

        // Return both eval groups in the canonical current-then-next order.
        OpeningBatch::new(current_evals, next_evals)
    }

    /// Samples a virtual evaluation on the full stacked polynomial.
    ///
    /// # Why heavier than prefix binding
    ///
    /// The stacked evaluation factors per column via the selector:
    ///
    /// ```text
    ///     stacked(point) = sum_{i}  eq(selector_i, point_selector_part)
    ///                               * col_i(point_local_part)
    /// ```
    ///
    /// # Flow
    ///
    /// - Each column is evaluated at its local sub-point.
    /// - Per-column partials are collected on the fly.
    /// - Those partials feed the SVO accumulator batcher.
    #[tracing::instrument(skip_all)]
    fn add_virtual_eval<Ch>(&mut self, challenger: &mut Ch) -> EF
    where
        Ch: FieldChallenger<F> + GrindingChallenger<Witness = F>,
    {
        // Sample a challenge point covering every stacked variable.
        let point = Point::expand_from_univariate(
            challenger.sample_algebra_element(),
            self.claims.num_variables,
        );

        // Per-column accumulation state:
        //
        // - eval    : running stacked evaluation.
        // - openings: one virtual opening per column, carrying SVO partials.
        // - weights : per-column selector-equality scalars.
        let mut eval = EF::ZERO;
        let mut openings = Vec::new();
        let mut weights = Vec::new();

        for placement in &self.claims.placements {
            for (poly_idx, selector) in placement.selectors().iter().enumerate() {
                // Source column behind this slot.
                let poly = self.claims.tables[placement.idx()].poly(poly_idx);

                // Split the challenge into (selector_bits, local_bits).
                let (selector_part, local_part) = point.split_at(selector.num_variables());

                // Scalar weight: eq(selector, selector_part) for this column.
                let weight =
                    Point::eval_eq::<EF>(selector.point().as_slice(), selector_part.as_slice());

                // Factorise the local part with the suffix split, then evaluate.
                let local_svo =
                    SvoPoint::new_unpacked(self.claims.folding, &local_part, VariableOrder::Suffix);
                let (column_eval, partial_evals) = local_svo.eval(poly);

                // Record a virtual opening (no source column tag) with partials.
                let opening = Opening {
                    poly_idx: None,
                    eval: column_eval,
                    data: partial_evals,
                };

                // Add the weighted column evaluation into the stacked total.
                eval += weight * column_eval;

                // Stash opening and weight for the accumulator-batcher call.
                openings.push(opening);
                weights.push(weight);
            }
        }

        // Batch every per-column opening into per-round SVO accumulators.
        let accumulators = calculate_accumulators_batch(
            &ProverMultiClaim::new(
                SvoPoint::new_unpacked(self.claims.folding, &point, VariableOrder::Suffix),
                openings,
                Vec::new(),
            ),
            &weights,
        );

        // Debug-only consistency check:
        //
        // - hand-rolled weighted sum must equal the direct stacked evaluation.
        // - accumulators batched per column must equal the single-opening batch.
        #[cfg(debug_assertions)]
        {
            // Materialise the stacked polynomial with no challenges applied.
            let poly = &self.compress_stacked(&Point::default());
            // Check 1: weighted sum equals the direct evaluation.
            assert_eq!(eval, poly.eval_base(&point));

            // Build the reference opening by evaluating the materialised poly directly.
            let ref_svo = SvoPoint::<EF, EF>::new_unpacked(
                self.claims.folding,
                &point,
                VariableOrder::Suffix,
            );
            let (ref_eval, ref_partials) = ref_svo.eval(poly.as_view());
            let opening = Opening {
                poly_idx: None,
                eval: ref_eval,
                data: ref_partials,
            };
            // Check 2: the reference evaluation matches the weighted one.
            assert_eq!(eval, ref_eval);
            // Check 3: accumulators from per-column batching match the single-opening batch.
            assert_eq!(
                accumulators,
                calculate_accumulators_batch(
                    &ProverMultiClaim::new(
                        SvoPoint::new_unpacked(self.claims.folding, &point, VariableOrder::Suffix),
                        vec![opening],
                        Vec::new(),
                    ),
                    &[EF::ONE],
                ),
            );
        }

        // Commit the evaluation to the transcript and record the claim.
        challenger.observe_algebra_element(eval);
        self.claims.virtual_claims.push(Claim {
            point,
            eval,
            data: accumulators,
        });

        eval
    }

    /// Finalises SVO preprocessing and returns the residual sumcheck prover.
    ///
    /// # Returns
    ///
    /// - Residual sumcheck prover over the unpacked product polynomial.
    /// - Folding challenges sampled during preprocessing.
    ///
    /// # Algorithm
    ///
    /// ```text
    ///     Phase | Action
    ///     ------+------------------------------------------------------------
    ///       1   | Sample batching challenge  a; flatten alphas by opening_idx.
    ///       2   | Pre-batch per-claim accumulators with the a-powers.
    ///       3   | Loop over preprocessing rounds:
    ///               a. (h(0), h(inf)) = dot(accumulators, Lagrange weights).
    ///               b. Sample challenge r; extrapolate the running sum.
    ///       4   | Compose the residual product polynomial from compressed slots.
    /// ```
    #[tracing::instrument(skip_all)]
    fn into_sumcheck<Ch>(
        self,
        sumcheck_data: &mut SumcheckData<F, EF>,
        pow_bits: usize,
        challenger: &mut Ch,
    ) -> (SumcheckProver<F, EF>, Point<EF>)
    where
        Ch: FieldChallenger<F> + GrindingChallenger<Witness = F>,
    {
        // Sanity: preprocessing cannot consume more rounds than the stacked arity.
        assert!(self.claims.folding <= self.claims.num_variables);
        let alpha: EF = challenger.sample_algebra_element();
        let n_claims = self.num_claims();

        // Stage A: batch per-claim accumulators using insertion-order alpha powers.
        //
        // - Iteration order is placement order, matching `sum` and `combine_weights`.
        // - Each claim consumes exactly `claim.len()` consecutive powers from
        //   the shared iterator, so the per-claim alpha vector is aligned with
        //   the claim's opening list by construction.
        let mut alphas = alpha.powers();
        let accumulators: Vec<_> = self
            .claims
            .concrete_claims()
            .map(|claim| {
                let per_claim: Vec<EF> = alphas.by_ref().take(claim.len()).collect();
                calculate_accumulators_batch(claim, &per_claim)
            })
            .collect();

        // Stage C: drive the preprocessing rounds from the accumulators.
        let mut sum = self.claims.sum(alpha);
        let mut rs: Vec<EF> = vec![];

        // First alpha power assigned to the virtual claims, sitting just past the concrete claims.
        // The claim count is fixed for the whole fold, so this exponentiation is loop-invariant.
        let alpha_base = alpha.exp_u64(n_claims as u64);

        for round_idx in 0..self.claims.folding {
            // Lagrange weights at the challenges sampled so far.
            let weights = lagrange_weights_01inf_multi(&rs);

            // Round-coefficient identity (linearity of the dot product):
            //
            //     c0    = sum_c  dot(claim_c.accs[0], weights)
            //           + sum_v  alpha_v * dot(virtual_v.accs[0], weights)
            //     c_inf = same with accs[1]
            //
            // - Concrete claims carry alpha pre-batched in stage B.
            // - Virtual claims keep a separate scalar per claim.
            // - No intermediate element-wise accumulator is needed.
            let mut c0 = EF::ZERO;
            let mut c_inf = EF::ZERO;

            for accs in &accumulators {
                c0 += dot_product::<EF, _, _>(
                    accs[round_idx][0].iter().copied(),
                    weights.iter().copied(),
                );
                c_inf += dot_product::<EF, _, _>(
                    accs[round_idx][1].iter().copied(),
                    weights.iter().copied(),
                );
            }

            // Virtual-claim contributions: scale each claim's dot by its alpha power.
            for (vc, alpha_i) in self
                .claims
                .virtual_claims
                .iter()
                .zip(alpha.shifted_powers(alpha_base))
            {
                let vc_accs = &vc.data;
                c0 += alpha_i
                    * dot_product::<EF, _, _>(
                        vc_accs[round_idx][0].iter().copied(),
                        weights.iter().copied(),
                    );
                c_inf += alpha_i
                    * dot_product::<EF, _, _>(
                        vc_accs[round_idx][1].iter().copied(),
                        weights.iter().copied(),
                    );
            }

            // Observe coefficients, sample r, extrapolate the running sum.
            let r = sumcheck_data.observe_and_sample(challenger, c0, c_inf, pow_bits);
            sum = extrapolate_01inf(c0, sum - c0, c_inf, r);
            rs.push(r);
        }

        // Stage D: materialise the residual product polynomial.
        //
        // - Suffix binding folds variables in reverse.
        // - The residual poly therefore lives in the reversed-challenges frame.
        let rs = Point::new(rs);
        // Reverse the challenges before handing them to the compressors.
        let reversed = rs.reversed();
        // Factor 1 of the product: the compressed stacked poly at rs.
        // No external scaling here; the plain path keeps the running sum unchanged.
        let compressed = self.compress_stacked(&reversed);
        // Factor 2 of the product: the batched equality-weight poly.
        let weights = self.combine_weights(&reversed, alpha);
        // Pair them; the product polynomial drives the remaining rounds.
        let poly = ProductPolynomial::new_unpacked(VariableOrder::Suffix, compressed, weights);
        // Cross-check: the dot product of the two factors must equal the
        // running sum accumulated across the preprocessing rounds.
        debug_assert_eq!(poly.dot_product(), sum);

        (SumcheckProver::new(poly, sum), rs)
    }

    fn strategy() -> LayoutStrategy {
        LayoutStrategy::new(false, VariableOrder::Suffix)
    }
}

impl<F: Field, EF: ExtensionField<F>> SuffixProver<F, EF> {
    /// Compress every stacked-table slot by fixing the suffix challenges.
    #[tracing::instrument(skip_all)]
    pub(crate) fn compress_stacked(&self, rs: &Point<EF>) -> Poly<EF> {
        self.compress_stacked_scaled(rs, EF::ONE)
    }

    /// Compress every stacked-table slot, folding `scale` into the equality table.
    ///
    /// ```text
    ///     out[slot, x_rest] = sum_{y in {0,1}^|r|}  scale * eq(r, y) * col(x_rest, y)
    /// ```
    ///
    /// One output slot per column.
    /// Writes never overlap.
    /// Output arity equals the stacked arity minus the challenge count.
    ///
    /// # Arguments
    ///
    /// - `rs` — suffix challenges already sampled.
    /// - `scale` — extra factor folded into the equality table.
    ///
    /// # Why a scale parameter
    ///
    /// - A non-unit `scale` lets the caller absorb a combining challenge into
    ///   the residual factor without a second pass.
    ///
    /// # Panics
    ///
    /// - `scale` is zero: a zero scale silently zeroes the residual.
    #[tracing::instrument(skip_all)]
    pub(crate) fn compress_stacked_scaled(&self, rs: &Point<EF>, scale: EF) -> Poly<EF> {
        assert!(rs.num_variables() <= self.claims.num_variables);
        assert!(scale != EF::ZERO, "compress scale must be non-zero");
        // Output spans the residual stacked space.
        // Size is 2^(num_variables - |rs|).
        let mut out = Poly::<EF>::zero(self.claims.num_variables - rs.num_variables());
        // Bake the scalar into the prefix-half equality table.
        // Each slot compression then returns scale * eq(r, y) * col(...) in one pass.
        let rs = SplitEq::new_unpacked(rs, scale);

        for placement in &self.claims.placements {
            for (poly_idx, selector) in placement.selectors().iter().enumerate() {
                let poly = self.claims.tables[placement.idx()].poly(poly_idx);
                assert!(rs.num_variables() <= poly.num_variables());
                // Slot start in the compressed output.
                let off = selector.index() << (poly.num_variables() - rs.num_variables());
                // Write this column's compression into its own slot.
                rs.compress_suffix_into(
                    &mut out.as_mut_slice()
                        [off..off + (1 << (poly.num_variables() - rs.num_variables()))],
                    poly.as_view(),
                );
            }
        }
        out
    }

    /// Builds the residual weight polynomial after the SVO rounds.
    ///
    /// # Contributions
    ///
    /// - Concrete claim: factored equality table scaled by
    ///   `alpha^i * eq(svo_part, rs)`, written into the owning slot only.
    /// - Virtual claim: scaled equality table written across the full output.
    #[tracing::instrument(skip_all)]
    pub(crate) fn combine_weights(&self, rs: &Point<EF>, alpha: EF) -> Poly<EF> {
        // Preconditions: challenge count matches the folding depth.
        assert_eq!(rs.num_variables(), self.claims.folding);
        // Output arity: stacked arity minus the folded challenges.
        let mut out = Poly::<EF>::zero(self.claims.num_variables - rs.num_variables());

        let mut alphas = alpha.powers();

        // Concrete claims: write each into the slot its column's selector addresses.
        // Walk order matches the batched claim, so alpha powers stay aligned.
        for placement in &self.claims.placements {
            let num_variables_table = self.num_variables_table(placement.idx());
            let slot_size = 1usize << num_variables_table;
            for claim in &self.claims.claim_map[placement.idx()] {
                // Current group: equality weight of the column at the claim point.
                for opening in claim.current_openings() {
                    // The opening's column tells us which selector picks the slot.
                    let col = opening.poly_idx().unwrap();
                    let off = placement.selectors()[col].index() << num_variables_table;
                    // Fold the scalar slot range down by the SVO depth.
                    let folded_range =
                        (off >> self.claims.folding)..((off + slot_size) >> self.claims.folding);
                    let scale = alphas.next().unwrap();
                    claim
                        .point()
                        .accumulate_into(&mut out.as_mut_slice()[folded_range], rs, scale);
                }
                // Next group: same slot, but the repeat-last successor weight.
                for opening in claim.next_openings() {
                    let col = opening.poly_idx().unwrap();
                    let off = placement.selectors()[col].index() << num_variables_table;
                    let folded_range =
                        (off >> self.claims.folding)..((off + slot_size) >> self.claims.folding);
                    let scale = alphas.next().unwrap();
                    claim.point().accumulate_next_suffix_into(
                        &mut out.as_mut_slice()[folded_range],
                        rs,
                        scale,
                    );
                }
            }
        }

        // Virtual claims: span the full output; alpha continues after concrete ones.
        let mut alpha_i = alpha.exp_u64(self.num_claims() as u64);
        for claim in &self.claims.virtual_claims {
            // Split the claim point into (rest-of-space, svo-sub-point).
            let (rest, svo) = claim
                .point
                .split_at(claim.point.num_variables() - rs.num_variables());
            // Scalar weight: alpha^i times the equality between svo part and rs.
            let scale = alpha_i * Point::eval_eq(svo.as_slice(), rs.as_slice());
            // Contribute the scaled equality table across the whole output.
            SplitEq::new_packed(&rest, scale).accumulate_into(out.as_mut_slice(), None);
            // Advance alpha for the next virtual claim.
            alpha_i *= alpha;
        }

        out
    }
}