From 8060288ed0703be5ad8ad297ac852e001d95a90c Mon Sep 17 00:00:00 2001
From: Marko Atanasievski <markoatana@gmail.com>
Date: Wed, 24 Jul 2024 13:59:33 +0200
Subject: [PATCH] feat: add prover config argumetns to scripts

---
 zero_bin/tools/prove_rpc.sh   | 12 +++++++++++-
 zero_bin/tools/prove_stdio.sh | 16 +++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/zero_bin/tools/prove_rpc.sh b/zero_bin/tools/prove_rpc.sh
index 6da293b4f..477698301 100755
--- a/zero_bin/tools/prove_rpc.sh
+++ b/zero_bin/tools/prove_rpc.sh
@@ -64,6 +64,16 @@ OUTPUT_TO_TERMINAL="${OUTPUT_TO_TERMINAL:-false}"
 # Only generate proof by default
 RUN_VERIFICATION="${RUN_VERIFICATION:-false}"
 
+# Prover config. Override the defaults if needed by setting the env variables.
+PROVER_BATCH_SIZE="${PROVER_BATCH_SIZE:-1}"
+PROVER_SEGMENT_CHUNK_SIZE="${PROVER_SEGMENT_CHUNK_SIZE:-64}"
+PROVER_MAC_CPU_LEN_LOG="${PROVER_MAC_CPU_LEN_LOG:-20}"
+if [[ $PROVER_SAVE_INPUTS_ON_ERROR == "true" ]]; then
+    PROVER_SAVE_INPUTS_ON_ERROR="--save-inputs-on-error"
+else
+    PROVER_SAVE_INPUTS_ON_ERROR=""
+fi
+
 mkdir -p $PROOF_OUTPUT_DIR
 
 if [ $IGNORE_PREVIOUS_PROOFS ]; then
@@ -122,7 +132,7 @@ if [[ $8 == "test_only" ]]; then
 else
     # normal run
     echo "Proving blocks ${BLOCK_INTERVAL} now... (Total: ${TOT_BLOCKS})"
-    command='cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
+    command='cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand --batch-size $PROVER_BATCH_SIZE --max-cpu-len-log $PROVER_MAC_CPU_LEN_LOG --segment-chunk-size $PROVER_SEGMENT_CHUNK_SIZE $PROVER_SAVE_INPUTS_ON_ERROR rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$3" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
     if [ "$OUTPUT_TO_TERMINAL" = true ]; then
         eval $command
         echo -e "Proof generation finished with result: $?"
diff --git a/zero_bin/tools/prove_stdio.sh b/zero_bin/tools/prove_stdio.sh
index 0ae4e72a6..86b4d8bf8 100755
--- a/zero_bin/tools/prove_stdio.sh
+++ b/zero_bin/tools/prove_stdio.sh
@@ -91,6 +91,18 @@ else
 fi
 
 
+# Prover config. Override the defaults if needed by setting the env variables.
+PROVER_BATCH_SIZE="${PROVER_BATCH_SIZE:-1}"
+PROVER_SEGMENT_CHUNK_SIZE="${PROVER_SEGMENT_CHUNK_SIZE:-64}"
+PROVER_MAC_CPU_LEN_LOG="${PROVER_MAC_CPU_LEN_LOG:-20}"
+if [[ $PROVER_SAVE_INPUTS_ON_ERROR == "true" ]]; then
+    PROVER_SAVE_INPUTS_ON_ERROR="--save-inputs-on-error"
+else
+    PROVER_SAVE_INPUTS_ON_ERROR=""
+fi
+
+
+
 # If we run ./prove_stdio.sh <witness file name> test_only, we'll generate a dummy
 # proof. This is useful for quickly testing decoding and all of the
 # other non-proving code.
@@ -108,7 +120,9 @@ fi
 cargo build --release --jobs "$num_procs"
 
 start_time=$(date +%s%N)
-"${TOOLS_DIR}/../../target/release/leader" --runtime in-memory --load-strategy on-demand stdio < $INPUT_FILE | tee $LEADER_OUT_PATH
+"${TOOLS_DIR}/../../target/release/leader" --runtime in-memory --load-strategy on-demand --batch-size $PROVER_BATCH_SIZE \
+   --max-cpu-len-log $PROVER_MAC_CPU_LEN_LOG --segment-chunk-size $PROVER_SEGMENT_CHUNK_SIZE \
+   $PROVER_SAVE_INPUTS_ON_ERROR stdio < $INPUT_FILE | tee $LEADER_OUT_PATH
 end_time=$(date +%s%N)
 
 tail -n 1 $LEADER_OUT_PATH > $PROOFS_JSON_PATH