Address comments

zero_bin/ops/src/lib.rs

@@ -4,14 +4,13 @@ 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;
+#[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},
@@ -30,15 +29,18 @@ registry!();
 
 #[cfg(feature = "test_only")]
 #[derive(Deserialize, Serialize, RemoteExecute)]
-pub struct BatchTestOnly {}
+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> {
-        let _ = SegmentDataIterator::<Field>::new(&inputs.0, Some(inputs.1)).collect::<Vec<_>>();
+        simulate_execution_all_segments::<Field>(inputs.0, inputs.1)
+            .map_err(|err| FatalError::from_anyhow(err, FatalStrategy::Terminate))?;
 
         Ok(())
     }

zero_bin/prover/src/lib.rs

@@ -144,20 +144,28 @@ impl BlockProverInput {
         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,
+            batch_size,
         )?;
 
-        let batch_ops = ops::BatchTestOnly {};
+        let batch_ops = ops::BatchTestOnly {
+            save_inputs_on_error,
+        };
 
         let simulation = Directive::map(
-            IndexedStream::from(txs.into_iter().zip(repeat(prover_config.max_cpu_len_log))),
+            IndexedStream::from(txs.into_iter().zip(repeat(max_cpu_len_log))),
             &batch_ops,
         );