Parallelize test_only segment simulation

zero_bin/ops/src/lib.rs

@@ -4,12 +4,15 @@ use std::time::Instant;
 #[cfg(not(feature = "test_only"))]
 use evm_arithmetization::generation::TrimmedGenerationInputs;
 use evm_arithmetization::proof::PublicValues;
+use evm_arithmetization::prover::SegmentDataIterator;
+use evm_arithmetization::GenerationInputs;
 #[cfg(not(feature = "test_only"))]
 use paladin::operation::FatalStrategy;
 use paladin::{
     operation::{FatalError, Monoid, Operation, Result},
     registry, RemoteExecute,
 };
+use proof_gen::types::Field;
 use proof_gen::{
     proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof},
     proof_types::{
@@ -25,6 +28,22 @@ 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 {}
+
+#[cfg(feature = "test_only")]
+impl Operation for BatchTestOnly {
+    type Input = (GenerationInputs, usize);
+    type Output = ();
+
+    fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
+        let _ = SegmentDataIterator::<Field>::new(&inputs.0, Some(inputs.1)).collect::<Vec<_>>();
+
+        Ok(())
+    }
+}
+
 #[derive(Deserialize, Serialize, RemoteExecute)]
 pub struct SegmentProof {
     pub save_inputs_on_error: bool,
@@ -73,8 +92,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

@@ -135,12 +135,14 @@ 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::field::goldilocks_field::GoldilocksField;
+        use std::iter::repeat;
+
+        use futures::StreamExt;
+        use paladin::directive::{Directive, IndexedStream};
 
         let block_number = self.get_block_number();
         info!("Testing witness generation for block {block_number}.");
@@ -152,10 +154,16 @@ impl BlockProverInput {
             prover_config.batch_size,
         )?;
 
-        type F = GoldilocksField;
-        for txn in txs.into_iter() {
-            simulate_execution_all_segments::<F>(txn, prover_config.max_cpu_len_log)?;
-        }
+        let batch_ops = ops::BatchTestOnly {};
+
+        let simulation = Directive::map(
+            IndexedStream::from(txs.into_iter().zip(repeat(prover_config.max_cpu_len_log))),
+            &batch_ops,
+        );
+
+        let result = simulation.run(runtime).await?;
+
+        result.collect::<Vec<_>>().await;
 
         // Wait for previous block proof
         let _prev = match previous {