separate segment proof and segment agg proof

0xPolygonZero · Sep 10, 2024 · e0231ad · e0231ad
1 parent 3701b1e
commit e0231ad
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 15 deletions.
diff --git 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
@@ -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: $?"