Revert "refactor: remove proofs and commitments from state manager (#1065)"

This reverts commit 025c012.

Author: andrewmilson

jolt-core/benches/e2e_profiling.rs

@@ -1,5 +1,6 @@
 use ark_serialize::CanonicalSerialize;
 use jolt_core::host;
+use jolt_core::zkvm::dag::state_manager::ProofKeys;
 use jolt_core::zkvm::JoltVerifierPreprocessing;
 use jolt_core::zkvm::{Jolt, JoltRV64IMAC};
 use std::fs;
@@ -280,11 +281,9 @@ fn prove_example_with_trace(
     });
     let proof_size = jolt_proof.serialized_size(ark_serialize::Compress::Yes);
     let proof_size_full_compressed = proof_size
-        - jolt_proof
-            .reduced_opening_proof
+        - jolt_proof.proofs[&ProofKeys::ReducedOpeningProof]
             .serialized_size(ark_serialize::Compress::Yes)
-        + (jolt_proof
-            .reduced_opening_proof
+        + (jolt_proof.proofs[&ProofKeys::ReducedOpeningProof]
             .serialized_size(ark_serialize::Compress::No)
             / 3)
         - jolt_proof

jolt-core/src/zkvm/bytecode/mod.rs

@@ -145,7 +145,7 @@ impl<F: JoltField, PCS: CommitmentScheme<Field = F>, T: Transcript> SumcheckStag
 {
     fn stage6_instances(
         &mut self,
-        sm: &mut StateManager<'_, F, PCS>,
+        sm: &mut StateManager<'_, F, T, PCS>,
         opening_accumulator: &mut ProverOpeningAccumulator<F>,
         transcript: &mut T,
     ) -> Vec<Box<dyn SumcheckInstanceProver<F, T>>> {
@@ -175,7 +175,7 @@ impl<F: JoltField, PCS: CommitmentScheme<Field = F>, T: Transcript>
 {
     fn stage6_instances(
         &mut self,
-        sm: &mut StateManager<'_, F, PCS>,
+        sm: &mut StateManager<'_, F, T, PCS>,
         opening_accumulator: &mut VerifierOpeningAccumulator<F>,
         transcript: &mut T,
     ) -> Vec<Box<dyn SumcheckInstanceVerifier<F, T>>> {
@@ -189,10 +189,10 @@ impl<F: JoltField, PCS: CommitmentScheme<Field = F>, T: Transcript>
     }
 }
 
-fn gen_ra_one_hot_provers<F: JoltField>(
-    state_manager: &mut StateManager<'_, F, impl CommitmentScheme<Field = F>>,
+fn gen_ra_one_hot_provers<F: JoltField, T: Transcript>(
+    state_manager: &mut StateManager<'_, F, T, impl CommitmentScheme<Field = F>>,
     opening_accumulator: &ProverOpeningAccumulator<F>,
-    transcript: &mut impl Transcript,
+    transcript: &mut T,
 ) -> (HammingWeightSumcheckProver<F>, BooleanitySumcheckProver<F>) {
     let (preprocessing, _, trace, _, _) = state_manager.get_prover_data();
     let bytecode_preprocessing = &preprocessing.shared.bytecode;
@@ -247,9 +247,9 @@ fn gen_ra_one_hot_provers<F: JoltField>(
     )
 }
 
-fn new_ra_one_hot_verifiers<F: JoltField>(
-    state_manager: &mut StateManager<'_, F, impl CommitmentScheme<Field = F>>,
-    transcript: &mut impl Transcript,
+fn new_ra_one_hot_verifiers<F: JoltField, T: Transcript>(
+    state_manager: &mut StateManager<'_, F, T, impl CommitmentScheme<Field = F>>,
+    transcript: &mut T,
 ) -> (
     HammingWeightSumcheckVerifier<F>,
     BooleanitySumcheckVerifier<F>,

jolt-core/src/zkvm/bytecode/read_raf_checking.rs

@@ -122,10 +122,10 @@ pub struct ReadRafSumcheckProver<F: JoltField> {
 
 impl<F: JoltField> ReadRafSumcheckProver<F> {
     #[tracing::instrument(skip_all, name = "BytecodeReadRafSumcheckProver::gen")]
-    pub fn gen(
-        state_manager: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+    pub fn gen<T: Transcript>(
+        state_manager: &mut StateManager<F, T, impl CommitmentScheme<Field = F>>,
         opening_accumulator: &ProverOpeningAccumulator<F>,
-        transcript: &mut impl Transcript,
+        transcript: &mut T,
     ) -> Self {
         let params = ReadRafSumcheckParams::gen(state_manager, opening_accumulator, transcript);
 
@@ -520,10 +520,10 @@ pub struct ReadRafSumcheckVerifier<F: JoltField> {
 }
 
 impl<F: JoltField> ReadRafSumcheckVerifier<F> {
-    pub fn gen(
-        state_manager: &mut StateManager<'_, F, impl CommitmentScheme<Field = F>>,
+    pub fn gen<T: Transcript>(
+        state_manager: &mut StateManager<'_, F, T, impl CommitmentScheme<Field = F>>,
         opening_accumulator: &VerifierOpeningAccumulator<F>,
-        transcript: &mut impl Transcript,
+        transcript: &mut T,
     ) -> Self {
         Self {
             params: ReadRafSumcheckParams::gen(state_manager, opening_accumulator, transcript),
@@ -642,10 +642,10 @@ struct ReadRafSumcheckParams<F: JoltField> {
 }
 
 impl<F: JoltField> ReadRafSumcheckParams<F> {
-    fn gen(
-        state_manager: &mut StateManager<'_, F, impl CommitmentScheme<Field = F>>,
+    fn gen<T: Transcript>(
+        state_manager: &mut StateManager<'_, F, T, impl CommitmentScheme<Field = F>>,
         opening_accumulator: &dyn OpeningAccumulator<F>,
-        transcript: &mut impl Transcript,
+        transcript: &mut T,
     ) -> Self {
         let preprocessing = state_manager.get_shared_preprocessing();
         let K = preprocessing.bytecode.code_size;
@@ -757,11 +757,11 @@ impl<F: JoltField> ReadRafSumcheckParams<F> {
         }
     }
 
-    fn compute_val_rv(
-        sm: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+    fn compute_val_rv<T: Transcript>(
+        sm: &mut StateManager<F, T, impl CommitmentScheme<Field = F>>,
         opening_accumulator: &dyn OpeningAccumulator<F>,
         val_type: ReadCheckingValType,
-        transcript: &mut impl Transcript,
+        transcript: &mut T,
     ) -> (Vec<F>, F) {
         match val_type {
             ReadCheckingValType::Stage1 => {
@@ -807,7 +807,7 @@ impl<F: JoltField> ReadRafSumcheckParams<F> {
     ///             + gamma^2 * circuit_flags[0](k) + gamma^3 * circuit_flags[1](k) + ...
     /// This particular Val virtualizes claims output by Spartan's "outer" sumcheck
     fn compute_val_1(
-        sm: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+        sm: &mut StateManager<F, impl Transcript, impl CommitmentScheme<Field = F>>,
         gamma_powers: &[F],
     ) -> Vec<F> {
         sm.get_bytecode()
@@ -880,7 +880,7 @@ impl<F: JoltField> ReadRafSumcheckParams<F> {
     ///       write_lookup_output_to_rd_flag(k) = 1 if instruction k writes lookup output to rd, 0 otherwise.
     /// This Val matches the fused product sumcheck.
     fn compute_val_2(
-        sm: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+        sm: &mut StateManager<F, impl Transcript, impl CommitmentScheme<Field = F>>,
         gamma_powers: &[F],
     ) -> Vec<F> {
         sm.get_bytecode()
@@ -948,7 +948,7 @@ impl<F: JoltField> ReadRafSumcheckParams<F> {
     ///             + gamma^3 * left_operand_is_pc(k) + ...
     /// This particular Val virtualizes claims output by the ShiftSumcheck.
     fn compute_val_3(
-        sm: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+        sm: &mut StateManager<F, impl Transcript, impl CommitmentScheme<Field = F>>,
         gamma_powers: &[F],
     ) -> Vec<F> {
         sm.get_bytecode()
@@ -1067,7 +1067,7 @@ impl<F: JoltField> ReadRafSumcheckParams<F> {
     /// and analogously for rs1(k, k') and rs2(k, k').
     /// This particular Val virtualizes claims output by the registers read/write checking sumcheck.
     fn compute_val_4(
-        sm: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+        sm: &mut StateManager<F, impl Transcript, impl CommitmentScheme<Field = F>>,
         opening_accumulator: &dyn OpeningAccumulator<F>,
         gamma_powers: &[F],
     ) -> Vec<F> {
@@ -1125,7 +1125,7 @@ impl<F: JoltField> ReadRafSumcheckParams<F> {
     /// This particular Val virtualizes the claim output by the registers val-evaluation sumcheck
     /// and the instruction lookups sumcheck.
     fn compute_val_5(
-        sm: &mut StateManager<F, impl CommitmentScheme<Field = F>>,
+        sm: &mut StateManager<F, impl Transcript, impl CommitmentScheme<Field = F>>,
         opening_accumulator: &dyn OpeningAccumulator<F>,
         gamma_powers: &[F],
     ) -> Vec<F> {