stage 5

moodlezoup · moodlezoup · commit 56d26b8947cd · 2025-11-21T19:12:45.000-05:00
diff --git a/jolt-core/src/zkvm/prover.rs b/jolt-core/src/zkvm/prover.rs
@@ -19,10 +19,8 @@ use crate::{
             commitment_scheme::StreamingCommitmentScheme,
             dory::{DoryContext, DoryGlobals},
         },
-        multilinear_polynomial::{MultilinearPolynomial, PolynomialEvaluation},
-        opening_proof::{
-            OpeningAccumulator, ProverOpeningAccumulator, ReducedOpeningProof, SumcheckId,
-        },
+        multilinear_polynomial::MultilinearPolynomial,
+        opening_proof::{ProverOpeningAccumulator, ReducedOpeningProof},
         rlc_polynomial::RLCStreamingData,
     },
     pprof_scope,
@@ -35,6 +33,7 @@ use crate::{
     zkvm::{
         config::{get_log_k_chunk, OneHotParams},
         ram::{
+            hamming_booleanity::HammingBooleanityParams,
             output_check::OutputSumcheckParams,
             populate_memory_states,
             raf_evaluation::RafEvaluationSumcheckParams,
@@ -45,7 +44,10 @@ use crate::{
             },
             val_final::{ValFinalSumcheckParams, ValFinalSumcheckProver},
         },
-        registers::read_write_checking::RegistersReadWriteCheckingParams,
+        registers::{
+            read_write_checking::RegistersReadWriteCheckingParams,
+            val_evaluation::RegistersValEvaluationSumcheckParams,
+        },
         spartan::{instruction_input::InstructionInputParams, shift::ShiftSumcheckParams},
         verifier::JoltVerifierPreprocessing,
     },
@@ -80,7 +82,7 @@ use crate::{
             product::ProductVirtualRemainderProver, prove_stage1_uni_skip, prove_stage2_uni_skip,
             shift::ShiftSumcheckProver,
         },
-        witness::{AllCommittedPolynomials, CommittedPolynomial, VirtualPolynomial},
+        witness::{AllCommittedPolynomials, CommittedPolynomial},
         ProverDebugInfo, Serializable,
     },
 };
