Feat/hyrax returns (#594)

* Bring back Hyrax * Delete BatchType * Delete CommitShape * add comment
a16z · Feb 21, 2025 · 107d937 · 107d937
1 parent dae559d
commit 107d937
Show file tree

Hide file tree

Showing 17 changed files with 465 additions and 198 deletions.
diff --git a/jolt-core/benches/commit.rs b/jolt-core/benches/commit.rs
@@ -1,9 +1,8 @@
 use ark_bn254::{Bn254, Fr};
 use criterion::Criterion;
 use jolt_core::field::JoltField;
-use jolt_core::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use jolt_core::poly::commitment::commitment_scheme::CommitmentScheme;
 use jolt_core::poly::commitment::hyperkzg::HyperKZG;
-use jolt_core::poly::commitment::kzg::CommitMode;
 use jolt_core::poly::commitment::zeromorph::Zeromorph;
 use jolt_core::poly::multilinear_polynomial::MultilinearPolynomial;
 use jolt_core::utils::transcript::{KeccakTranscript, Transcript};
@@ -58,7 +57,7 @@ where
     // Compute known products (one per layer)
     let known_products: Vec<F> = leaves.iter().map(|layer| layer.iter().product()).collect();
 
-    let setup = PCS::setup(&[CommitShape::new(SRS_SIZE, BatchType::Big)]);
+    let setup = PCS::setup(SRS_SIZE);
 
     (leaves, setup, known_products)
 }
@@ -69,7 +68,6 @@ fn benchmark_commit<PCS, F, ProofTranscript>(
     num_layer: usize,
     layer_size: usize,
     threshold: u32,
-    batch_type: BatchType,
 ) where
     PCS: CommitmentScheme<ProofTranscript, Field = F>, // Generic over PCS implementing CommitmentScheme for field F
     F: JoltField,                                      // Generic over a field F
@@ -81,15 +79,11 @@ fn benchmark_commit<PCS, F, ProofTranscript>(
         .into_iter()
         .map(|layer| MultilinearPolynomial::from(layer))
         .collect::<Vec<_>>();
-    let mode = match batch_type {
-        BatchType::GrandProduct => CommitMode::GrandProduct,
-        _ => CommitMode::Default,
-    };
     c.bench_function(
         &format!("{} Commit(mode:{:?}): {}% Ones", name, mode, threshold),
         |b| {
             b.iter(|| {
-                PCS::batch_commit(&leaves, &setup, batch_type.clone());
+                PCS::batch_commit(&leaves, &setup);
             });
         },
     );
@@ -108,23 +102,13 @@ fn main() {
         num_layers,
         layer_size,
         90,
-        BatchType::Big,
     );
     benchmark_commit::<HyperKZG<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
         &mut criterion,
         "HyperKZG",
         num_layers,
         layer_size,
         90,
-        BatchType::GrandProduct,
-    );
-    benchmark_commit::<HyperKZG<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
-        &mut criterion,
-        "HyperKZG",
-        num_layers,
-        layer_size,
-        90,
-        BatchType::Big,
     );
 
     criterion.final_summary();

diff --git a/jolt-core/benches/grand_product.rs b/jolt-core/benches/grand_product.rs
@@ -1,7 +1,7 @@
 use ark_bn254::{Bn254, Fr};
 use criterion::Criterion;
 use jolt_core::field::JoltField;
-use jolt_core::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use jolt_core::poly::commitment::commitment_scheme::CommitmentScheme;
 use jolt_core::poly::commitment::hyperkzg::HyperKZG;
 use jolt_core::poly::opening_proof::{ProverOpeningAccumulator, VerifierOpeningAccumulator};
 use jolt_core::subprotocols::grand_product::{
@@ -69,7 +69,7 @@ where
     // Compute known products (one per layer)
     let known_products: Vec<F> = leaves.iter().map(|layer| layer.iter().product()).collect();
 
-    let setup = PCS::setup(&[CommitShape::new(SRS_SIZE, BatchType::Big)]);
+    let setup = PCS::setup(SRS_SIZE);
 
     ((leaves.concat(), batch_size), setup, known_products)
 }

diff --git a/jolt-core/src/jolt/vm/bytecode.rs b/jolt-core/src/jolt/vm/bytecode.rs
@@ -13,7 +13,7 @@ use crate::jolt::instruction::JoltInstructionSet;
 use crate::lasso::memory_checking::{
     Initializable, NoExogenousOpenings, StructuredPolynomialData, VerifierComputedOpening,
 };
