Better Sum-check Docstrings For spartan (Stages 1-3)

jolt-core/src/zkvm/claim_reductions/instruction_lookups.rs

@@ -1,3 +1,75 @@
+//! # InstructionClaimReduction (Stage 2)
+//!
+//! Source: `jolt-core/src/zkvm/claim_reductions/instruction_lookups.rs`
+//!
+//!
+//! ## Schwartz–Zippel randomness
+//!
+//! - Re-uses `r^(1)_cycle ∈ F^{log₂ T}` from Stage 1 (SpartanOuter)
+//!
+//!
+//! ## Batching randomness
+//!
+//! - `γ ∈ F`: fresh batching challenge (powers γ, γ², γ³, γ⁴ used)
+//!
+//!
+//! ## Sumcheck
+//!
+//! `eq(r^(1)_cycle, X_j)` is the [multilinear Lagrange basis polynomial][ml].
+//!
+//! [ml]: https://en.wikipedia.org/wiki/Multilinear_polynomial
+//!
+//! ```text
+//! LHS := Σ_{X_j}  eq(r^(1)_cycle, X_j)
+//!                · (LookupOutput(X_j)
+//!                  + γ   · LeftLookupOperand(X_j)
+//!                  + γ²  · RightLookupOperand(X_j)
+//!                  + γ³  · LeftInstructionInput(X_j)
+//!                  + γ⁴  · RightInstructionInput(X_j))
+//!
+//! RHS := LookupOutput(r^(1)_cycle)
+//!      + γ   · LeftLookupOperand(r^(1)_cycle)
+//!      + γ²  · RightLookupOperand(r^(1)_cycle)
+//!      + γ³  · LeftInstructionInput(r^(1)_cycle)
+//!      + γ⁴  · RightInstructionInput(r^(1)_cycle)
+//!
+//! where  X_j ∈ {0,1}^{log₂ T}
+//! ```
+//!
+//! Dimensions: `log₂ T` rounds (cycle only).
+//!
+//! The RHS is known: all five polynomials were opened at
+//! `r^(1)_cycle` in Stage 1.
+//!
+//!
+//! ## Opening point
+//!
+//! After sumcheck: `r^(2)_cycle ∈ F^{log₂ T}`.
+//!
+//!
+//! ## Verifier opening claim
+//!
+//! The verifier checks that the final sumcheck message equals:
+//!
+//! ```text
+//! eq(r^(1)_cycle, r^(2)_cycle)
+//!   · (LookupOutput(r^(2)_cycle)
+//!     + γ   · LeftLookupOperand(r^(2)_cycle)
+//!     + γ²  · RightLookupOperand(r^(2)_cycle)
+//!     + γ³  · LeftInstructionInput(r^(2)_cycle)
+//!     + γ⁴  · RightInstructionInput(r^(2)_cycle))
+//! ```
+//!
+//! The `eq` term is computable by the verifier. The prover supplies
+//! openings for all five VPs below.
+//!
+//!
+//! ## VirtualPolynomials opened at `r^(2)_cycle`
+//!
+//! ```text
+//! LookupOutput, LeftLookupOperand, RightLookupOperand,
+//! LeftInstructionInput, RightInstructionInput
+//! ```
 use std::array;
 use std::sync::Arc;
 

jolt-core/src/zkvm/claim_reductions/registers.rs

