From c410daaff205b0f12f7a628f37df1daeb5eb82fb Mon Sep 17 00:00:00 2001
From: Arpit Temani <temaniarpit27@gmail.com>
Date: Tue, 10 Sep 2024 21:49:50 +0530
Subject: [PATCH] separate segment proof and segment agg proof

---
 zero_bin/prover/src/lib.rs  | 32 +++++++++++++++++++-------------
 zero_bin/tools/prove_rpc.sh |  4 ++--
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/zero_bin/prover/src/lib.rs b/zero_bin/prover/src/lib.rs
index 97a80df6a..cb0fe13ba 100644
--- a/zero_bin/prover/src/lib.rs
+++ b/zero_bin/prover/src/lib.rs
@@ -9,8 +9,7 @@ use anyhow::{Context, Result};
 use evm_arithmetization::SegmentDataIterator;
 use futures::StreamExt;
 use futures::{
-    future, future::BoxFuture, stream::FuturesOrdered, stream::FuturesUnordered, FutureExt,
-    TryFutureExt, TryStreamExt,
+    future, future::BoxFuture, stream::FuturesUnordered, FutureExt, TryFutureExt, TryStreamExt,
 };
 use num_traits::ToPrimitive as _;
 use paladin::directive::{Directive, IndexedStream};
@@ -86,11 +85,9 @@ impl BlockProverInput {
             save_inputs_on_error,
         };
 
-        // Segment the batches, prove segments, and aggregate the segment proofs.
-        let vec_segment_agg_proofs: FuturesOrdered<_> = block_generation_inputs
+        let vec_segment_agg_proofs: FuturesUnordered<_> = block_generation_inputs
             .iter()
-            .enumerate()
-            .map(|(_, txn_batch)| {
+            .map(|txn_batch| {
                 let segment_data_iterator = SegmentDataIterator::<proof_gen::types::Field>::new(
                     txn_batch,
                     Some(max_cpu_len_log),
@@ -102,13 +99,22 @@ impl BlockProverInput {
             })
             .collect();
 
-        // Collect results of segment proofs and prepare for batch aggregation
-        let batch_proof_futs = vec_segment_agg_proofs.map(|segment_proofs| {
-            let x = segment_proofs.unwrap();
-            Directive::fold(IndexedStream::new(x), &seg_agg_ops)
-                .run(&proof_runtime.block_proof_runtime)
-                .map(move |e| e.map(|p| (proof_gen::proof_types::BatchAggregatableProof::from(p))))
-        });
+        // Await all the futures concurrently
+        let vec_segment_agg_proofs: Vec<_> = vec_segment_agg_proofs.collect().await;
+
+        let batch_proof_futs: FuturesUnordered<_> = vec_segment_agg_proofs
+            .into_iter()
+            .enumerate()
+            .map(|(idx, result)| {
+                // todo - check how to handle the unwrap
+                let res = result.unwrap();
+                Directive::fold(IndexedStream::new(res), &seg_agg_ops)
+                    .run(&proof_runtime.block_proof_runtime)
+                    .map(move |e| {
+                        e.map(|p| (idx, proof_gen::proof_types::BatchAggregatableProof::from(p)))
+                    })
+            })
+            .collect();
 
         // Fold the batch aggregated proof stream into a single proof.
         let final_batch_proof =
diff --git a/zero_bin/tools/prove_rpc.sh b/zero_bin/tools/prove_rpc.sh
index 0c5e19c96..1c7420491 100755
--- a/zero_bin/tools/prove_rpc.sh
+++ b/zero_bin/tools/prove_rpc.sh
@@ -110,7 +110,7 @@ fi
 if [[ $8 == "test_only" ]]; then
     # test only run
     echo "Proving blocks ${BLOCK_INTERVAL} in a test_only mode now... (Total: ${TOT_BLOCKS})"
-    command='cargo r --release --bin leader -- --test-only --worker-run-mode split --load-strategy on-demand --proof-output-dir $PROOF_OUTPUT_DIR --block-batch-size $BLOCK_BATCH_SIZE rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL  $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
+    command='cargo r --release --bin leader -- --test-only --runtime in-memory --load-strategy on-demand --proof-output-dir $PROOF_OUTPUT_DIR --block-batch-size $BLOCK_BATCH_SIZE rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL  $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
     if [ "$OUTPUT_TO_TERMINAL" = true ]; then
         eval $command
         retVal=$?
@@ -133,7 +133,7 @@ if [[ $8 == "test_only" ]]; then
 else
     # normal run
     echo "Proving blocks ${BLOCK_INTERVAL} now... (Total: ${TOT_BLOCKS})"
-    command='cargo r --release --bin leader --  --worker-run-mode split --load-strategy on-demand --proof-output-dir $PROOF_OUTPUT_DIR --block-batch-size $BLOCK_BATCH_SIZE rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --block-interval $BLOCK_INTERVAL $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
+    command='cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand --proof-output-dir $PROOF_OUTPUT_DIR --block-batch-size $BLOCK_BATCH_SIZE rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --block-interval $BLOCK_INTERVAL $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
     if [ "$OUTPUT_TO_TERMINAL" = true ]; then
         eval $command
         echo -e "Proof generation finished with result: $?"