-use crate::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::poly::compact_polynomial::{CompactPolynomial, SmallScalar};
 use crate::poly::multilinear_polynomial::{MultilinearPolynomial, PolynomialEvaluation};
 use common::constants::{BYTES_PER_INSTRUCTION, RAM_START_ADDRESS};
@@ -470,18 +470,6 @@ where
             );
         }
     }
-
-    /// Computes the shape of all commitment for use in PCS::setup().
-    pub fn commit_shapes(max_bytecode_size: usize, max_trace_length: usize) -> Vec<CommitShape> {
-        // Account for no-op prepended to bytecode
-        let max_bytecode_size = (max_bytecode_size + 1).next_power_of_two();
-        let max_trace_length = max_trace_length.next_power_of_two();
-
-        let read_write_shape = CommitShape::new(max_trace_length, BatchType::Big);
-        let init_final_shape = CommitShape::new(max_bytecode_size, BatchType::Small);
-
-        vec![read_write_shape, init_final_shape]
-    }
 }
 
 impl<F, PCS, ProofTranscript> MemoryCheckingProver<F, PCS, ProofTranscript>

diff --git a/jolt-core/src/jolt/vm/instruction_lookups.rs b/jolt-core/src/jolt/vm/instruction_lookups.rs
@@ -20,7 +20,7 @@ use crate::lasso::memory_checking::{
     Initializable, MultisetHashes, NoExogenousOpenings, StructuredPolynomialData,
     VerifierComputedOpening,
 };
-use crate::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::utils::transcript::Transcript;
 use crate::{
     lasso::memory_checking::{MemoryCheckingProof, MemoryCheckingProver, MemoryCheckingVerifier},
@@ -1252,16 +1252,6 @@ where
         subtable_lookup_indices
     }
 