@@ -1,3 +1,67 @@
+//! # RegistersClaimReduction (Stage 3)
+//!
+//! Source: `jolt-core/src/zkvm/claim_reductions/registers.rs`
+//!
+//!
+//! ## Schwartz–Zippel randomness
+//!
+//! - Re-uses `r^(1)_cycle ∈ F^{log₂ T}` from Stage 1 (SpartanOuter)
+//!
+//!
+//! ## Batching randomness
+//!
+//! - `γ ∈ F`: fresh batching challenge (powers γ, γ² used)
+//!
+//!
+//! ## Sumcheck
+//!
+//! ```text
+//! LHS := Σ_{X_j}  eq(r^(1)_cycle, X_j)
+//!                · (RdWriteValue(X_j)
+//!                  + γ   · Rs1Value(X_j)
+//!                  + γ²  · Rs2Value(X_j))
+//!
+//! RHS := RdWriteValue(r^(1)_cycle)
+//!      + γ   · Rs1Value(r^(1)_cycle)
+//!      + γ²  · Rs2Value(r^(1)_cycle)
+//!
+//! where  X_j ∈ {0,1}^{log₂ T}
+//! ```
+//!
+//! Dimensions: `log₂ T` rounds (cycle only).
+//!
+//! The RHS is known: `RdWriteValue`, `Rs1Value`, and `Rs2Value`
+//! were opened at `r^(1)_cycle` in Stage 1 (SpartanOuter).
+//!
+//!
+//! ## Opening point
+//!
+//! After sumcheck: `r^(3)_cycle ∈ F^{log₂ T}` (shared with Shift
+//! and InstructionInput via Stage 3 batching).
+//!
+//!
+//! ## Verifier opening claim
+//!
+//! The verifier checks that the final sumcheck message equals:
+//!
+//! ```text
+//! eq(r^(1)_cycle, r^(3)_cycle)
+//!   · (RdWriteValue(r^(3)_cycle)
+//!     + γ   · Rs1Value(r^(3)_cycle)
+//!     + γ²  · Rs2Value(r^(3)_cycle))
+//! ```
+//!
+//! `eq` is computable by the verifier.
+//! The prover supplies openings for the 3 VPs below.
+//!
+//!
+//! ## VirtualPolynomials opened at `r^(3)_cycle`
+//!
+//! ```text
+//! RdWriteValue,
+//! Rs1Value,
+//! Rs2Value
+//! ```
 use std::array;
 use std::sync::Arc;
 

jolt-core/src/zkvm/ram/output_check.rs

@@ -1,3 +1,62 @@
+//! # OutputCheck (Stage 2)
+//!
+//! Source: `jolt-core/src/zkvm/ram/output_check.rs`
+//!
+//!
+//! ## Schwartz–Zippel randomness
+//!
+//! - `r_address ∈ F^{log₂ K_RAM}`: fresh address challenge vector
+//!
+//!
+//! ## Sumcheck
+//!
+//! `eq(r_address, X_k)` is the [multilinear Lagrange basis polynomial][ml].
+//!
+//! [ml]: https://en.wikipedia.org/wiki/Multilinear_polynomial
+//!
+//! ```text
+//! LHS := Σ_{X_k}  eq(r_address, X_k) · io_mask(X_k)
+//!                · (RamValFinal(X_k) - ValIO(X_k))
+//!
+//! RHS := 0   (zero-check)
+//!
+//! where  X_k ∈ {0,1}^{log₂ K_RAM}
+//! ```
+//!
+//! Dimensions: `log₂ K_RAM` rounds (address only).
+//!
+//! - `io_mask(X_k)`: 1 if address `X_k` is in the I/O region,
+//!   0 otherwise. Verifier-computable.
+//! - `ValIO(X_k)`: publicly claimed output value at address `X_k`.
+//!   Verifier-computable (public input).
+//! - `RamValFinal(X_k)`: final RAM state at address `X_k` after
+//!   all cycles (virtual).
+//!
+//!
+//! ## Opening point
+//!
+//! After sumcheck: `r^(2)_{K_RAM} ∈ F^{log₂ K_RAM}`.
+//!
+//!
+//! ## Verifier opening claim
+//!
+//! The verifier checks that the final sumcheck message equals:
+//!
+//! ```text
+//! eq(r_address, r^(2)_{K_RAM}) · io_mask(r^(2)_{K_RAM})
+//!   · (RamValFinal(r^(2)_{K_RAM}) - ValIO(r^(2)_{K_RAM}))
+//! ```
+//!
+//! `eq`, `io_mask`, and `ValIO` are computable by the verifier.
+//! The prover supplies the opening for `RamValFinal`.
+//!
+//!
+//! ## VirtualPolynomials opened at `r^(2)_{K_RAM}`
+//!
+//! ```text
+//! RamValFinal
+//! ```
+//!
 #[cfg(feature = "zk")]
 use crate::poly::opening_proof::OpeningId;
 #[cfg(feature = "zk")]
