fix: proving logic

common/src/block_interval.rs

@@ -105,7 +105,7 @@ impl BlockInterval {
     /// Convert the block interval into a stream of block numbers.
     pub fn into_bounded_stream(self) -> Result<impl Stream<Item = u64>, BlockIntervalError> {
         match self {
-            BlockInterval::Single(num) => Ok(tokio_stream::iter(num..num)),
+            BlockInterval::Single(num) => Ok(tokio_stream::iter(num..num + 1)),
             BlockInterval::Range(range) => Ok(tokio_stream::iter(range)),
             _ => Err(BlockIntervalError::IntoBoundedStreamError),
         }
@@ -238,4 +238,18 @@ mod test {
             )))
         )
     }
+
+    #[tokio::test]
+    async fn can_into_bounded_stream() {
+        use tokio_stream::StreamExt;
+        let mut result = Vec::new();
+        let mut stream = BlockInterval::new("1..10")
+            .unwrap()
+            .into_bounded_stream()
+            .unwrap();
+        while let Some(val) = stream.next().await {
+            result.push(val);
+        }
+        assert_eq!(result, Vec::from_iter(1u64..10u64));
+    }
 }

leader/src/cli.rs

@@ -46,10 +46,10 @@ pub(crate) enum Command {
         /// The previous proof output.
         #[arg(long, short = 'f', value_hint = ValueHint::FilePath)]
         previous_proof: Option<PathBuf>,
-        /// If provided, write the generated proof to this file instead of
+        /// If provided, write the generated proofs to this directory instead of
         /// stdout.
         #[arg(long, short = 'o', value_hint = ValueHint::FilePath)]
-        proof_output_path: Option<PathBuf>,
+        proof_output_dir: Option<PathBuf>,
         /// If true, save the public inputs to disk on error.
         #[arg(short, long, default_value_t = false)]
         save_inputs_on_error: bool,

leader/src/jerigon.rs

@@ -18,7 +18,7 @@ pub(crate) async fn jerigon_main(
     block_interval: BlockInterval,
     checkpoint_block_number: u64,
     previous_proof: Option<PlonkyProofIntern>,
-    proof_output_path_opt: Option<PathBuf>,
+    proof_output_dir_opt: Option<PathBuf>,
     save_inputs_on_error: bool,
 ) -> Result<()> {
     let prover_input = rpc::prover_input(
@@ -39,21 +39,20 @@ pub(crate) async fn jerigon_main(
     runtime.close().await?;
 
     for block_proof in block_proofs {
-        //todo fix proof dump to disk (for every block there it may be different file
-        // in one directory)
+        let blokck_proof_str = serde_json::to_vec(&block_proof.intern)?;
         write_proof(
-            serde_json::to_vec(&block_proof.intern)?,
-            proof_output_path_opt.clone().map(|mut path| {
-                path.push(format!("_block_{}", block_proof.b_height));
+            blokck_proof_str,
+            proof_output_dir_opt.clone().map(|mut path| {
+                path.push(format!("b{}.zkproof", block_proof.b_height));
                 path
             }),
         )?;
     }
     Ok(())
 }
 
-fn write_proof(proof: Vec<u8>, proof_output_path_opt: Option<PathBuf>) -> Result<()> {
-    match proof_output_path_opt {
+fn write_proof(proof: Vec<u8>, proof_output_dir_opt: Option<PathBuf>) -> Result<()> {
+    match proof_output_dir_opt {
         Some(p) => {
             if let Some(parent) = p.parent() {
                 create_dir_all(parent)?;

leader/src/main.rs

@@ -95,7 +95,7 @@ async fn main() -> Result<()> {
             block_interval,
             checkpoint_block_number,
             previous_proof,
-            proof_output_path,
+            proof_output_dir,
             save_inputs_on_error,
         } => {
             let previous_proof = get_previous_proof(previous_proof)?;
@@ -118,7 +118,7 @@ async fn main() -> Result<()> {
                 block_interval,
                 checkpoint_block_number,
                 previous_proof,
-                proof_output_path,
+                proof_output_dir,
                 save_inputs_on_error,
             )
             .await?;

prover/src/lib.rs

@@ -130,40 +130,43 @@ pub struct ProverInput {
 }
 
 impl ProverInput {
-    #[cfg(not(feature = "test_only"))]
     pub async fn prove(
         self,
         runtime: &Runtime,
         checkpoint_proof: Option<GeneratedBlockProof>,
         save_inputs_on_error: bool,
     ) -> Result<Vec<GeneratedBlockProof>> {
-        let results: Arc<Mutex<HashMap<u64, PlonkyProofIntern>>> =
+        let results: Arc<Mutex<HashMap<u64, Option<PlonkyProofIntern>>>> =
             Arc::new(Mutex::new(HashMap::new()));
         if let Some(checkpoint_proof) = checkpoint_proof {
             results
                 .lock()
                 .await
-                .insert(checkpoint_proof.b_height, checkpoint_proof.intern);
+                .insert(checkpoint_proof.b_height, Some(checkpoint_proof.intern));
         };
 
         for block in self.blocks {
-            //todo this will be further reorganized with the new BlockProofFuture
-
+            //todo this will be further reorganized with the new BlockProofFuture,
+            // running multiple block proof generation in parallel, and awaiting for the
+            // previous block proof when needed to prove the next block.
+            // For now prove blocks one by one sequentially and assume previous proof is
+            // already available
             let results = results.clone();
-            // for now prove blocks one by one
             async move {
-                //todo handle genesis block case for previous block number
-                let previous_block_number = block
-                    .get_block_number()
+                let block_number = block.get_block_number();
+                info!("Proving block {block_number}");
+                // For genesis block we don't have a previous proof, so
+                // previous_block_number would be None
+                let previous_block_number: Option<u64> = block_number
                     .checked_sub(U256::from(1))
-                    .ok_or(anyhow::Error::msg("block number smaller than 1"))?
-                    .to_u64();
-                let previous_proof = if let Some(previous_block_number) = previous_block_number {
-                    //todo optimize previous block proof clone
-                    results.lock().await.get(&previous_block_number).cloned()
+                    .and_then(|n| n.to_u64());
+                let previous_proof = if let Some(number) = previous_block_number {
+                    //todo could we optimize this to avoid this large proof cloning?
+                    results.lock().await.get(&number).cloned().flatten()
                 } else {
                     None
                 };
+                // Prove the block
                 let block_proof = block
                     .prove(runtime, previous_proof, save_inputs_on_error)
                     .await?;
@@ -172,34 +175,25 @@ impl ProverInput {
                         .b_height
                         .to_u64()
                         .ok_or(anyhow::Error::msg("block number u64 overflow"))?,
-                    block_proof.intern,
+                    Some(block_proof.intern),
                 );
+                info!("Proving block {block_number} finished!");
                 Result::<(), anyhow::Error>::Ok(())
             }
             .await?;
         }
 
         let mut results = results.lock().await;
-        Ok(results
+        results
             .drain()
-            .map(|(block_number, intern)| GeneratedBlockProof {
-                b_height: block_number,
-                intern,
+            .map(|(block_number, intern)| {
+                Ok::<GeneratedBlockProof, anyhow::Error>(GeneratedBlockProof {
+                    b_height: block_number,
+                    intern: intern.ok_or_else(|| {
+                        anyhow::Error::msg("missing proof for block {block_number}")
+                    })?,
+                })
             })
-            .collect())
-    }
-
-    #[cfg(feature = "test_only")]
-    pub async fn prove(
-        self,
-        runtime: &Runtime,
-        _previous: Option<PlonkyProofIntern>,
-        save_inputs_on_error: bool,
-    ) -> Result<Vec<GeneratedBlockProof>> {
-        //todo test and update this
-        self.blocks
-            .into_iter()
-            .map(|block| block.prove(runtime, None, save_inputs_on_error))?
-            .collect()
+            .collect::<Result<Vec<_>, _>>()
     }
 }

tools/prove_blocks.sh

@@ -5,48 +5,94 @@
 # 2 --> End block index (inclusive)
 # 3 --> Rpc endpoint:port (eg. http://35.246.1.96:8545)
 # 4 --> Ignore previous proofs (boolean)
+# 5 --> Test only flag
 
+export RUST_MIN_STACK=33554432
 export RUST_BACKTRACE=1
-export RUST_LOG=mpt_trie=info,trace_decoder=info,plonky2=info,evm_arithmetization=trace,leader=info
-export RUSTFLAGS='-Ctarget-cpu=native'
+export RUST_LOG=info
+# Disable the lld linker for now, as it's causing issues with the linkme package.
+# https://github.com/rust-lang/rust/pull/124129
+# https://github.com/dtolnay/linkme/pull/88
+export RUSTFLAGS='-C target-cpu=native -Zlinker-features=-lld'
 
-export ARITHMETIC_CIRCUIT_SIZE="16..23"
-export BYTE_PACKING_CIRCUIT_SIZE="9..21"
-export CPU_CIRCUIT_SIZE="12..25"
-export KECCAK_CIRCUIT_SIZE="14..20"
-export KECCAK_SPONGE_CIRCUIT_SIZE="9..15"
-export LOGIC_CIRCUIT_SIZE="12..18"
-export MEMORY_CIRCUIT_SIZE="17..28"
+if [[ $5 == "test_only" ]]; then
+  # Circuit sizes don't matter in test_only mode, so we keep them minimal.
+  export ARITHMETIC_CIRCUIT_SIZE="16..17"
+  export BYTE_PACKING_CIRCUIT_SIZE="9..10"
+  export CPU_CIRCUIT_SIZE="12..13"
+  export KECCAK_CIRCUIT_SIZE="14..15"
+  export KECCAK_SPONGE_CIRCUIT_SIZE="9..10"
+  export LOGIC_CIRCUIT_SIZE="12..13"
+  export MEMORY_CIRCUIT_SIZE="17..18"
+else
+  export ARITHMETIC_CIRCUIT_SIZE="16..23"
+  export BYTE_PACKING_CIRCUIT_SIZE="9..21"
+  export CPU_CIRCUIT_SIZE="12..25"
+  export KECCAK_CIRCUIT_SIZE="14..20"
+  export KECCAK_SPONGE_CIRCUIT_SIZE="9..15"
+  export LOGIC_CIRCUIT_SIZE="12..18"
+  export MEMORY_CIRCUIT_SIZE="17..28"
+fi
 
 PROOF_OUTPUT_DIR="proofs"
+OUT_LOG_PATH="${PROOF_OUTPUT_DIR}/b${i}.log"
 ALWAYS_WRITE_LOGS=0 # Change this to `1` if you always want logs to be written.
-
 TOT_BLOCKS=$(($2-$1+1))
+
+START_BLOCK=$1
+END_BLOCK=$2
+NODE_RPC_URL=$3
 IGNORE_PREVIOUS_PROOFS=$4
 
-echo "Proving blocks ${1}..=${2}... (Total: ${TOT_BLOCKS})"
+
 mkdir -p $PROOF_OUTPUT_DIR
 
-for ((i=$1; i<=$2; i++))
-do
-    echo "Proving block ${i}..."
 
-    OUT_PROOF_PATH="${PROOF_OUTPUT_DIR}/b${i}.zkproof"
-    OUT_LOG_PATH="${PROOF_OUTPUT_DIR}/b${i}.log"
+if [ $IGNORE_PREVIOUS_PROOFS ]; then
+    # Set checkpoint height to previous block number for the first block in range
+    prev_proof_num=$(($1-1))
+    PREV_PROOF_EXTRA_ARG="--checkpoint-block-number ${prev_proof_num}"
+else
+    if [ $1 -gt 1 ]; then
+        prev_proof_num=$(($1-1))
+        PREV_PROOF_EXTRA_ARG="-f ${PROOF_OUTPUT_DIR}/b${prev_proof_num}.zkproof"
+    fi
+fi
+
+# Convert hex do decimal parameters
+if [[ $START_BLOCK == 0x* ]]; then
+    START_BLOCK=$((16#${START_BLOCK#"0x"}))
+fi
+if [[ $END_BLOCK == 0x* ]]; then
+    END_BLOCK=$((16#${END_BLOCK#"0x"}))
+fi
+
+if [ $START_BLOCK == $END_BLOCK ]; then
+      BLOCK_INTERVAL=$((16#${START_BLOCK#"0x"}))
+else
+    BLOCK_INTERVAL=$START_BLOCK..=$END_BLOCK
+fi
 
-    if [ $IGNORE_PREVIOUS_PROOFS ]; then
-        # Set checkpoint height to previous block number
-        prev_proof_num=$((i-1))
-        PREV_PROOF_EXTRA_ARG="--checkpoint-block-number ${prev_proof_num}"
+
+# If we set test_only flag, we'll generate a dummy
+# proof. This is useful for quickly testing decoding and all of the
+# other non-proving code.
+if [[ $5 == "test_only" ]]; then
+    # test only run
+    echo "Proving blocks ${BLOCK_INTERVAL} in a test_only mode now... (Total: ${TOT_BLOCKS})"
+    cargo r --release --features test_only --bin leader -- --runtime in-memory --load-strategy on-demand jerigon --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG | tee test-jerigon.out
+    if grep 'Successfully generated witness for block' test-jerigon.out; then
+        echo "Success - Note this was just a test, not a proof"
+        exit
     else
-        if [ $i -gt 1 ]; then
-            prev_proof_num=$((i-1))
-            PREV_PROOF_EXTRA_ARG="-f ${PROOF_OUTPUT_DIR}/b${prev_proof_num}.zkproof"
-        fi
+        echo "Failed to create a witness"
+        exit 1
     fi
+else
+    # normal run
+    echo "Proving blocks ${BLOCK_INTERVAL} now... (Total: ${TOT_BLOCKS})"
+    cargo r --release --bin leader -- --runtime in-memory --load-strategy on-demand jerigon --rpc-url "$3" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG > $OUT_LOG_PATH 2>&1
 
-    cargo r --release --bin leader -- --runtime in-memory jerigon --rpc-url "$3" --block-number $i --proof-output-path $OUT_PROOF_PATH $PREV_PROOF_EXTRA_ARG > $OUT_LOG_PATH 2>&1
-
     retVal=$?
     if [ $retVal -ne 0 ]; then
         # Some error occured.
@@ -58,6 +104,10 @@ do
             rm $OUT_LOG_PATH
         fi
     fi
-done
 
-echo "Successfully generated ${TOT_BLOCKS} proofs!"
+    echo "Successfully generated ${TOT_BLOCKS} proofs!"
+fi
+
+
+
+