Parallelize test_only segment simulation (#498)

* Parallelize test_only segment simulation

* Address comments

evm_arithmetization/src/generation/mod.rs

@@ -24,7 +24,6 @@ use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::generation::state::{GenerationState, State};
 use crate::generation::trie_extractor::{get_receipt_trie, get_state_trie, get_txn_trie};
-use crate::memory::columns::PREINITIALIZED_SEGMENTS;
 use crate::memory::segments::{Segment, PREINITIALIZED_SEGMENTS_INDICES};
 use crate::proof::{
     BlockHashes, BlockMetadata, ExtraBlockData, MemCap, PublicValues, RegistersData, TrieRoots,

evm_arithmetization/src/generation/mpt.rs

@@ -133,7 +133,7 @@ fn parse_storage_value(value_rlp: &[u8]) -> Result<Vec<U256>, ProgramError> {
     Ok(vec![value])
 }
 
-fn parse_storage_value_no_return(value_rlp: &[u8]) -> Result<Vec<U256>, ProgramError> {
+fn parse_storage_value_no_return(_value_rlp: &[u8]) -> Result<Vec<U256>, ProgramError> {
     Ok(vec![])
 }
 

zero_bin/ops/src/lib.rs

@@ -4,12 +4,14 @@ use std::time::Instant;
 #[cfg(not(feature = "test_only"))]
 use evm_arithmetization::generation::TrimmedGenerationInputs;
 use evm_arithmetization::{proof::PublicValues, AllData};
-#[cfg(not(feature = "test_only"))]
-use paladin::operation::FatalStrategy;
+#[cfg(feature = "test_only")]
+use evm_arithmetization::{prover::testing::simulate_execution_all_segments, GenerationInputs};
 use paladin::{
-    operation::{FatalError, Monoid, Operation, Result},
+    operation::{FatalError, FatalStrategy, Monoid, Operation, Result},
     registry, RemoteExecute,
 };
+#[cfg(feature = "test_only")]
+use proof_gen::types::Field;
 use proof_gen::{
     proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof},
     proof_types::{
@@ -24,6 +26,25 @@ use zero_bin_common::{debug_utils::save_inputs_to_disk, prover_state::p_state};
 
 registry!();
 
+#[cfg(feature = "test_only")]
+#[derive(Deserialize, Serialize, RemoteExecute)]
+pub struct BatchTestOnly {
+    pub save_inputs_on_error: bool,
+}
+
+#[cfg(feature = "test_only")]
+impl Operation for BatchTestOnly {
+    type Input = (GenerationInputs, usize);
+    type Output = ();
+
+    fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
+        simulate_execution_all_segments::<Field>(inputs.0, inputs.1)
+            .map_err(|err| FatalError::from_anyhow(err, FatalStrategy::Terminate))?;
+
+        Ok(())
+    }
+}
+
 #[derive(Deserialize, Serialize, RemoteExecute)]
 pub struct SegmentProof {
     pub save_inputs_on_error: bool,
@@ -72,8 +93,8 @@ impl Operation for SegmentProof {
     type Input = AllData;
     type Output = ();
 
-    fn execute(&self, _input: Self::Input) -> Result<Self::Output> {
-        todo!() // currently unused, change or remove
+    fn execute(&self, _all_data: Self::Input) -> Result<Self::Output> {
+        Ok(())
     }
 }
 

zero_bin/prover/src/lib.rs

@@ -9,7 +9,6 @@ use futures::{future::BoxFuture, stream::FuturesOrdered, FutureExt, TryFutureExt
 use num_traits::ToPrimitive as _;
 use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
-use proof_gen::types::Field;
 use serde::{Deserialize, Serialize};
 use tokio::io::AsyncWriteExt;
 use tokio::sync::oneshot;
@@ -83,8 +82,10 @@ impl BlockProverInput {
             .iter()
             .enumerate()
             .map(|(idx, txn_batch)| {
-                let segment_data_iterator =
-                    SegmentDataIterator::<Field>::new(txn_batch, Some(max_cpu_len_log));
+                let segment_data_iterator = SegmentDataIterator::<proof_gen::types::Field>::new(
+                    txn_batch,
+                    Some(max_cpu_len_log),
+                );
 
                 Directive::map(IndexedStream::from(segment_data_iterator), &seg_prove_ops)
                     .fold(&seg_agg_ops)
@@ -129,34 +130,48 @@ impl BlockProverInput {
     #[cfg(feature = "test_only")]
     pub async fn prove(
         self,
-        _runtime: &Runtime,
+        runtime: &Runtime,
         previous: Option<impl Future<Output = Result<GeneratedBlockProof>>>,
         prover_config: ProverConfig,
     ) -> Result<GeneratedBlockProof> {
-        use evm_arithmetization::prover::testing::simulate_execution_all_segments;
-        use plonky2_maybe_rayon::{MaybeIntoParIter, ParallelIterator};
+        use std::iter::repeat;
+
+        use futures::StreamExt;
+        use paladin::directive::{Directive, IndexedStream};
+
+        let ProverConfig {
+            max_cpu_len_log,
+            batch_size,
+            save_inputs_on_error,
+        } = prover_config;
 
         let block_number = self.get_block_number();
         info!("Testing witness generation for block {block_number}.");
 
-        let other_data = self.other_data;
-        let txs = self.block_trace.into_txn_proof_gen_ir(
-            &ProcessingMeta::new(resolve_code_hash_fn),
-            other_data.clone(),
-            prover_config.batch_size,
+        let block_generation_inputs = trace_decoder::entrypoint(
+            self.block_trace,
+            self.other_data,
+            batch_size,
+            |_| unimplemented!(),
         )?;
 
-        type F = GoldilocksField;
-        for txn in txs.into_iter() {
-            simulate_execution_all_segments::<F>(txn, prover_config.max_cpu_len_log)?;
-        }
-
-        // Wait for previous block proof
-        let _prev = match previous {
-            Some(it) => Some(it.await?),
-            None => None,
+        let batch_ops = ops::BatchTestOnly {
+            save_inputs_on_error,
         };
 
+        let simulation = Directive::map(
+            IndexedStream::from(
+                block_generation_inputs
+                    .into_iter()
+                    .zip(repeat(max_cpu_len_log)),
+            ),
+            &batch_ops,
+        );
+
+        let result = simulation.run(runtime).await?;
+
+        result.collect::<Vec<_>>().await;
+
         info!("Successfully generated witness for block {block_number}.");
 
         // Wait for previous block proof