Skip to main content

p3_sumcheck/zk/prover/
layout.rs

1//! Zero-knowledge extension of the stacked sumcheck layout.
2//!
3//! Captures the data the masking layer needs from any layout, and the one
4//! piece of arithmetic that genuinely branches on the binding direction.
5
6use p3_field::{ExtensionField, Field, PackedValue, TwoAdicField};
7use p3_multilinear_util::point::Point;
8use p3_util::log2_strict_usize;
9
10use crate::layout::{Layout, PrefixProver, ProverMultiClaim, ProverVirtualClaim, SuffixProver};
11use crate::product_polynomial::ProductPolynomial;
12use crate::strategy::VariableOrder;
13
14/// Per-mode hooks consumed by the zero-knowledge prover.
15///
16/// Every implementor exposes the data the masking layer reads.
17/// The residual handoff is the only operation that branches on the binding direction.
18pub trait ZkLayout<F, EF>: Layout<F, EF>
19where
20    F: TwoAdicField,
21    EF: ExtensionField<F>,
22{
23    /// Walks concrete claims in placement order.
24    fn concrete_claims(&self) -> impl Iterator<Item = &ProverMultiClaim<F, EF>>;
25
26    /// Returns the virtual-claim slice.
27    fn virtual_claims(&self) -> &[ProverVirtualClaim<EF>];
28
29    /// Returns the alpha-batched plain sum.
30    fn batched_sum(&self, alpha: EF) -> EF;
31
32    /// Builds the residual product polynomial, scaled by the combining challenge.
33    ///
34    /// Consumes the layout; the residual factor is its last consumer.
35    fn zk_residual_handoff(self, rs: &Point<EF>, alpha: EF, eps: EF) -> ProductPolynomial<F, EF>
36    where
37        EF: TwoAdicField;
38}
39
40impl<F, EF> ZkLayout<F, EF> for PrefixProver<F, EF>
41where
42    F: TwoAdicField,
43    EF: ExtensionField<F>,
44{
45    fn concrete_claims(&self) -> impl Iterator<Item = &ProverMultiClaim<F, EF>> {
46        // Flatten placement order across all per-table claim lists.
47        self.placements
48            .iter()
49            .flat_map(|placement| self.claim_map[placement.idx()].iter())
50    }
51
52    fn virtual_claims(&self) -> &[ProverVirtualClaim<EF>] {
53        &self.virtual_claims
54    }
55
56    fn batched_sum(&self, alpha: EF) -> EF {
57        self.sum(alpha)
58    }
59
60    fn zk_residual_handoff(self, rs: &Point<EF>, alpha: EF, eps: EF) -> ProductPolynomial<F, EF>
61    where
62        EF: TwoAdicField,
63    {
64        // Prefix packs the residual weights: one full SIMD lane must survive the fold.
65        // `Poly::pack` requires `num_variables - k >= k_pack`, else it panics.
66        // Suffix is unpacked and unconstrained, so this guard is prefix-only.
67        // Phrased as `k + k_pack <= num_variables` to avoid `usize` underflow.
68        let k_pack = log2_strict_usize(<F as Field>::Packing::WIDTH);
69        assert!(
70            rs.num_variables() + k_pack <= self.num_variables(),
71            "prefix packed residual needs num_variables - folding >= k_pack",
72        );
73
74        // Fold the stacked polynomial low-to-high.
75        // The combining challenge is baked into the compression scale.
76        let compressed = tracing::info_span!("compress_prefix_to_packed")
77            .in_scope(|| self.poly.compress_prefix_to_packed(rs, eps));
78        // Pack the equality weights for the SIMD-friendly residual rounds.
79        let weights = self.combine_eqs(rs, alpha).pack::<F, EF>();
80        ProductPolynomial::new_packed(VariableOrder::Prefix, compressed, weights)
81    }
82}
83
84impl<F, EF> ZkLayout<F, EF> for SuffixProver<F, EF>
85where
86    F: TwoAdicField,
87    EF: ExtensionField<F>,
88{
89    fn concrete_claims(&self) -> impl Iterator<Item = &ProverMultiClaim<F, EF>> {
90        // Flatten placement order across all per-table claim lists.
91        self.placements
92            .iter()
93            .flat_map(|placement| self.claim_map[placement.idx()].iter())
94    }
95
96    fn virtual_claims(&self) -> &[ProverVirtualClaim<EF>] {
97        &self.virtual_claims
98    }
99
100    fn batched_sum(&self, alpha: EF) -> EF {
101        self.sum(alpha)
102    }
103
104    fn zk_residual_handoff(self, rs: &Point<EF>, alpha: EF, eps: EF) -> ProductPolynomial<F, EF>
105    where
106        EF: TwoAdicField,
107    {
108        // Reverse the challenges to match the suffix-binding frame.
109        let reversed = rs.reversed();
110        // Walk per-table slots; the combining challenge rides on the slot compression.
111        let compressed = tracing::info_span!("compress_stacked_with_eps")
112            .in_scope(|| self.compress_stacked_scaled(&reversed, eps));
113        // The SVO preprocessing covers what packing would help; no packing here.
114        let weights = self.combine_eqs(&reversed, alpha);
115        ProductPolynomial::new_unpacked(VariableOrder::Suffix, compressed, weights)
116    }
117}