fix: discard intermediate proofs (#106)

* fix: discard intermediary proofs

* fix: write intermediate proofs on creation

* fix: return list of blocks proved

* fix: remove files

* fix: comment

* fix: reviews

* fix: rename jerigon params to proof params

* fix: error

* fix: refactor

* fix: comment

common/src/fs.rs

@@ -0,0 +1,7 @@
+use std::path::PathBuf;
+
+pub fn generate_block_proof_file_name(directory: &Option<&str>, block_height: u64) -> PathBuf {
+    let mut path = PathBuf::from(directory.unwrap_or(""));
+    path.push(format!("b{}.zkproof", block_height));
+    path
+}

common/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod block_interval;
 pub mod debug_utils;
+pub mod fs;
 pub mod parsing;
 pub mod prover_state;

leader/src/cli.rs

@@ -54,6 +54,15 @@ pub(crate) enum Command {
         /// to determine the blockchain node polling interval.
         #[arg(short, long, env = "ZERO_BIN_BLOCK_TIME", default_value_t = 2000)]
         block_time: u64,
+        /// Keep intermediate proofs. Default action is to
+        /// delete them after the final proof is generated.
+        #[arg(
+            short,
+            long,
+            env = "ZERO_BIN_KEEP_INTERMEDIATE_PROOFS",
+            default_value_t = false
+        )]
+        keep_intermediate_proofs: bool,
     },
     /// Reads input from HTTP and writes output to a directory.
     Http {

leader/src/jerigon.rs

@@ -1,61 +1,64 @@
-use std::{
-    fs::{create_dir_all, File},
-    io::Write,
-    path::PathBuf,
-};
+use std::path::PathBuf;
 
 use alloy::providers::RootProvider;
 use anyhow::Result;
 use common::block_interval::BlockInterval;
+use common::fs::generate_block_proof_file_name;
 use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
+use tracing::{error, warn};
+
+#[derive(Debug, Default)]
+pub struct ProofParams {
+    pub checkpoint_block_number: u64,
+    pub previous_proof: Option<GeneratedBlockProof>,
+    pub proof_output_dir: Option<PathBuf>,
+    pub save_inputs_on_error: bool,
+    pub keep_intermediate_proofs: bool,
+}
 
 /// The main function for the jerigon mode.
 pub(crate) async fn jerigon_main(
     runtime: Runtime,
     rpc_url: &str,
     block_interval: BlockInterval,
-    checkpoint_block_number: u64,
-    previous_proof: Option<GeneratedBlockProof>,
-    proof_output_dir_opt: Option<PathBuf>,
-    save_inputs_on_error: bool,
+    mut params: ProofParams,
 ) -> Result<()> {
     let prover_input = rpc::prover_input(
         RootProvider::new_http(rpc_url.parse()?),
         block_interval,
-        checkpoint_block_number.into(),
+        params.checkpoint_block_number.into(),
     )
     .await?;
 
-    let block_proofs = prover_input
-        .prove(&runtime, previous_proof, save_inputs_on_error)
+    // If `keep_intermediate_proofs` is not set we only keep the last block
+    // proof from the interval. It contains all the necessary information to
+    // verify the whole sequence.
+    let proved_blocks = prover_input
+        .prove(
+            &runtime,
+            params.previous_proof.take(),
+            params.save_inputs_on_error,
+            params.proof_output_dir.clone(),
+        )
         .await?;
     runtime.close().await?;
 
-    for block_proof in block_proofs {
-        let block_proof_str = serde_json::to_vec(&block_proof)?;
-        write_proof(
-            block_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_dir_opt: Option<PathBuf>) -> Result<()> {
-    match proof_output_dir_opt {
-        Some(p) => {
-            if let Some(parent) = p.parent() {
-                create_dir_all(parent)?;
-            }
-
-            let mut f = File::create(p)?;
-            f.write_all(&proof)?;
-        }
-        None => std::io::stdout().write_all(&proof)?,
+    if params.keep_intermediate_proofs {
+        warn!("Skipping cleanup, intermediate proofs are kept");
+    } else if let Some(proof_output_dir) = params.proof_output_dir.as_ref() {
+        proved_blocks
+            .into_iter()
+            .rev()
+            .skip(1)
+            .map(|b| generate_block_proof_file_name(&proof_output_dir.to_str(), b))
+            .for_each(|path| {
+                if let Err(e) = std::fs::remove_file(path) {
+                    error!("Failed to remove intermediate proof file: {e}");
+                }
+            });
+    } else {
+        // Proofs are written to stdio, so no need to clean up
     }
 
     Ok(())

leader/src/main.rs

@@ -11,6 +11,7 @@ use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
 use tracing::info;
 
+use crate::jerigon::{jerigon_main, ProofParams};
 use crate::utils::get_package_version;
 
 mod cli;
@@ -93,6 +94,7 @@ async fn main() -> Result<()> {
             proof_output_dir,
             save_inputs_on_error,
             block_time,
+            keep_intermediate_proofs,
         } => {
             let previous_proof = get_previous_proof(previous_proof)?;
             let mut block_interval = BlockInterval::new(&block_interval)?;
@@ -106,15 +108,17 @@ async fn main() -> Result<()> {
             }
 
             info!("Proving interval {block_interval}");
-
-            jerigon::jerigon_main(
+            jerigon_main(
                 runtime,
                 &rpc_url,
                 block_interval,
-                checkpoint_block_number,
-                previous_proof,
-                proof_output_dir,
-                save_inputs_on_error,
+                ProofParams {
+                    checkpoint_block_number,
+                    previous_proof,
+                    proof_output_dir,
+                    save_inputs_on_error,
+                    keep_intermediate_proofs,
+                },
             )
             .await?;
         }

prover/Cargo.toml

@@ -18,6 +18,7 @@ anyhow = { workspace = true }
 futures = { workspace = true }
 alloy.workspace = true
 tokio = {workspace = true}
+serde_json = {workspace = true}
 ruint = { version = "1.12.1", features = ["num-traits", "primitive-types"] }
 ops = { path = "../ops" }
 common = { path = "../common" }

prover/src/lib.rs

@@ -1,7 +1,9 @@
 use std::future::Future;
+use std::path::PathBuf;
 
-use alloy::primitives::U256;
-use anyhow::Result;
+use alloy::primitives::{BlockNumber, U256};
+use anyhow::{Context, Result};
+use common::fs::generate_block_proof_file_name;
 use futures::{future::BoxFuture, stream::FuturesOrdered, FutureExt, TryFutureExt, TryStreamExt};
 use num_traits::ToPrimitive as _;
 use ops::TxProof;
@@ -11,6 +13,7 @@ use paladin::{
 };
 use proof_gen::proof_types::GeneratedBlockProof;
 use serde::{Deserialize, Serialize};
+use tokio::io::AsyncWriteExt;
 use tokio::sync::oneshot;
 use trace_decoder::{
     processed_block_trace::ProcessingMeta,
@@ -133,7 +136,8 @@ impl ProverInput {
         runtime: &Runtime,
         previous_proof: Option<GeneratedBlockProof>,
         save_inputs_on_error: bool,
-    ) -> Result<Vec<GeneratedBlockProof>> {
+        proof_output_dir: Option<PathBuf>,
+    ) -> Result<Vec<BlockNumber>> {
         let mut prev: Option<BoxFuture<Result<GeneratedBlockProof>>> =
             previous_proof.map(|proof| Box::pin(futures::future::ok(proof)) as BoxFuture<_>);
 
@@ -147,16 +151,21 @@ impl ProverInput {
                 let (tx, rx) = oneshot::channel::<GeneratedBlockProof>();
 
                 // Prove the block
+                let proof_output_dir = proof_output_dir.clone();
                 let fut = block
                     .prove(runtime, prev.take(), save_inputs_on_error)
-                    .then(|proof| async {
+                    .then(move |proof| async move {
                         let proof = proof?;
+                        let block_number = proof.b_height;
 
-                        if tx.send(proof.clone()).is_err() {
+                        // Write latest generated proof to disk or stdout
+                        ProverInput::write_proof(proof_output_dir, &proof).await?;
+
+                        if tx.send(proof).is_err() {
                             anyhow::bail!("Failed to send proof");
                         }
 
-                        Ok(proof)
+                        Ok(block_number)
                     })
                     .boxed();
 
@@ -168,4 +177,30 @@ impl ProverInput {
 
         results.try_collect().await
     }
+
+    /// Write the proof to the disk (if `output_dir` is provided) or stdout.
+    pub(crate) async fn write_proof(
+        output_dir: Option<PathBuf>,
+        proof: &GeneratedBlockProof,
+    ) -> Result<()> {
+        let proof_serialized = serde_json::to_vec(proof)?;
+        let block_proof_file_path =
+            output_dir.map(|path| generate_block_proof_file_name(&path.to_str(), proof.b_height));
+        match block_proof_file_path {
+            Some(p) => {
+                if let Some(parent) = p.parent() {
+                    tokio::fs::create_dir_all(parent).await?;
+                }
+
+                let mut f = tokio::fs::File::create(p).await?;
+                f.write_all(&proof_serialized)
+                    .await
+                    .context("Failed to write proof to disk")
+            }
+            None => tokio::io::stdout()
+                .write_all(&proof_serialized)
+                .await
+                .context("Failed to write proof to stdout"),
+        }
+    }
 }