p3-sumcheck 0.8.0

//! Shared per-table claim state for the stacked-sumcheck provers.
//!
//! - Both binding modes record the same openings.
//! - Both batch them the same way under one challenge.
//! - The committed polynomial is the only mode-specific datum.
//! - The prefix prover keeps that polynomial separately and embeds this state.

use alloc::vec::Vec;

use p3_field::{ExtensionField, Field, dot_product};

use crate::Claim;
use crate::lagrange::lagrange_weights_01inf_multi;
use crate::layout::opening::{EqSvoPartials, NextSvoPartials, Opening};
use crate::layout::witness::{Table, TablePlacement};
use crate::layout::{ProverMultiClaim, ProverVirtualClaim};
use crate::svo::{SvoAccumulators, SvoPoint, calculate_accumulators_batch};
use crate::table::{OpeningEvals, OpeningRequest, TableShape};

/// Opening claims recorded against one stacked polynomial, shared by both binding modes.
///
/// - Holds the layout context needed to batch the recorded claims under one challenge.
/// - Both provers embed this and add only their mode-specific state on top.
#[derive(Debug, Clone)]
pub struct StackedClaims<F: Field, EF: ExtensionField<F>> {
    /// Source tables behind the stacked polynomial.
    pub(crate) tables: Vec<Table<F>>,
    /// Per-table placement metadata inside the stacked polynomial.
    pub(crate) placements: Vec<TablePlacement>,
    /// Number of variables of the stacked polynomial.
    pub(crate) num_variables: usize,
    /// Number of preprocessing rounds consumed before residual sumcheck.
    pub(crate) folding: usize,
    /// Concrete claims recorded per source table.
    ///
    /// # Invariants
    ///
    /// - Every opening stored here is tied to a concrete source column.
    /// - Virtual openings never enter this map.
    /// - Claims are appended in insertion order.
    pub(crate) claim_map: Vec<Vec<ProverMultiClaim<F, EF>>>,
    /// Virtual claims sampled directly on the stacked polynomial.
    pub(crate) virtual_claims: Vec<ProverVirtualClaim<EF>>,
}

impl<F: Field, EF: ExtensionField<F>> StackedClaims<F, EF> {
    /// Creates empty claim state for a freshly committed stacked polynomial.
    ///
    /// # Arguments
    ///
    /// - `tables` — source tables behind the stacked polynomial.
    /// - `placements` — per-table slot assignment inside the stacked polynomial.
    /// - `num_variables` — arity of the stacked polynomial.
    /// - `folding` — preprocessing rounds consumed before residual sumcheck.
    pub(crate) fn new(
        tables: Vec<Table<F>>,
        placements: Vec<TablePlacement>,
        num_variables: usize,
        folding: usize,
    ) -> Self {
        // One empty claim list per source table.
        // Virtual claims live in their own bucket.
        let claim_map = (0..tables.len()).map(|_| Vec::new()).collect();
        Self {
            tables,
            placements,
            num_variables,
            folding,
            claim_map,
            virtual_claims: Vec::new(),
        }
    }

    /// Returns the number of preprocessing rounds.
    pub(crate) const fn folding(&self) -> usize {
        self.folding
    }

    /// Returns the number of variables of the stacked polynomial.
    pub(crate) const fn num_variables(&self) -> usize {
        self.num_variables
    }

    /// Returns the number of variables of table `id`.
    pub(crate) fn num_variables_table(&self, id: usize) -> usize {
        self.tables[id].num_variables()
    }

    /// Returns every source table's `(arity, width)`, in caller order.
    pub(crate) fn table_shapes(&self) -> Vec<TableShape> {
        self.tables.iter().map(Table::shape).collect()
    }

    /// Returns source table `id`.
    pub fn table(&self, id: usize) -> &Table<F> {
        &self.tables[id]
    }

    /// Returns the total number of concrete openings recorded so far.
    pub(crate) fn num_claims(&self) -> usize {
        self.claim_map
            .iter()
            .flat_map(|claims| claims.iter().map(ProverMultiClaim::len))
            .sum()
    }