@@ -759,15 +761,19 @@ impl<'a, F: JoltField, PCS: StreamingCommitmentScheme<Field = F>, ProofTranscrip
     fn prove_stage5(&mut self) -> SumcheckInstanceProof<F, ProofTranscript> {
         #[cfg(not(target_arch = "wasm32"))]
         print_current_memory_usage("Stage 5 baseline");
+        let registers_val_evaluation_params =
+            RegistersValEvaluationSumcheckParams::new(&self.opening_accumulator);
+        let ram_hamming_booleanity_params = HammingBooleanityParams::new(&self.opening_accumulator);
 
-        let registers_val_evaluation = RegistersValEvaluationSumcheckProver::gen(
+        let registers_val_evaluation = RegistersValEvaluationSumcheckProver::initialize(
+            registers_val_evaluation_params,
             &self.trace,
             &self.preprocessing.bytecode,
             &self.program_io.memory_layout,
-            &self.opening_accumulator,
         );
         let ram_hamming_booleanity =
-            HammingBooleanitySumcheckProver::gen(&self.trace, &self.opening_accumulator);
+            HammingBooleanitySumcheckProver::initialize(ram_hamming_booleanity_params, &self.trace);
+
         let ram_ra_virtual = RamRaSumcheckProver::gen(
             &self.trace,
             &self.program_io.memory_layout,
diff --git a/jolt-core/src/zkvm/ram/hamming_booleanity.rs b/jolt-core/src/zkvm/ram/hamming_booleanity.rs
@@ -8,15 +8,13 @@ use crate::poly::opening_proof::{
 use crate::poly::split_eq_poly::GruenSplitEqPolynomial;
 use crate::poly::unipoly::UniPoly;
 use crate::subprotocols::sumcheck_prover::SumcheckInstanceProver;
-use crate::subprotocols::sumcheck_verifier::SumcheckInstanceVerifier;
+use crate::subprotocols::sumcheck_verifier::{SumcheckInstanceParams, SumcheckInstanceVerifier};
 use crate::transcripts::Transcript;
-use crate::utils::math::Math;
 use crate::zkvm::witness::VirtualPolynomial;
 use allocative::Allocative;
 #[cfg(feature = "allocative")]
 use allocative::FlameGraphBuilder;
 use rayon::prelude::*;
-use std::marker::PhantomData;
 use tracer::instruction::Cycle;
 
 // RAM Hamming booleanity sumcheck
@@ -30,53 +28,74 @@ use tracer::instruction::Cycle;
 /// Degree bound of the sumcheck round polynomials in [`HammingBooleanitySumcheckVerifier`].
 const DEGREE_BOUND: usize = 3;
 
+#[derive(Allocative)]
+pub struct HammingBooleanityParams<F: JoltField> {
+    r_cycle: OpeningPoint<BIG_ENDIAN, F>,
+}
+
+impl<F: JoltField> HammingBooleanityParams<F> {
+    pub fn new(opening_accumulator: &dyn OpeningAccumulator<F>) -> Self {
+        let (r_cycle, _) = opening_accumulator.get_virtual_polynomial_opening(
+            VirtualPolynomial::LookupOutput,
+            SumcheckId::SpartanOuter,
+        );
+
+        Self { r_cycle }
+    }
+}
+
+impl<F: JoltField> SumcheckInstanceParams<F> for HammingBooleanityParams<F> {
+    fn degree(&self) -> usize {
+        DEGREE_BOUND
+    }
+
+    fn num_rounds(&self) -> usize {
+        self.r_cycle.len()
+    }
+
+    fn input_claim(&self, _: &dyn OpeningAccumulator<F>) -> F {
+        F::zero()
+    }
+
+    fn normalize_opening_point(
+        &self,
+        challenges: &[<F as JoltField>::Challenge],
+    ) -> OpeningPoint<BIG_ENDIAN, F> {
+        OpeningPoint::<LITTLE_ENDIAN, F>::new(challenges.to_vec()).match_endianness()
+    }
+}
+
 #[derive(Allocative)]
 pub struct HammingBooleanitySumcheckProver<F: JoltField> {
     eq_r_cycle: GruenSplitEqPolynomial<F>,
     H: MultilinearPolynomial<F>,
-    log_T: usize,
+    params: HammingBooleanityParams<F>,
 }
 
 impl<F: JoltField> HammingBooleanitySumcheckProver<F> {
-    #[tracing::instrument(skip_all, name = "RamHammingBooleanitySumcheckProver::gen")]
-    pub fn gen(trace: &[Cycle], opening_accumulator: &ProverOpeningAccumulator<F>) -> Self {
-        let T = trace.len();
-        let log_T = T.log_2();
-
+    #[tracing::instrument(skip_all, name = "RamHammingBooleanitySumcheckProver::initialize")]
+    pub fn initialize(params: HammingBooleanityParams<F>, trace: &[Cycle]) -> Self {
         let H = trace
             .par_iter()
             .map(|cycle| cycle.ram_access().address() != 0)
             .collect::<Vec<bool>>();
         let H = MultilinearPolynomial::from(H);
 
-        let (r_cycle, _) = opening_accumulator.get_virtual_polynomial_opening(
-            VirtualPolynomial::LookupOutput,
-            SumcheckId::SpartanOuter,
-        );
-
-        let eq_r_cycle = GruenSplitEqPolynomial::new(&r_cycle.r, BindingOrder::LowToHigh);
+        let eq_r_cycle = GruenSplitEqPolynomial::new(&params.r_cycle.r, BindingOrder::LowToHigh);
 
         Self {
             eq_r_cycle,
             H,
-            log_T,
+            params,
         }
     }
 }
 
 impl<F: JoltField, T: Transcript> SumcheckInstanceProver<F, T>
     for HammingBooleanitySumcheckProver<F>
 {
-    fn degree(&self) -> usize {
-        DEGREE_BOUND
-    }
-
-    fn num_rounds(&self) -> usize {
-        self.log_T
-    }
-
-    fn input_claim(&self, _accumulator: &ProverOpeningAccumulator<F>) -> F {
-        F::zero()
+    fn get_params(&self) -> Box<&dyn SumcheckInstanceParams<F>> {
+        Box::new(&self.params)
     }
 
     #[tracing::instrument(skip_all, name = "RamHammingBooleanitySumcheckProver::compute_message")]
@@ -113,7 +132,7 @@ impl<F: JoltField, T: Transcript> SumcheckInstanceProver<F, T>
             transcript,
             VirtualPolynomial::RamHammingWeight,
             SumcheckId::RamHammingBooleanity,
-            get_opening_point(sumcheck_challenges),
+            self.params.normalize_opening_point(sumcheck_challenges),
             self.H.final_sumcheck_claim(),
         );
     }
@@ -125,33 +144,22 @@ impl<F: JoltField, T: Transcript> SumcheckInstanceProver<F, T>
 }
 
 pub struct HammingBooleanitySumcheckVerifier<F: JoltField> {
-    log_T: usize,
-    _phantom: PhantomData<F>,
+    params: HammingBooleanityParams<F>,
 }
 
 impl<F: JoltField> HammingBooleanitySumcheckVerifier<F> {
-    pub fn new(n_cycle_vars: usize) -> Self {
+    pub fn new(opening_accumulator: &dyn OpeningAccumulator<F>) -> Self {
         Self {
-            // TODO: Make the name for this consistent across the codebase.
-            log_T: n_cycle_vars,
-            _phantom: PhantomData,
+            params: HammingBooleanityParams::new(opening_accumulator),
         }
     }
 }
 
 impl<F: JoltField, T: Transcript> SumcheckInstanceVerifier<F, T>
     for HammingBooleanitySumcheckVerifier<F>
 {
-    fn degree(&self) -> usize {
-        DEGREE_BOUND
-    }
-
-    fn num_rounds(&self) -> usize {
-        self.log_T
-    }
-
-    fn input_claim(&self, _accumulator: &VerifierOpeningAccumulator<F>) -> F {
-        F::zero()
+    fn get_params(&self) -> Box<&dyn SumcheckInstanceParams<F>> {
+        Box::new(&self.params)
     }
 
     fn expected_output_claim(
@@ -194,13 +202,7 @@ impl<F: JoltField, T: Transcript> SumcheckInstanceVerifier<F, T>
             transcript,
             VirtualPolynomial::RamHammingWeight,
             SumcheckId::RamHammingBooleanity,
-            get_opening_point(sumcheck_challenges),
+            self.params.normalize_opening_point(sumcheck_challenges),
         );
     }
 }
-
-fn get_opening_point<F: JoltField>(
-    sumcheck_challenges: &[F::Challenge],
-) -> OpeningPoint<BIG_ENDIAN, F> {
-    OpeningPoint::<LITTLE_ENDIAN, F>::new(sumcheck_challenges.to_vec()).match_endianness()
-}
diff --git a/jolt-core/src/zkvm/ram/ra_virtual.rs b/jolt-core/src/zkvm/ram/ra_virtual.rs
@@ -49,7 +49,6 @@ pub struct RaSumcheckProver<F: JoltField> {
     ra_i_polys: Vec<RaPolynomial<u16, F>>,
     /// eq poly
     eq_poly: MultilinearPolynomial<F>,
-    #[allocative(skip)]
     params: RaSumcheckParams<F>,
 }
 
@@ -271,6 +270,7 @@ impl<F: JoltField, T: Transcript> SumcheckInstanceVerifier<F, T> for RaSumcheckV
     }
 }
 
+#[derive(Allocative)]
 struct RaSumcheckParams<F: JoltField> {
     gamma_powers: [F; 3],
     /// Random challenge r_cycle
diff --git a/jolt-core/src/zkvm/registers/val_evaluation.rs b/jolt-core/src/zkvm/registers/val_evaluation.rs
@@ -20,7 +20,6 @@ use crate::{
         sumcheck_verifier::{SumcheckInstanceParams, SumcheckInstanceVerifier},
     },
     transcripts::Transcript,
-    utils::math::Math,
     zkvm::{
         bytecode::BytecodePreprocessing,
         witness::{CommittedPolynomial, VirtualPolynomial},
@@ -52,23 +51,30 @@ const LOG_K: usize = REGISTER_COUNT.ilog2() as usize;
 const DEGREE_BOUND: usize = 3;
 
 #[derive(Allocative)]
-struct ValEvaluationSumcheckParams {
-    n_cycle_vars: usize,
+pub struct RegistersValEvaluationSumcheckParams<F: JoltField> {
+    r_address: OpeningPoint<BIG_ENDIAN, F>,
+    r_cycle: OpeningPoint<BIG_ENDIAN, F>,
 }
 
-impl ValEvaluationSumcheckParams {
-    pub fn new(n_cycle_vars: usize) -> Self {
-        Self { n_cycle_vars }
+impl<F: JoltField> RegistersValEvaluationSumcheckParams<F> {
+    pub fn new(opening_accumulator: &dyn OpeningAccumulator<F>) -> Self {
+        // The opening point is r_address || r_cycle
+        let (r, _) = opening_accumulator.get_virtual_polynomial_opening(
+            VirtualPolynomial::RegistersVal,
+            SumcheckId::RegistersReadWriteChecking,
+        );
+        let (r_address, r_cycle) = r.split_at(LOG_K);
+        Self { r_address, r_cycle }
     }
 }
 
-impl<F: JoltField> SumcheckInstanceParams<F> for ValEvaluationSumcheckParams {
+impl<F: JoltField> SumcheckInstanceParams<F> for RegistersValEvaluationSumcheckParams<F> {
     fn degree(&self) -> usize {
         DEGREE_BOUND
     }
 
     fn num_rounds(&self) -> usize {
-        self.n_cycle_vars
+        self.r_cycle.len()
     }
 
     fn input_claim(&self, accumulator: &dyn OpeningAccumulator<F>) -> F {
@@ -92,33 +98,25 @@ pub(crate) struct ValEvaluationSumcheckProver<F: JoltField> {
     inc: MultilinearPolynomial<F>,
     wa: RaPolynomial<u8, F>,
     lt: LtPolynomial<F>,
-    params: ValEvaluationSumcheckParams,
+    params: RegistersValEvaluationSumcheckParams<F>,
 }
 
 impl<F: JoltField> ValEvaluationSumcheckProver<F> {
-    #[tracing::instrument(skip_all, name = "RegistersValEvaluationSumcheckProver::gen")]
-    pub fn gen(
+    #[tracing::instrument(skip_all, name = "RegistersValEvaluationSumcheckProver::initialize")]
+    pub fn initialize(
+        params: RegistersValEvaluationSumcheckParams<F>,
         trace: &[Cycle],
         bytecode_preprocessing: &BytecodePreprocessing,
         memory_layout: &MemoryLayout,
-        opening_accumulator: &ProverOpeningAccumulator<F>,
     ) -> Self {
-        // The opening point is r_address || r_cycle
-        let registers_val_input_sample = opening_accumulator.get_virtual_polynomial_opening(
-            VirtualPolynomial::RegistersVal,
-            SumcheckId::RegistersReadWriteChecking,
-        );
-        let (r_address, r_cycle) = registers_val_input_sample.0.split_at(LOG_K);
-
-        let params = ValEvaluationSumcheckParams::new(trace.len().log_2());
         let inc = CommittedPolynomial::RdInc.generate_witness(
             bytecode_preprocessing,
             memory_layout,
             trace,
             None,
         );
 
-        let eq_r_address = EqPolynomial::evals(&r_address.r);
+        let eq_r_address = EqPolynomial::evals(&params.r_address.r);
         let wa: Vec<Option<u8>> = trace
             .par_iter()
             .map(|cycle| {
@@ -127,7 +125,7 @@ impl<F: JoltField> ValEvaluationSumcheckProver<F> {
             })
             .collect();
         let wa = RaPolynomial::new(Arc::new(wa), eq_r_address);
-        let lt = LtPolynomial::new(&r_cycle);
+        let lt = LtPolynomial::new(&params.r_cycle);
 
         Self {
             inc,
@@ -231,18 +229,20 @@ impl<F: JoltField, T: Transcript> SumcheckInstanceProver<F, T> for ValEvaluation
     }
 }
 
-pub struct ValEvaluationSumcheckVerifier {
-    params: ValEvaluationSumcheckParams,
+pub struct ValEvaluationSumcheckVerifier<F: JoltField> {
+    params: RegistersValEvaluationSumcheckParams<F>,
 }
 
-impl ValEvaluationSumcheckVerifier {
-    pub fn new(n_cycle_vars: usize) -> Self {
-        let params = ValEvaluationSumcheckParams::new(n_cycle_vars);
+impl<F: JoltField> ValEvaluationSumcheckVerifier<F> {
+    pub fn new(opening_accumulator: &VerifierOpeningAccumulator<F>) -> Self {
+        let params = RegistersValEvaluationSumcheckParams::new(opening_accumulator);
         Self { params }
     }
 }
 
-impl<F: JoltField, T: Transcript> SumcheckInstanceVerifier<F, T> for ValEvaluationSumcheckVerifier {
+impl<F: JoltField, T: Transcript> SumcheckInstanceVerifier<F, T>
+    for ValEvaluationSumcheckVerifier<F>
+{
     fn get_params(&self) -> Box<&dyn SumcheckInstanceParams<F>> {
         Box::new(&self.params)
     }
diff --git a/jolt-core/src/zkvm/verifier.rs b/jolt-core/src/zkvm/verifier.rs
@@ -335,8 +335,10 @@ impl<'a, F: JoltField, PCS: CommitmentScheme<Field = F>, ProofTranscript: Transc
 
     fn verify_stage5(&mut self) -> Result<(), anyhow::Error> {
         let n_cycle_vars = self.proof.trace_length.log_2();
-        let registers_val_evaluation = RegistersValEvaluationSumcheckVerifier::new(n_cycle_vars);
-        let ram_hamming_booleanity = HammingBooleanitySumcheckVerifier::new(n_cycle_vars);
+        let registers_val_evaluation =
+            RegistersValEvaluationSumcheckVerifier::new(&self.opening_accumulator);
+        let ram_hamming_booleanity =
+            HammingBooleanitySumcheckVerifier::new(&self.opening_accumulator);
         let ram_ra_virtual = RamRaSumcheckVerifier::new(
             self.proof.trace_length,
             &self.one_hot_params,

Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,6 @@ pub struct RaSumcheckProver<F: JoltField> {`
`49`	`49`	`ra_i_polys: Vec<RaPolynomial<u16, F>>,`
`50`	`50`	`/// eq poly`
`51`	`51`	`eq_poly: MultilinearPolynomial<F>,`
`52`		`- #[allocative(skip)]`
`53`	`52`	`params: RaSumcheckParams<F>,`
`54`	`53`	`}`
`55`	`54`
`@@ -271,6 +270,7 @@ impl<F: JoltField, T: Transcript> SumcheckInstanceVerifier<F, T> for RaSumcheckV`
`271`	`270`	`}`
`272`	`271`	`}`
`273`	`272`
	`273`	`+#[derive(Allocative)]`
`274`	`274`	`struct RaSumcheckParams<F: JoltField> {`
`275`	`275`	`gamma_powers: [F; 3],`
`276`	`276`	`/// Random challenge r_cycle`