-    /// Computes the shape of all commitments.
-    pub fn commitment_shapes(max_trace_length: usize) -> Vec<CommitShape> {
-        let max_trace_length = max_trace_length.next_power_of_two();
-
-        let read_write_shape = CommitShape::new(max_trace_length, BatchType::Big);
-        let init_final_shape = CommitShape::new(M, BatchType::Small);
-
-        vec![read_write_shape, init_final_shape]
-    }
-
     #[tracing::instrument(skip_all, name = "InstructionLookupsProof::compute_lookup_outputs")]
     fn compute_lookup_outputs(instructions: &Vec<JoltTraceStep<InstructionSet>>) -> Vec<u32> {
         instructions

diff --git a/jolt-core/src/jolt/vm/mod.rs b/jolt-core/src/jolt/vm/mod.rs
@@ -29,7 +29,7 @@ use crate::lasso::memory_checking::{
     Initializable, MemoryCheckingProver, MemoryCheckingVerifier, StructuredPolynomialData,
 };
 use crate::msm::icicle;
-use crate::poly::commitment::commitment_scheme::{BatchType, CommitmentScheme};
+use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::r1cs::inputs::{ConstraintInput, R1CSPolynomials, R1CSProof, R1CSStuff};
 use crate::utils::errors::ProofVerifyError;
 use crate::utils::thread::drop_in_background_thread;
@@ -248,8 +248,7 @@ impl<F: JoltField> JoltPolynomials<F> {
         drop(span);
 
         let trace_polys = self.read_write_values();
-        let trace_commitments =
-            PCS::batch_commit(&trace_polys, &preprocessing.generators, BatchType::Big);
+        let trace_commitments = PCS::batch_commit(&trace_polys, &preprocessing.generators);
 
         commitments
             .read_write_values_mut()
@@ -276,7 +275,6 @@ impl<F: JoltField> JoltPolynomials<F> {
         commitments.instruction_lookups.final_cts = PCS::batch_commit(
             &self.instruction_lookups.final_cts,
             &preprocessing.generators,
-            BatchType::Big,
         );
         drop(_guard);
         drop(span);
@@ -308,27 +306,6 @@ where
         F::initialize_lookup_tables(small_value_lookup_tables.clone());
         icicle::icicle_init();
 
-        let bytecode_commitment_shapes = BytecodeProof::<F, PCS, ProofTranscript>::commit_shapes(
-            max_bytecode_size,
-            max_trace_length,
-        );
-        let ram_commitment_shapes = ReadWriteMemoryPolynomials::<F>::commitment_shapes(
-            max_memory_address,
-            max_trace_length,
-        );
-        let timestamp_range_check_commitment_shapes =
-            TimestampValidityProof::<F, PCS, ProofTranscript>::commitment_shapes(max_trace_length);
-
-        let instruction_lookups_commitment_shapes = InstructionLookupsProof::<
-            C,
-            M,
-            F,
-            PCS,
-            Self::InstructionSet,
-            Self::Subtables,
-            ProofTranscript,
-        >::commitment_shapes(max_trace_length);
-
         let instruction_lookups_preprocessing = InstructionLookupsPreprocessing::preprocess::<
             M,
             Self::InstructionSet,
@@ -358,14 +335,16 @@ where
             .collect();
         let bytecode_preprocessing = BytecodePreprocessing::<F>::preprocess(bytecode_rows);
 
-        let commitment_shapes = [
-            bytecode_commitment_shapes,
-            ram_commitment_shapes,
-            timestamp_range_check_commitment_shapes,
-            instruction_lookups_commitment_shapes,
+        let max_poly_len: usize = [
+            (max_bytecode_size + 1).next_power_of_two(), // Account for no-op prepended to bytecode
+            max_trace_length.next_power_of_two(),
+            max_memory_address.next_power_of_two(),
+            M,
         ]
-        .concat();
-        let generators = PCS::setup(&commitment_shapes);
+        .into_iter()
+        .max()
+        .unwrap();
+        let generators = PCS::setup(max_poly_len);
 
         JoltPreprocessing {
             generators,

diff --git a/jolt-core/src/jolt/vm/read_write_memory.rs b/jolt-core/src/jolt/vm/read_write_memory.rs
@@ -12,7 +12,7 @@ use rayon::prelude::*;
 use std::collections::HashSet;
 use std::marker::PhantomData;
 
-use crate::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::utils::transcript::Transcript;
 use crate::{
     lasso::memory_checking::{
@@ -492,20 +492,6 @@ impl<F: JoltField> ReadWriteMemoryPolynomials<F> {
             identity: None,
         }
     }
-
-    /// Computes the shape of all commitments.
-    pub fn commitment_shapes(
-        max_memory_address: usize,
-        max_trace_length: usize,
-    ) -> Vec<CommitShape> {
-        let max_memory_address = max_memory_address.next_power_of_two();
-        let max_trace_length = max_trace_length.next_power_of_two();
-
-        let read_write_shape = CommitShape::new(max_trace_length, BatchType::Big);
-        let init_final_shape = CommitShape::new(max_memory_address, BatchType::Small);
-
-        vec![read_write_shape, init_final_shape]
-    }
 }
 
 impl<F, PCS, ProofTranscript> MemoryCheckingProver<F, PCS, ProofTranscript>

diff --git a/jolt-core/src/jolt/vm/timestamp_range_check.rs b/jolt-core/src/jolt/vm/timestamp_range_check.rs
@@ -3,7 +3,7 @@ use crate::field::JoltField;
 use crate::lasso::memory_checking::{
     ExogenousOpenings, Initializable, StructuredPolynomialData, VerifierComputedOpening,
 };
-use crate::poly::commitment::commitment_scheme::{BatchType, CommitShape, CommitmentScheme};
+use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::poly::compact_polynomial::{CompactPolynomial, SmallScalar};
 use crate::poly::multilinear_polynomial::{MultilinearPolynomial, PolynomialEvaluation};
 use crate::poly::opening_proof::{ProverOpeningAccumulator, VerifierOpeningAccumulator};
@@ -914,13 +914,6 @@ where
         Ok(())
     }
 
-    /// Computes the shape of all commitments.
-    pub fn commitment_shapes(max_trace_length: usize) -> Vec<CommitShape> {
-        let max_trace_length = max_trace_length.next_power_of_two();
-
-        vec![CommitShape::new(max_trace_length, BatchType::Big)]
-    }
-
     fn protocol_name() -> &'static [u8] {
         b"Timestamp Validity Proof"
     }

diff --git a/jolt-core/src/lasso/surge.rs b/jolt-core/src/lasso/surge.rs
@@ -2,7 +2,6 @@ use crate::{
     field::JoltField,
     jolt::vm::{JoltCommitments, JoltPolynomials, ProverDebugInfo},
     poly::{
-        commitment::commitment_scheme::BatchType,
         compact_polynomial::{CompactPolynomial, SmallScalar},
         multilinear_polynomial::{MultilinearPolynomial, PolynomialEvaluation},
         opening_proof::{ProverOpeningAccumulator, VerifierOpeningAccumulator},
@@ -416,18 +415,13 @@ where
 
         let mut commitments = SurgeCommitments::<PCS, ProofTranscript>::initialize(preprocessing);
         let trace_polys = polynomials.read_write_values();
-        let trace_comitments =
-            PCS::batch_commit(&trace_polys, generators, BatchType::SurgeReadWrite);
+        let trace_comitments = PCS::batch_commit(&trace_polys, generators);
         commitments
             .read_write_values_mut()
             .into_iter()
             .zip(trace_comitments.into_iter())
             .for_each(|(dest, src)| *dest = src);
-        commitments.final_cts = PCS::batch_commit(
-            &polynomials.final_cts,
-            generators,
-            BatchType::SurgeInitFinal,
-        );
+        commitments.final_cts = PCS::batch_commit(&polynomials.final_cts, generators);
 
         let num_rounds = num_lookups.log_2();
         let instruction = Instruction::default();
@@ -679,10 +673,7 @@ mod tests {
     use crate::{
         jolt::instruction::xor::XORInstruction,
         lasso::surge::SurgeProof,
-        poly::commitment::{
-            commitment_scheme::{BatchType, CommitShape, CommitmentScheme},
-            hyperkzg::HyperKZG,
-        },
+        poly::commitment::{commitment_scheme::CommitmentScheme, hyperkzg::HyperKZG},
     };
     use ark_bn254::{Bn254, Fr};
     use ark_std::test_rng;
@@ -703,10 +694,7 @@ mod tests {
         .collect();
 
         let preprocessing = SurgePreprocessing::preprocess();
-        let generators = HyperKZG::<_, KeccakTranscript>::setup(&[CommitShape::new(
-            M,
-            BatchType::SurgeReadWrite,
-        )]);
+        let generators = HyperKZG::<_, KeccakTranscript>::setup(M);
         let (proof, debug_info) = SurgeProof::<
             Fr,
             HyperKZG<Bn254, KeccakTranscript>,
@@ -735,10 +723,7 @@ mod tests {
         .collect();
 
         let preprocessing = SurgePreprocessing::preprocess();
-        let generators = HyperKZG::<_, KeccakTranscript>::setup(&[CommitShape::new(
-            M,
-            BatchType::SurgeReadWrite,
-        )]);
+        let generators = HyperKZG::<_, KeccakTranscript>::setup(M);
         let (proof, debug_info) = SurgeProof::<
             Fr,
             HyperKZG<Bn254, KeccakTranscript>,

diff --git a/jolt-core/src/poly/commitment/binius.rs b/jolt-core/src/poly/commitment/binius.rs
@@ -1,7 +1,5 @@
 #![allow(dead_code)]
 
-use crate::poly::commitment::commitment_scheme::BatchType;
-use crate::poly::commitment::commitment_scheme::CommitShape;
 use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::poly::multilinear_polynomial::MultilinearPolynomial;
 use crate::utils::errors::ProofVerifyError;
@@ -42,7 +40,7 @@ impl<ProofTranscript: Transcript> CommitmentScheme<ProofTranscript>
     type Proof = BiniusProof;
     type BatchedProof = BiniusBatchedProof;
 
-    fn setup(_shapes: &[CommitShape]) -> Self::Setup {
+    fn setup(_max_poly_len: usize) -> Self::Setup {
         None {}
     }
     fn commit(
@@ -51,11 +49,7 @@ impl<ProofTranscript: Transcript> CommitmentScheme<ProofTranscript>
     ) -> Self::Commitment {
         todo!()
     }
-    fn batch_commit<P>(
-        _polys: &[P],
-        _gens: &Self::Setup,
-        _batch_type: BatchType,
-    ) -> Vec<Self::Commitment>
+    fn batch_commit<P>(_polys: &[P], _gens: &Self::Setup) -> Vec<Self::Commitment>
     where
         P: Borrow<MultilinearPolynomial<Self::Field>>,
     {