    /// Returns the number of out-of-domain claims recorded on the stacked polynomial.
    pub(crate) const fn num_virtual_claims(&self) -> usize {
        self.virtual_claims.len()
    }

    /// Appends one batch's claims from evaluations the caller already holds.
    ///
    /// The caller factorises the point in its own binding mode and passes it in.
    ///
    /// # Arguments
    ///
    /// - `table_idx` — source table the request names.
    /// - `batch` — columns opened directly and through the successor view.
    /// - `point` — the batch's opening point, already factorised.
    /// - `evals` — one evaluation per opened column, direct entries first.
    ///
    /// # Panics
    ///
    /// When preprocessing rounds are configured. Each such round reads a residual of the
    /// column, which an evaluation on its own does not carry.
    pub(crate) fn record_known(
        &mut self,
        table_idx: usize,
        batch: &OpeningRequest,
        point: SvoPoint<F, EF>,
        evals: &OpeningEvals<EF>,
    ) {
        assert_eq!(
            self.folding, 0,
            "a supplied evaluation carries no preprocessing residual for a round to read"
        );
        debug_assert!(batch.has_same_shape(evals));

        // No preprocessing round runs, so every opening's residual list is empty.
        let current = batch
            .current()
            .iter()
            .zip(evals.current())
            .map(|(&poly_idx, &eval)| {
                Opening::new_with_data(poly_idx, eval, EqSvoPartials::new(Vec::new()))
            })
            .collect();
        let next = batch
            .next()
            .iter()
            .zip(evals.next())
            .map(|(&poly_idx, &eval)| {
                Opening::new_with_data(poly_idx, eval, NextSvoPartials::new(Vec::new()))
            })
            .collect();

        self.claim_map[table_idx].push(ProverMultiClaim::new(point, current, next));
    }

    /// Walks concrete claims in placement order.
    ///
    /// - Outer: placements, in the order the witness laid them out.
    /// - Inner: the claims recorded against each placement's source table.
    pub(crate) fn concrete_claims(&self) -> impl Iterator<Item = &ProverMultiClaim<F, EF>> {
        self.placements
            .iter()
            .flat_map(|placement| self.claim_map[placement.idx()].iter())
    }

    /// Computes the batched claimed sum from concrete and virtual openings.
    ///
    /// # Identity
    ///
    /// ```text
    ///     sum = sum_{i}  alpha^i * eval_i
    /// ```
    ///
    /// # Alpha ordering
    ///
    /// Powers of `alpha` are handed out in placement order, not insertion order:
    ///
    /// - Outer: placements, in the order the witness laid them out.
    /// - Middle: claims recorded against that placement's source table.
    /// - Inner: openings inside each claim, in the order they were recorded.
    ///
    /// # Virtual claims
    ///
    /// - Virtual evaluations continue the same alpha sequence.
    /// - They start at `alpha^n`, with `n` the total concrete opening count.
    ///
    /// # Verifier agreement
    ///
    /// The verifier walks its claim registry with the same three-loop order,
    /// so both sides assign the same `alpha^i` to the same claim point.
    pub(crate) fn sum(&self, alpha: EF) -> EF {
        let mut sum = EF::ZERO;
        let mut alphas = alpha.powers();

        // Walk every concrete opening in the canonical batching order.
        //     placements -> claims -> current openings -> next openings
        // Each opening consumes the next power of alpha, matching the verifier.
        for placement in &self.placements {
            for claim in &self.claim_map[placement.idx()] {
                // Current group first.
                for opening in claim.current_openings() {
                    sum += opening.eval() * alphas.next().unwrap();
                }
                // Next group second, continuing the same power sequence.
                for opening in claim.next_openings() {
                    sum += opening.eval() * alphas.next().unwrap();
                }
            }
        }

        // Virtual claims continue the alpha sequence right after the concrete ones.
        sum += dot_product::<EF, _, _>(
            self.virtual_claims.iter().map(Claim::eval),
            alpha.shifted_powers(alpha.exp_u64(self.num_claims() as u64)),
        );

        sum
    }

