feat: extract decoded transaction from the block (#504)

* feat: extract transaction from block interval

* fix: rename extract-tx to extract subcommand

* fix: output GenerationInputs

* fix: refactor extract command

* fix: comment

* fix: error printout

* fix: remove clone

* fix: update for continuations

* fix: remove accidental file

* fix: full sentences

* fix: full sentences 2

zero_bin/README.md

@@ -230,7 +230,7 @@ Options:
   -k, --keep-intermediate-proofs
           Keep intermediate proofs. Default action is to delete them after the final proof is generated [env: ZERO_BIN_KEEP_INTERMEDIATE_PROOFS=]
       --backoff <BACKOFF>
-          Backoff in milliseconds for request retries [default: 0]
+          Backoff in milliseconds for retry requests [default: 0]
       --max-retries <MAX_RETRIES>
           The maximum number of retries [default: 0]
   -h, --help
@@ -272,7 +272,7 @@ Options:
   -k, --keep-intermediate-proofs
           Keep intermediate proofs. Default action is to delete them after the final proof is generated [env: ZERO_BIN_KEEP_INTERMEDIATE_PROOFS=]
       --backoff <BACKOFF>
-          Backoff in milliseconds for request retries [default: 0]
+          Backoff in milliseconds for retry requests [default: 0]
       --max-retries <MAX_RETRIES>
           The maximum number of retries [default: 0]
   -h, --help

zero_bin/leader/src/cli.rs

@@ -68,7 +68,7 @@ pub(crate) enum Command {
             default_value_t = false
         )]
         keep_intermediate_proofs: bool,
-        /// Backoff in milliseconds for request retries
+        /// Backoff in milliseconds for retry requests
         #[arg(long, default_value_t = 0)]
         backoff: u64,
         /// The maximum number of retries

zero_bin/leader/src/client.rs

@@ -52,7 +52,7 @@ pub(crate) async fn client_main(
         &params.previous_proof,
         block_interval.get_start_block()?,
     )?;
