Apply comments

Cargo.toml

@@ -19,6 +19,7 @@ serde = "1.0.166"
 serde_json = "1.0.96"
 serde-big-array = "0.5.1"
 thiserror = "1.0.49"
+hashbrown = "0.14.0"
 
 # plonky2-related dependencies
 plonky2 = { git = "https://github.com/0xPolygonZero/plonky2.git", rev = "dc77c77f2b06500e16ad4d7f1c2b057903602eed" }

evm_arithmetization/Cargo.toml

@@ -36,7 +36,7 @@ rlp = { workspace = true }
 rlp-derive = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 static_assertions = "1.1.0"
-hashbrown = { version = "0.14.0" }
+hashbrown = { workspace = true }
 tiny-keccak = "2.0.2"
 serde_json = { workspace = true }
 serde-big-array = { workspace = true }

evm_arithmetization/src/cpu/kernel/tests/init_exc_stop.rs

@@ -15,6 +15,8 @@ use crate::generation::NUM_EXTRA_CYCLES_AFTER;
 use crate::generation::NUM_EXTRA_CYCLES_BEFORE;
 use crate::memory::segments::Segment;
 use crate::proof::BlockMetadata;
+use crate::proof::RegistersData;
+use crate::proof::RegistersIdx;
 use crate::proof::TrieRoots;
 use crate::witness::memory::MemoryAddress;
 use crate::witness::state::RegistersState;