    /// Batches every recorded claim's accumulators under `alpha`.
    ///
    /// Powers follow the order of [`Self::sum`]:
    /// placements, claims, current openings, successor openings, then virtual claims.
    pub(crate) fn batched_accumulators(&self, alpha: EF) -> BatchedAccumulators<'_, EF> {
        // Iteration order is placement order, matching `sum`. 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 concrete = self
            .concrete_claims()
            .map(|claim| {
                let per_claim: Vec<EF> = alphas.by_ref().take(claim.len()).collect();
                calculate_accumulators_batch(claim, &per_claim)
            })
            .collect();

        // Virtual claims continue the alpha sequence right after the concrete ones.
        let virtual_alphas = alpha
            .shifted_powers(alpha.exp_u64(self.num_claims() as u64))
            .take(self.virtual_claims.len())
            .collect();

        BatchedAccumulators {
            concrete,
            virtual_claims: &self.virtual_claims,
            virtual_alphas,
        }
    }
}

/// SVO accumulators of every recorded claim, batched under one challenge.
///
/// - Concrete claims carry their alpha powers inside their accumulators.
/// - Virtual claims keep their own accumulators and one alpha power each.
pub(crate) struct BatchedAccumulators<'a, EF: Field> {
    /// One accumulator set per concrete claim, in placement order.
    concrete: Vec<SvoAccumulators<EF>>,
    /// Virtual claims, in recording order.
    virtual_claims: &'a [ProverVirtualClaim<EF>],
    /// Alpha power of each virtual claim, continuing after the concrete openings.
    virtual_alphas: Vec<EF>,
}

impl<EF: Field> BatchedAccumulators<'_, EF> {
    /// Round polynomial `(h(0), h(inf))` of the next preprocessing round.
    ///
    /// `rs` holds the challenges drawn so far, so its length is the round index.
    ///
    /// # Identity
    ///
    /// Linearity of the dot product gives:
    ///
    /// ```text
    ///     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 by [`StackedClaims::batched_accumulators`].
    /// - Virtual claims keep a separate scalar per claim.
    /// - No intermediate element-wise accumulator is needed.
    pub(crate) fn round_coefficients(&self, rs: &[EF]) -> (EF, EF) {
        let weights = lagrange_weights_01inf_multi(rs);
        let round_idx = rs.len();

        let mut c0 = EF::ZERO;
        let mut c_inf = EF::ZERO;

        for accs in &self.concrete {
            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.virtual_claims.iter().zip(self.virtual_alphas.iter()) {
            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(),
                );
        }

        (c0, c_inf)
    }
}

#[cfg(test)]
mod tests {
    use alloc::vec;

    use p3_baby_bear::BabyBear;
    use p3_field::PrimeCharacteristicRing;
    use p3_field::extension::BinomialExtensionField;
    use p3_multilinear_util::point::Point;

    use super::*;
    use crate::layout::opening::{EqSvoPartials, NextSvoPartials, Opening};
    use crate::strategy::VariableOrder;
    use crate::svo::SvoPoint;

    type F = BabyBear;
    type EF = BinomialExtensionField<F, 4>;

    // A single zero column of the given arity, enough to drive the table getters.
    fn table(arity: usize) -> Table<F> {
        Table::zero(1, arity)
    }

    // Placement pointing back at source table `idx`.
    // Selectors stay empty: only the table index is read by these methods.
    fn placement(idx: usize) -> TablePlacement {
        TablePlacement::new(idx, vec![])
    }

    // Concrete claim with the given current-group and next-group opening evaluations.
    // The shared point and per-round payloads are empty.
    // `sum`, `num_claims`, and `concrete_claims` read only the opening evals.
    fn claim(current: &[u64], next: &[u64]) -> ProverMultiClaim<F, EF> {
        // Empty SVO point: zero folded variables.
        let point = SvoPoint::<F, EF>::new_unpacked(0, &Point::new(vec![]), VariableOrder::Prefix);
        // Current openings carry the eq-weight payload type.
        let current_openings = current
            .iter()
            .map(|&e| Opening::new_with_data(0, EF::from_u64(e), EqSvoPartials::new(vec![])))
            .collect();
        // Next openings carry the successor payload type.
        let next_openings = next
            .iter()
            .map(|&e| Opening::new_with_data(0, EF::from_u64(e), NextSvoPartials::new(vec![])))
            .collect();
        ProverMultiClaim::new(point, current_openings, next_openings)
    }