-    // Grab interval checkpoint block state trie
+    // Grab interval checkpoint block state trie.
     let checkpoint_state_trie_root = cached_provider
         .get_block(
             params.checkpoint_block_number.into(),

zero_bin/rpc/src/lib.rs

@@ -34,6 +34,14 @@ pub enum RpcType {
     Native,
 }
 
+#[derive(Clone, Debug, Copy)]
+pub struct RpcParams {
+    pub start_block: u64,
+    pub end_block: u64,
+    pub checkpoint_block_number: Option<u64>,
+    pub rpc_type: RpcType,
+}
+
 /// Obtain the prover input for one block
 pub async fn block_prover_input<ProviderT, TransportT>(
     cached_provider: Arc<CachedProvider<ProviderT, TransportT>>,

zero_bin/rpc/src/main.rs

@@ -1,98 +1,201 @@
 use std::env;
 use std::sync::Arc;
 
+use alloy::primitives::B256;
+use alloy::providers::Provider;
 use alloy::rpc::types::eth::BlockId;
 use alloy::rpc::types::{BlockNumberOrTag, BlockTransactionsKind};
-use clap::{Parser, ValueHint};
+use alloy::transports::Transport;
+use anyhow::anyhow;
+use clap::{Args, Parser, Subcommand, ValueHint};
 use futures::StreamExt;
+use prover::BlockProverInput;
 use rpc::provider::CachedProvider;
-use rpc::{retry::build_http_retry_provider, RpcType};
+use rpc::{retry::build_http_retry_provider, RpcParams, RpcType};
 use tracing_subscriber::{prelude::*, EnvFilter};
 use url::Url;
 use zero_bin_common::pre_checks::check_previous_proof_and_checkpoint;
 use zero_bin_common::version;
 use zero_bin_common::{block_interval::BlockInterval, prover_state::persistence::CIRCUIT_VERSION};
 
-#[derive(Parser)]
-pub enum Cli {
-    /// Fetch and generate prover input from the RPC endpoint
+#[derive(Args, Clone, Debug)]
+pub(crate) struct RpcConfig {
+    /// The RPC URL.
+    #[arg(short = 'u', long, value_hint = ValueHint::Url)]
+    rpc_url: Url,
+    /// The RPC Tracer Type.
+    #[arg(short = 't', long, default_value = "jerigon")]
+    rpc_type: RpcType,
+    /// Backoff in milliseconds for retry requests.
+    #[arg(long, default_value_t = 0)]
+    backoff: u64,
+    /// The maximum number of retries.
+    #[arg(long, default_value_t = 0)]
+    max_retries: u32,
+}
+
+#[derive(Subcommand)]
+pub(crate) enum Command {
     Fetch {
-        // Starting block of interval to fetch
+        /// Starting block of interval to fetch.
         #[arg(short, long)]
         start_block: u64,
-        // End block of interval to fetch
+        /// End block of interval to fetch.
         #[arg(short, long)]
         end_block: u64,
-        /// The RPC URL.
-        #[arg(short = 'u', long, value_hint = ValueHint::Url)]
-        rpc_url: Url,
-        /// The RPC Tracer Type
-        #[arg(short = 't', long, default_value = "jerigon")]
-        rpc_type: RpcType,
         /// The checkpoint block number. If not provided,
-        /// block before the `start_block` is the checkpoint
+        /// the block before the `start_block` is the checkpoint.
         #[arg(short, long)]
         checkpoint_block_number: Option<u64>,
-        /// Backoff in milliseconds for request retries
-        #[arg(long, default_value_t = 0)]
-        backoff: u64,
-        /// The maximum number of retries
-        #[arg(long, default_value_t = 0)]
-        max_retries: u32,
     },
+    Extract {
+        /// Transaction hash.
+        #[arg(long, short)]
+        tx: String,
+        /// Number of transactions in a batch to process at once.
+        #[arg(short, long, default_value_t = 1)]
+        batch_size: usize,
+    },
+}
+
+#[derive(Parser)]
+pub(crate) struct Cli {
+    #[clap(flatten)]
+    pub(crate) config: RpcConfig,
+
+    /// Fetch and generate prover input from the RPC endpoint.
+    #[command(subcommand)]
+    pub(crate) command: Command,
+}
+
+pub(crate) async fn retrieve_block_prover_inputs<ProviderT, TransportT>(
+    cached_provider: Arc<CachedProvider<ProviderT, TransportT>>,
+    params: RpcParams,
+) -> Result<Vec<BlockProverInput>, anyhow::Error>
+where
+    ProviderT: Provider<TransportT>,
+    TransportT: Transport + Clone,
+{
+    let checkpoint_block_number = params
+        .checkpoint_block_number
+        .unwrap_or(params.start_block - 1);
+    check_previous_proof_and_checkpoint(checkpoint_block_number, &None, params.start_block)?;
+
+    // Grab interval checkpoint block state trie
+    let checkpoint_state_trie_root = cached_provider
+        .get_block(
+            BlockId::Number(checkpoint_block_number.into()),
+            BlockTransactionsKind::Hashes,
+        )
+        .await?
+        .header
+        .state_root;
+
+    let block_interval = BlockInterval::Range(params.start_block..params.end_block + 1);
+    let mut block_prover_inputs = Vec::new();
+    let mut block_interval = block_interval.clone().into_bounded_stream()?;
+    while let Some(block_num) = block_interval.next().await {
+        let block_id = BlockId::Number(BlockNumberOrTag::Number(block_num));
+        // Get the prover input for particular block.
+        let result = rpc::block_prover_input(
+            cached_provider.clone(),
+            block_id,
+            checkpoint_state_trie_root,
+            params.rpc_type,
+        )
+        .await?;
+
+        block_prover_inputs.push(result);
+    }
+    Ok(block_prover_inputs)
 }
 
 impl Cli {
     /// Execute the cli command.
     pub async fn execute(self) -> anyhow::Result<()> {
-        match self {
-            Self::Fetch {
+        let cached_provider = Arc::new(CachedProvider::new(build_http_retry_provider(
+            self.config.rpc_url.clone(),
+            self.config.backoff,
+            self.config.max_retries,
+        )));
+
+        match self.command {
+            Command::Fetch {
                 start_block,
                 end_block,
-                rpc_url,
-                rpc_type,
                 checkpoint_block_number,
-                backoff,
-                max_retries,
             } => {
-                let checkpoint_block_number = checkpoint_block_number.unwrap_or(start_block - 1);
-                check_previous_proof_and_checkpoint(checkpoint_block_number, &None, start_block)?;
-                let block_interval = BlockInterval::Range(start_block..end_block + 1);
-
-                let cached_provider = Arc::new(CachedProvider::new(build_http_retry_provider(
-                    rpc_url.clone(),
-                    backoff,
-                    max_retries,
-                )));
-
-                // Grab interval checkpoint block state trie
-                let checkpoint_state_trie_root = cached_provider
-                    .get_block(
-                        BlockId::Number(checkpoint_block_number.into()),
-                        BlockTransactionsKind::Hashes,
-                    )
-                    .await?
-                    .header
-                    .state_root;
-
-                let mut block_prover_inputs = Vec::new();
-                let mut block_interval = block_interval.clone().into_bounded_stream()?;
-                while let Some(block_num) = block_interval.next().await {
-                    let block_id = BlockId::Number(BlockNumberOrTag::Number(block_num));
-                    // Get the prover input for particular block.
-                    let result = rpc::block_prover_input(
-                        cached_provider.clone(),
-                        block_id,
-                        checkpoint_state_trie_root,
-                        rpc_type,
-                    )
-                    .await?;
-
-                    block_prover_inputs.push(result);
-                }
+                let params = RpcParams {
+                    start_block,
+                    end_block,
+                    checkpoint_block_number,
+                    rpc_type: self.config.rpc_type,
+                };
 
+                let block_prover_inputs =
+                    retrieve_block_prover_inputs(cached_provider, params).await?;
                 serde_json::to_writer_pretty(std::io::stdout(), &block_prover_inputs)?;
             }
+            Command::Extract { tx, batch_size } => {
+                let tx_hash: B256 = tx.parse()?;
+                // Get transaction info.
+                match cached_provider
+                    .clone()
+                    .as_provider()
+                    .get_transaction_by_hash(tx_hash)
+                    .await?
+                {
+                    Some(tx_info) => {
+                        let block_number = tx_info.block_number.ok_or(anyhow!(
+                            "transaction {} does not have block number",
+                            tx_hash
+                        ))?;
+                        let params = RpcParams {
+                            start_block: block_number,
+                            end_block: block_number,
+                            checkpoint_block_number: None,
+                            rpc_type: self.config.rpc_type,
+                        };
+
+                        let block_prover_inputs =
+                            retrieve_block_prover_inputs(cached_provider, params).await?;
+
+                        let block_prover_input =
+                            block_prover_inputs.into_iter().next().ok_or(anyhow!(
+                                "error, block prover input for block {} not retrieved",
+                                block_number
+                            ))?;
+
+                        let generation_inputs = trace_decoder::entrypoint(
+                            block_prover_input.block_trace,
+                            block_prover_input.other_data,
+                            batch_size,
+                        )?;
+
+                        if let Some(index) = tx_info.transaction_index {
+                            let generation_input_index = if batch_size == 1 {
+                                // If batch size is 1, it means one transaction per
+                                // GenerationInputs. Take element
+                                // with txn index from the GenerationInput array.
+                                index as usize
+                            } else {
+                                // Batch size bigger than one, meaning multiple transactions in one
+                                // GenerationInput. Find GenerationInput
+                                // where the transaction is placed.
+                                index as usize / batch_size
+                            };
+                            let extracted_generation_input =
+                                generation_inputs.get(generation_input_index).cloned();
+                            serde_json::to_writer(std::io::stdout(), &extracted_generation_input)?;
+                        } else {
+                            anyhow::bail!("invalid transaction index for transaction {}", tx_hash);
+                        }
+                    }
+                    None => {
+                        anyhow::bail!("unable to find transaction {}", tx_hash);
+                    }
+                }
+            }
         }
         Ok(())
     }