Make computation of Bytecode commitments streaming

jolt-core/src/jolt/vm/bytecode.rs

@@ -8,7 +8,7 @@ use std::{collections::HashMap, marker::PhantomData};
 
 use crate::field::JoltField;
 use crate::jolt::instruction::JoltInstructionSet;
-use crate::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use crate::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme, StreamingCommitmentScheme};
 use crate::poly::eq_poly::EqPolynomial;
 use crate::utils::{
     streaming::map_state,
@@ -577,6 +577,55 @@ impl<C: CommitmentScheme> AppendToTranscript for BytecodeCommitment<C> {
     }
 }
 
+pub struct StreamingBytecodeCommitment<C: StreamingCommitmentScheme> {
+    a_read_write: C::State,
+    v_read_write: [C::State; 6],
+    t_read: C::State,
+}
+
+impl<C: StreamingCommitmentScheme> StreamingBytecodeCommitment<C> {
+    /// Initialize a streaming computation of a commitment.
+    pub fn initialize(size: usize, setup: &C::Setup, batch_type: &BatchType) -> Self {
+        let a_read_write = C::initialize(size, setup, batch_type);
+        let v_read_write = std::array::from_fn(|_| a_read_write.clone());
+        let t_read = a_read_write.clone();
+
+        StreamingBytecodeCommitment {
+            a_read_write,
+            v_read_write,
+            t_read,
+        }
+    }
+
+    /// Process one step to compute the commitment.
+    pub fn process(self, step: &BytecodePolynomialStep<C::Field, C>) -> Self {
+        let step_v_read_write = [
+            step.v_read_write.address,
+            step.v_read_write.bitflags,
+            step.v_read_write.rd,
+            step.v_read_write.rs1,
+            step.v_read_write.rs2,
+            step.v_read_write.imm,
+        ];
+        StreamingBytecodeCommitment {
+            a_read_write: C::process(self.a_read_write, step.a_read_write),
+            v_read_write: self.v_read_write.into_iter().zip(step_v_read_write)
+                .map(|(vrd, step)| C::process(vrd, step)).collect::<Vec<_>>().try_into().unwrap(),
+            t_read: C::process(self.t_read, step.t_read),
+        }
+    }
+
+    /// Return the trace commitments.
+    pub fn finalize(self) -> Vec<C::Commitment> {
+        [
+            C::finalize(self.a_read_write),
+            C::finalize(self.t_read),
+        ].into_iter().chain(
+            self.v_read_write.into_iter().map(|vrw| C::finalize(vrw))
+        ).collect()
+    }
+}
+
 impl<F, C> StructuredCommitment<C> for BytecodePolynomials<F, C>
 where
     F: JoltField,

jolt-core/src/jolt/vm/mod.rs

@@ -19,7 +19,10 @@ use crate::jolt::{
     },
     subtable::JoltSubtableSet,
     vm::{
-        bytecode::StreamingBytecodePolynomials,
+        bytecode::{
+            StreamingBytecodeCommitment,
+            StreamingBytecodePolynomials,
+        },
         timestamp_range_check::{
             RangeCheckPolynomials,
             TimestampValidityProof,
@@ -367,13 +370,13 @@ pub trait Jolt<F: JoltField, PCS: StreamingCommitmentScheme<Field = F>, const C:
         );
 
         let streaming_bytecode_polynomials = StreamingBytecodePolynomials::<F, PCS>::new(&preprocessing.bytecode, &mut trace2);
-        let initialized_commitment = PCS::initialize(streaming_bytecode_polynomials.length(), &preprocessing.generators, &BatchType::Big);
+        let initialized_commitment = StreamingBytecodeCommitment::initialize(streaming_bytecode_polynomials.length(), &preprocessing.generators, &BatchType::Big);
         // JP: `fold` likely isn't sufficient since we need to extract the internal state.
         let streaming_trace_commitments =
             streaming_bytecode_polynomials.fold(initialized_commitment, |state, step| {
-                PCS::process(state, step.a_read_write)
+                StreamingBytecodeCommitment::process(state, &step)
             });
-        let a_read_write_commitment = PCS::finalize(streaming_trace_commitments);
+        let bytecode_commitments = StreamingBytecodeCommitment::finalize(streaming_trace_commitments);
 
         let jolt_polynomials = JoltPolynomials {
             bytecode: bytecode_polynomials,
@@ -383,7 +386,9 @@ pub trait Jolt<F: JoltField, PCS: StreamingCommitmentScheme<Field = F>, const C:
         };
 
         let mut jolt_commitments = jolt_polynomials.commit(&preprocessing.generators);
-        assert_eq!(a_read_write_commitment, jolt_commitments.bytecode.trace_commitments[0]);
+
+        /// TODO: Temp, remove me XXX
+        assert_eq!(bytecode_commitments, jolt_commitments.bytecode.trace_commitments);
 
         let (witness_segments, r1cs_commitments, r1cs_builder) = Self::r1cs_setup(
             padded_trace_length,

jolt-core/src/poly/commitment/commitment_scheme.rs

@@ -1,4 +1,5 @@
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use std::fmt::Debug;
 
 use crate::{
     field::JoltField,
@@ -100,7 +101,7 @@ pub trait CommitmentScheme: Clone + Sync + Send + 'static {
 }
 
 pub trait StreamingCommitmentScheme: CommitmentScheme {
-    type State;
+    type State: Clone + Debug;
 
     fn initialize(size: usize, setup: &Self::Setup, batch_type: &BatchType) -> Self::State;
     fn process(state: Self::State, eval: Self::Field) -> Self::State;

jolt-core/src/poly/commitment/hyrax.rs

@@ -148,6 +148,7 @@ impl<F: JoltField, G: CurveGroup<ScalarField = F>> CommitmentScheme for HyraxSch
     }
 }
 
+#[derive(Clone, Debug)]
 pub struct HyraxSchemeState<G: CurveGroup> {
     row_commitments: Vec<G>,
     generators: Vec<<G as CurveGroup>::Affine>,