    #[test]
    fn new_starts_empty_and_records_shape() {
        // Invariant:
        //     new allocates one empty claim bucket per table and stores the shape verbatim.
        //
        // Fixture state:
        //     two tables (arity 3, arity 5), stacked arity 6, folding 2.
        let sc = StackedClaims::<F, EF>::new(
            vec![table(3), table(5)],
            vec![placement(1), placement(0)],
            6,
            2,
        );

        // The scalar shape fields come back exactly as supplied.
        assert_eq!(sc.folding(), 2);
        assert_eq!(sc.num_variables(), 6);

        // Per-table arities are read straight from the source tables.
        assert_eq!(sc.num_variables_table(0), 3);
        assert_eq!(sc.num_variables_table(1), 5);

        // No openings recorded yet.
        assert_eq!(sc.num_claims(), 0);
        assert!(sc.concrete_claims().next().is_none());

        // One bucket per source table, and no virtual claims.
        assert_eq!(sc.claim_map.len(), 2);
        assert!(sc.virtual_claims.is_empty());
    }

    #[test]
    fn num_claims_counts_openings_in_both_groups() {
        // Invariant:
        //     num_claims sums the openings of both groups across every recorded claim.
        //
        // Fixture state:
        //     one table; one claim of 2 current + 1 next, then a claim of 1 current.
        let mut sc = StackedClaims::<F, EF>::new(vec![table(2)], vec![placement(0)], 2, 0);
        sc.claim_map[0].push(claim(&[3, 5], &[7]));
        sc.claim_map[0].push(claim(&[11], &[]));

        // 3 openings + 1 opening = 4 consumed challenge powers.
        assert_eq!(sc.num_claims(), 4);
    }

    #[test]
    fn concrete_claims_walks_in_placement_order() {
        // Invariant:
        //     concrete_claims yields claims in placement order, not source-table order.
        //
        // Fixture state:
        //     placements list table 1 before table 0 (largest-first layout).
        //     bucket 0 holds a claim with current eval 100.
        //     bucket 1 holds a claim with current eval 200.
        let mut sc = StackedClaims::<F, EF>::new(
            vec![table(2), table(2)],
            vec![placement(1), placement(0)],
            3,
            0,
        );
        sc.claim_map[0].push(claim(&[100], &[]));
        sc.claim_map[1].push(claim(&[200], &[]));

        // Placement order is [table 1, table 0], so bucket 1's claim comes first.
        let evals: Vec<EF> = sc
            .concrete_claims()
            .map(|c| c.current_openings()[0].eval())
            .collect();
        assert_eq!(evals, vec![EF::from_u64(200), EF::from_u64(100)]);
    }

    #[test]
    fn sum_weights_openings_by_alpha_powers() {
        // Invariant:
        //     sum = Σ_i alpha^i * eval_i, walked as
        //     placements -> claims -> current openings -> next openings,
        //     with virtual claims continuing at alpha^(concrete opening count).
        //
        // Fixture state (alpha = 2):
        //     claim A: current [3, 5], next [7]  -> powers 0, 1, 2
        //     claim B: current [11]              -> power 3
        //     virtual: eval 13                   -> power 4 = num_claims
        //
        //     3*1 + 5*2 + 7*4 + 11*8 + 13*16 = 337
        let mut sc = StackedClaims::<F, EF>::new(vec![table(2)], vec![placement(0)], 2, 0);
        sc.claim_map[0].push(claim(&[3, 5], &[7]));
        sc.claim_map[0].push(claim(&[11], &[]));
        sc.virtual_claims.push(Claim {
            point: Point::new(vec![]),
            eval: EF::from_u64(13),
            data: Vec::new(),
        });

        // Concrete openings fix the virtual claim's starting power at alpha^4.
        assert_eq!(sc.num_claims(), 4);
        assert_eq!(sc.sum(EF::from_u64(2)), EF::from_u64(337));
    }
}