@@ -32,17 +91,6 @@ use common::{constants::RAM_START_ADDRESS, jolt_device::MemoryLayout};
 use rayon::prelude::*;
 use tracer::JoltDevice;
 
-// RAM output sumchecks
-//
-// OutputSumcheck:
-//   Proves the zero-check
-//     Σ_k eq(r_address, k) ⋅ io_mask(k) ⋅ (Val_final(k) − Val_io(k)) = 0,
-//   where:
-//   - r_address is a random address challenge vector.
-//   - io_mask is the MLE of the I/O-region indicator (1 on matching {0,1}-points).
-//   - Val_final(k) is the final memory value at address k.
-//   - Val_io(k) is the publicly claimed output value at address k.
-
 /// Degree bonud of the sumcheck round polynomials in [`OutputSumcheckVerifier`].
 const OUTPUT_SUMCHECK_DEGREE_BOUND: usize = 3;
 

jolt-core/src/zkvm/ram/raf_evaluation.rs

@@ -1,3 +1,64 @@
+//! # RafEvaluation (Stage 2)
+//!
+//! Source: `jolt-core/src/zkvm/ram/raf_evaluation.rs`
+//!
+//!
+//! ## Schwartz–Zippel randomness
+//!
+//! - Re-uses `r^(1)_cycle ∈ F^{log₂ T}` from Stage 1 (SpartanOuter)
+//!
+//!
+//! ## Sumcheck
+//!
+//! `unmap(X_k)` is the MLE of the identity function over `{0,1}^{log₂ K_RAM}`:
+//! it maps a boolean hypercube point to the integer it encodes.
+//! Definition: `jolt-core/src/poly/identity_poly.rs` (`UnmapRamAddressPolynomial`).
+//!
+//! ```text
+//! unmap(X_k) = Σ_{i=0}^{log₂ K_RAM - 1} 2^i · X_{k,i}
+//! ```
+//!
+//! The prover pre-computes
+//! `ra(X_k) := Σ_j eq(r^(1)_cycle, j) · 1[address(j) = k]`,
+//! aggregating access counts per remapped address `k`.
+//!
+//! ```text
+//! LHS := Σ_{X_k}  ra(X_k) · unmap(X_k)
+//!
+//! RHS := RamAddress(r^(1)_cycle)
+//!
+//! where  X_k ∈ {0,1}^{log₂ K_RAM}
+//! ```
+//!
+//! Dimensions: `log₂ K_RAM` rounds (address only).
+//!
+//! - `ra(X_k)`: `RamRa` partially evaluated at `r^(1)_cycle` (virtual).
+//! - `unmap(X_k)`: identity polynomial. Verifier-computable.
+//! - `RamAddress(r^(1)_cycle)`: opened in Stage 1 (SpartanOuter).
+//!
+//!
+//! ## Opening point
+//!
+//! After sumcheck: `r^(2)_{K_RAM} ∈ F^{log₂ K_RAM}`.
+//!
+//!
+//! ## Verifier opening claim
+//!
+//! The verifier checks that the final sumcheck message equals:
+//!
+//! ```text
+//! unmap(r^(2)_{K_RAM}) · RamRa(r^(2)_{K_RAM}, r^(1)_cycle)
+//! ```
+//!
+//! `unmap` is computable by the verifier.
+//! The prover supplies the opening for `RamRa`.
+//!
+//!
+//! ## VirtualPolynomials opened at `(r^(2)_{K_RAM}, r^(1)_cycle)`
+//!
+//! ```text
+//! RamRa
+//! ```
 use common::jolt_device::MemoryLayout;
 use num_traits::Zero;
 

jolt-core/src/zkvm/ram/read_write_checking.rs