@@ -101,22 +103,24 @@ fn test_init_exc_stop() {
     interpreter.set_is_kernel(true);
     interpreter.clock = 0;
 
-    // Set the program counter and `is_kernel` at the end of the execution. They
-    // have offsets 6 and 7 respectively in segment `RegistersStates`.
+    // Set the program counter and `is_kernel` at the end of the execution. The
+    // `registers_before` and `registers_after` are stored contiguously in the
+    // `RegistersState` segment. We need to update `registers_after` here, hence the
+    // offset by `RegistersData::SIZE`.
     let regs_to_set = [
         (
             MemoryAddress {
                 context: 0,
                 segment: Segment::RegistersStates.unscale(),
-                virt: 6,
+                virt: RegistersData::SIZE + RegistersIdx::ProgramCounter as usize,
             },
             pc_u256,
         ),
         (
             MemoryAddress {
                 context: 0,
                 segment: Segment::RegistersStates.unscale(),
-                virt: 7,
+                virt: RegistersData::SIZE + RegistersIdx::IsKernel as usize,
             },
             U256::one(),
         ),
@@ -128,7 +132,7 @@ fn test_init_exc_stop() {
     // The "-2" comes from the fact that:
     // - we stop 1 cycle before the max, to allow for one padding row, which is
     //   needed for CPU STARK.
-    // - we normally need one additional cycle to enter `exc_stop`.
+    // - we need one additional cycle to enter `exc_stop`.
     assert_eq!(
         interpreter.get_clock(),
         NUM_EXTRA_CYCLES_AFTER - 2,

evm_arithmetization/src/fixed_recursive_verifier.rs

@@ -900,8 +900,8 @@ where
             lhs_pv.mem_before.clone(),
         );
 
-        // If the rhs is a real proof: All the block metadata is the same for both
-        // segments. It is also the case for extra_block_data.
+        // If the rhs is a real proof, all the block metadata must be the same for both
+        // segments. It is also the case for the extra block data.
         TrieRootsTarget::conditional_assert_eq(
             &mut builder,
             is_not_dummy,
@@ -1695,35 +1695,29 @@ where
     ///
     /// - `lhs_is_agg`: a boolean indicating whether the left child proof is an
     ///   aggregation proof or a regular segment proof.
-    /// - `lhs_proof`: the left child proof.
-    /// - `lhs_public_values`: the public values associated to the right child
-    ///   proof.
+    /// - `lhs_proof`: the left child prover output data.
     /// - `rhs_is_agg`: a boolean indicating whether the right child proof is an
     ///   aggregation proof or a regular transaction proof.
-    /// - `rhs_proof`: the right child proof.
-    /// - `rhs_public_values`: the public values associated to the right child
-    ///   proof.
+    /// - `rhs_proof`: the right child prover output data.
     ///
     /// # Outputs
     ///
-    /// This method outputs a tuple of [`ProofWithPublicInputs<F, C, D>`] and
-    /// its [`PublicValues`]. Only the proof with public inputs is necessary
-    /// for a verifier to assert correctness of the computation,
-    /// but the public values are output for the prover convenience, as these
-    /// are necessary during proof aggregation.
+    /// This method outputs a [`ProverOutputData<F, C, D>`]. Only the proof with
+    /// public inputs is necessary for a verifier to assert correctness of
+    /// the computation, but the public values and `is_dummy` are output for the
+    /// prover convenience, as these are necessary during proof aggregation.
     pub fn prove_segment_aggregation(
         &self,
         lhs_is_agg: bool,
-        lhs_proof: &ProofWithPublicInputs<F, C, D>,
-        lhs_public_values: PublicValues,
-
+        lhs_prover_output: &ProverOutputData<F, C, D>,
         rhs_is_agg: bool,
-        rhs_is_dummy: bool,
-        rhs_proof: &ProofWithPublicInputs<F, C, D>,
-        rhs_public_values: PublicValues,
-    ) -> anyhow::Result<(ProofWithPublicInputs<F, C, D>, PublicValues)> {
+        rhs_prover_output: &ProverOutputData<F, C, D>,
+    ) -> anyhow::Result<ProverOutputData<F, C, D>> {
         let mut agg_inputs = PartialWitness::new();
 
+        let lhs_proof = &lhs_prover_output.proof_with_pis;
+        let rhs_proof = &rhs_prover_output.proof_with_pis;
+        let rhs_is_dummy = rhs_prover_output.is_dummy;
         Self::set_dummy_if_necessary(
             &self.segment_aggregation.lhs,
             lhs_is_agg,
@@ -1760,14 +1754,17 @@ where
 
         // Aggregates both `PublicValues` from the provided proofs into a single one.
 
+        let lhs_public_values = &lhs_prover_output.public_values;
+        let rhs_public_values = &rhs_prover_output.public_values;
+
         let real_public_values = if rhs_is_dummy {
             lhs_public_values.clone()
         } else {
             rhs_public_values.clone()
         };
 
         let agg_public_values = PublicValues {
-            trie_roots_before: lhs_public_values.trie_roots_before,
+            trie_roots_before: lhs_public_values.trie_roots_before.clone(),
             trie_roots_after: real_public_values.trie_roots_after,
             extra_block_data: ExtraBlockData {
                 checkpoint_state_trie_root: lhs_public_values
@@ -1780,9 +1777,9 @@ where
             },
             block_metadata: real_public_values.block_metadata,
             block_hashes: real_public_values.block_hashes,
-            registers_before: lhs_public_values.registers_before,
+            registers_before: lhs_public_values.registers_before.clone(),
             registers_after: real_public_values.registers_after,
-            mem_before: lhs_public_values.mem_before,
+            mem_before: lhs_public_values.mem_before.clone(),
             mem_after: real_public_values.mem_after,
         };
 
@@ -1796,7 +1793,12 @@ where
         })?;
 
         let aggregation_proof = self.segment_aggregation.circuit.prove(agg_inputs)?;
-        Ok((aggregation_proof, agg_public_values))
+        let agg_output = ProverOutputData {
+            is_dummy: false,
+            proof_with_pis: aggregation_proof,
+            public_values: agg_public_values,
+        };
+        Ok(agg_output)
     }
 
     pub fn verify_segment_aggregation(

evm_arithmetization/src/generation/state.rs

@@ -77,9 +77,9 @@ pub(crate) trait State<F: Field> {
     fn get_context(&self) -> usize;
 
     /// Checks whether we have reached the maximal cpu length.
-    fn at_end_segment(&self, opt_max_cpu_len: Option<usize>) -> bool {
-        if let Some(max_cpu_len_log) = opt_max_cpu_len {
-            self.get_clock() == (1 << max_cpu_len_log) - NUM_EXTRA_CYCLES_AFTER
+    fn at_end_segment(&self, opt_cycle_limit: Option<usize>) -> bool {
+        if let Some(cycle_limit) = opt_cycle_limit {
+            self.get_clock() == cycle_limit
         } else {
             false
         }
@@ -173,6 +173,9 @@ pub(crate) trait State<F: Field> {
     {
         let halt_offsets = self.get_halt_offsets();
 
+        let cycle_limit =
+            max_cpu_len_log.map(|max_len_log| (1 << max_len_log) - NUM_EXTRA_CYCLES_AFTER);
+
         let mut final_registers = RegistersState::default();
         let final_mem = self.get_generation_state().memory.clone();
         let mut running = true;
@@ -182,7 +185,7 @@ pub(crate) trait State<F: Field> {
             let pc = registers.program_counter;
 
             let halt_final = registers.is_kernel && halt_offsets.contains(&pc);
-            if running && (self.at_halt() || self.at_end_segment(max_cpu_len_log)) {
+            if running && (self.at_halt() || self.at_end_segment(cycle_limit)) {
                 running = false;
                 final_registers = registers;
 

evm_arithmetization/src/proof.rs

@@ -378,7 +378,19 @@ pub struct RegistersData {
     /// Gas used so far.
     pub gas_used: U256,
 }
+
+pub(crate) enum RegistersIdx {
+    ProgramCounter = 0,
+    IsKernel = 1,
+    StackLen = 2,
+    StackTop = 3,
+    Context = 4,
+    GasUsed = 5,
+}
+
 impl RegistersData {
+    pub(crate) const SIZE: usize = 6;
+
     pub fn from_public_inputs<F: RichField>(pis: &[F]) -> Self {
         assert!(pis.len() == RegistersDataTarget::SIZE);
 

evm_arithmetization/tests/empty_txn_list.rs

@@ -178,44 +178,43 @@ fn test_empty_txn_list() -> anyhow::Result<()> {
     );
 
     // We can duplicate the proofs here because the state hasn't mutated.
-    let (segmented_agg_proof, segmented_agg_public_values) = all_circuits
-        .prove_segment_aggregation(
-            false,
-            &segment_proofs_data[0].proof_with_pis,
-            segment_proofs_data[0].public_values.clone(),
-            false,
-            segment_proofs_data[1].is_dummy,
-            &segment_proofs_data[1].proof_with_pis,
-            segment_proofs_data[1].public_values.clone(),
-        )?;
-    all_circuits.verify_segment_aggregation(&segmented_agg_proof)?;
-
-    let (segmented_agg_proof, segmented_agg_public_values) = all_circuits
-        .prove_segment_aggregation(
-            true,
-            &segmented_agg_proof,
-            segmented_agg_public_values,
-            false,
-            segment_proofs_data[2].is_dummy,
-            &segment_proofs_data[2].proof_with_pis,
-            segment_proofs_data[2].public_values.clone(),
-        )?;
-    all_circuits.verify_segment_aggregation(&segmented_agg_proof)?;
+    let aggregation_output_data = all_circuits.prove_segment_aggregation(
+        false,
+        &segment_proofs_data[0],
+        false,
+        &segment_proofs_data[1],
+    )?;
+    all_circuits.verify_segment_aggregation(&aggregation_output_data.proof_with_pis)?;
+
+    let aggregation_output_data = all_circuits.prove_segment_aggregation(
+        true,
+        &aggregation_output_data,
+        false,
+        &segment_proofs_data[2],
+    )?;
+    all_circuits.verify_segment_aggregation(&aggregation_output_data.proof_with_pis)?;
 
     // Test retrieved public values from the proof public inputs.
     let retrieved_public_values = PublicValues::from_public_inputs(
-        &segmented_agg_proof.public_inputs,
-        segmented_agg_public_values.mem_before.mem_cap.len(),
+        &aggregation_output_data.proof_with_pis.public_inputs,
+        aggregation_output_data
+            .public_values
+            .mem_before
+            .mem_cap
+            .len(),
+    );
+    assert_eq!(
+        retrieved_public_values,
+        aggregation_output_data.public_values
     );
-    assert_eq!(retrieved_public_values, segmented_agg_public_values);
 
     let (txn_proof, txn_public_values) = all_circuits.prove_transaction_aggregation(
         false,
-        &segmented_agg_proof,
-        segmented_agg_public_values.clone(),
+        &aggregation_output_data.proof_with_pis,
+        aggregation_output_data.public_values.clone(),
         false,
-        &segmented_agg_proof,
-        segmented_agg_public_values,
+        &aggregation_output_data.proof_with_pis,
+        aggregation_output_data.public_values,
     )?;
     all_circuits.verify_txn_aggregation(&txn_proof)?;
 

evm_arithmetization/tests/log_opcode.rs

@@ -477,17 +477,14 @@ fn test_log_with_aggreg() -> anyhow::Result<()> {
 
     assert_eq!(segment_proofs_data_first.len(), 2); // second one is a dummy segment
 
-    let (segment_agg_proof_first, updated_agg_public_values_first) = all_circuits
-        .prove_segment_aggregation(
-            false,
-            &segment_proofs_data_first[0].proof_with_pis,
-            segment_proofs_data_first[0].public_values.clone(),
-            false,
-            segment_proofs_data_first[1].is_dummy,
-            &segment_proofs_data_first[1].proof_with_pis,
-            segment_proofs_data_first[1].public_values.clone(),
-        )?;
-    all_circuits.verify_segment_aggregation(&segment_agg_proof_first)?;
+    let segment_agg_prover_output_data_first = all_circuits.prove_segment_aggregation(
+        false,
+        &segment_proofs_data_first[0],
+        false,
+        &segment_proofs_data_first[1],
+    )?;
+    all_circuits
+        .verify_segment_aggregation(&segment_agg_prover_output_data_first.proof_with_pis)?;
 
     // The gas used and transaction number are fed to the next transaction, so the
     // two proofs can be correctly aggregated.
@@ -627,25 +624,22 @@ fn test_log_with_aggreg() -> anyhow::Result<()> {
         all_circuits.verify_root(proof.clone())?;
     }
 
-    let (segment_agg_proof_second, updated_agg_public_values_second) = all_circuits
-        .prove_segment_aggregation(
-            false,
-            &segment_proofs_data_second[0].proof_with_pis,
-            segment_proofs_data_second[0].public_values.clone(),
-            false,
-            segment_proofs_data_second[1].is_dummy,
-            &segment_proofs_data_second[1].proof_with_pis,
-            segment_proofs_data_second[1].public_values.clone(),
-        )?;
-    all_circuits.verify_segment_aggregation(&segment_agg_proof_second)?;
+    let segment_agg_prover_output_data_second = all_circuits.prove_segment_aggregation(
+        false,
+        &segment_proofs_data_second[0],
+        false,
+        &segment_proofs_data_second[1],
+    )?;
+    all_circuits
+        .verify_segment_aggregation(&segment_agg_prover_output_data_second.proof_with_pis)?;
 
     let (txn_proof, txn_pv) = all_circuits.prove_transaction_aggregation(
         false,
-        &segment_agg_proof_first,
-        updated_agg_public_values_first,
+        &segment_agg_prover_output_data_first.proof_with_pis,
+        segment_agg_prover_output_data_first.public_values,
         false,
-        &segment_agg_proof_second,
-        updated_agg_public_values_second,
+        &segment_agg_prover_output_data_second.proof_with_pis,
+        segment_agg_prover_output_data_second.public_values,
     )?;
 
     let (first_block_proof, _block_public_values) =
@@ -717,24 +711,21 @@ fn test_log_with_aggreg() -> anyhow::Result<()> {
         all_circuits.verify_root(proof.clone())?;
     }
 
-    let (segment_agg_proof, updated_agg_public_values) = all_circuits.prove_segment_aggregation(
+    let segment_agg_prover_output_data = all_circuits.prove_segment_aggregation(
         false,
-        &segment_proofs_data[0].proof_with_pis,
-        segment_proofs_data[0].public_values.clone(),
+        &segment_proofs_data[0],
         false,
-        segment_proofs_data[1].is_dummy,
-        &segment_proofs_data[1].proof_with_pis,
-        segment_proofs_data[1].public_values.clone(),
+        &segment_proofs_data[1],
     )?;
-    all_circuits.verify_segment_aggregation(&segment_agg_proof)?;
+    all_circuits.verify_segment_aggregation(&segment_agg_prover_output_data.proof_with_pis)?;
 
     let (second_txn_proof, second_txn_pvs) = all_circuits.prove_transaction_aggregation(
         false,
-        &segment_agg_proof,
-        updated_agg_public_values.clone(),
+        &segment_agg_prover_output_data.proof_with_pis,
+        segment_agg_prover_output_data.public_values.clone(),
         false,
-        &segment_agg_proof,
-        updated_agg_public_values,
+        &segment_agg_prover_output_data.proof_with_pis,
+        segment_agg_prover_output_data.public_values,
     )?;
     let (second_block_proof, _block_public_values) = all_circuits.prove_block(
         None, // We don't specify a previous proof, considering block 1 as the new checkpoint.

proof_gen/Cargo.toml

@@ -15,7 +15,7 @@ log = { workspace = true }
 paste = "1.0.14"
 plonky2 = { workspace = true }
 serde = { workspace = true }
-hashbrown = { version = "0.14.0" }
+hashbrown = { workspace = true }
 
 # Local dependencies
 evm_arithmetization = { version = "0.1.3", path = "../evm_arithmetization" }