@@ -1,3 +1,73 @@
+//! # RamReadWriteChecking (Stage 2)
+//!
+//! Source: `jolt-core/src/zkvm/ram/read_write_checking.rs`
+//!
+//!
+//! - `γ ∈ F`: Batching randomness
+//!
+//! ## Schwartz–Zippel randomness
+//!
+//! - Re-uses `r^(1)_cycle ∈ F^{log₂ T}` from Stage 1 (SpartanOuter)
+//!
+//!
+//! ## Sumcheck
+//!
+//! `eq(r^(1)_cycle, X_j)` is the [multilinear Lagrange basis polynomial][ml].
+//!
+//! [ml]: https://en.wikipedia.org/wiki/Multilinear_polynomial
+//!
+//! ```text
+//! LHS := Σ_{X_k, X_j}  eq(r^(1)_cycle, X_j) · RamRa(X_k, X_j)
+//!                      · (RamVal(X_k, X_j)
+//!                        + γ · (RamVal(X_k, X_j) + RamInc(X_j)))
+//!
+//! RHS := RamReadValue(r^(1)_cycle) + γ · RamWriteValue(r^(1)_cycle)
+//!
+//! where  X_k ∈ {0,1}^{log₂ K_RAM},  X_j ∈ {0,1}^{log₂ T}
+//! ```
+//!
+//! Dimensions: `log₂ K_RAM + log₂ T` rounds (address + cycle).
+//!
+//! The RHS is known: `RamReadValue` and `RamWriteValue` were opened
+//! at `r^(1)_cycle` in Stage 1.
+//!
+//! - `RamRa(X_k, X_j)`: 1 if address `X_k` is accessed at cycle `X_j` (virtual)
+//! - `RamVal(X_k, X_j)`: value at address `X_k` right before cycle `X_j` (virtual)
+//! - `RamInc(X_j)`: write increment at cycle `X_j` (COMMITTED)
+//!
+//!
+//! ## Opening point
+//!
+//! After sumcheck: `(r^(2)_{K_RAM}, r^(2)_cycle) ∈ F^{log₂ K_RAM + log₂ T}`.
+//! Both components are used downstream.
+//!
+//!
+//! ## Verifier opening claim
+//!
+//! The verifier checks that the final sumcheck message equals:
+//!
+//! ```text
+//! eq(r^(1)_cycle, r^(2)_cycle) · RamRa(r^(2)_{K_RAM}, r^(2)_cycle)
+//!   · (RamVal(r^(2)_{K_RAM}, r^(2)_cycle)
+//!     + γ · (RamVal(r^(2)_{K_RAM}, r^(2)_cycle) + RamInc(r^(2)_cycle)))
+//! ```
+//!
+//! The `eq` term is computable by the verifier. The prover supplies
+//! openings for `RamRa`, `RamVal`, and `RamInc`.
+//!
+//!
+//! ## Polynomials opened
+//!
+//! Virtual:
+//! ```text
+//! RamVal  at (r^(2)_{K_RAM}, r^(2)_cycle)
+//! RamRa   at (r^(2)_{K_RAM}, r^(2)_cycle)
+//! ```
+//!
+//! Committed:
+//! ```text
+//! RamInc  at r^(2)_cycle
+//! ```
 use common::jolt_device::MemoryLayout;
 use num::Integer;
 use num_traits::Zero;
@@ -45,21 +115,6 @@ use allocative::FlameGraphBuilder;
 use rayon::prelude::*;
 use tracer::instruction::Cycle;
 
-// RAM read-write checking sumcheck
-//
-// Proves the relation:
-//   Σ_{k,j} eq(r_cycle, j) ⋅ ra(k, j) ⋅ (Val(k, j) + γ ⋅ (inc(j) + Val(k, j)))
-//   = rv_claim + γ ⋅ wv_claim
-// where:
-// - r_cycle are the challenges for the cycle variables in this sumcheck (from Spartan outer)
-// - ra(k, j) = 1 if memory address k is accessed at cycle j, and 0 otherwise
-// - Val(k, j) is the value at memory address k right before cycle j
-// - inc(j) is the change in value at cycle j if a write occurs, and 0 otherwise
-// - rv_claim and wv_claim are the claimed read and write values from the Spartan outer sumcheck.
-//
-// This sumcheck ensures that the values read from and written to RAM are consistent
-// with the memory trace and the initial/final memory states.
-
 /// Degree bound of the sumcheck round polynomials in [`RamReadWriteCheckingVerifier`].
 const DEGREE_BOUND: usize = 3;