Simplify withdrawals logic

Cargo.toml

@@ -8,13 +8,13 @@ resolver = "2"
 ethers = "2.0.7"
 tokio = { version = "1", features = ["full"] }
 anyhow = "1.0.71"
-plonky2 = "0.2.1"
+plonky2 = "0.2.2"
 evm_arithmetization = { git = "https://github.com/0xPolygonZero/zk_evm.git", branch = "develop" }
 clap = { version = "4.4.11", features = ["derive", "env"] }
 trace_decoder = { git = "https://github.com/0xPolygonZero/zk_evm.git", branch = "develop" }
 proof_gen = { git = "https://github.com/0xPolygonZero/zk_evm.git", branch = "develop" }
 dotenvy = "0.15.7"
-paladin-core = "0.4.1"
+paladin-core = "0.4.2"
 serde = { version = "1.0.193", features = ["derive"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

leader/src/lib.rs

@@ -23,7 +23,7 @@ use padding_and_withdrawals::{
     add_withdrawals_to_txns, pad_gen_inputs_with_dummy_inputs_if_needed, BlockMetaAndHashes,
 };
 use rpc::{CliqueGetSignersAtHashResponse, EthChainIdResponse};
-use trace_decoder::types::{HashedAccountAddr, TrieRootHash, TxnProofGenIR};
+use trace_decoder::types::{HashedAccountAddr, TrieRootHash};
 
 use crate::utils::{has_storage_deletion, keccak};
 use crate::{
@@ -159,7 +159,7 @@ pub async fn gather_witness(
     tx: TxHash,
     provider: &Provider<Http>,
     request_miner_from_clique: bool,
-) -> Result<Vec<TxnProofGenIR>> {
+) -> Result<Vec<GenerationInputs>> {
     let tx = provider
         .get_transaction(tx)
         .await?
@@ -519,7 +519,7 @@ pub async fn gather_witness(
         })
         .unwrap_or_else(|| initial_extra_data.clone());
 
-    let dummies_added = pad_gen_inputs_with_dummy_inputs_if_needed(
+    pad_gen_inputs_with_dummy_inputs_if_needed(
         &mut proof_gen_ir,
         &b_data,
         &final_extra_data,
@@ -529,14 +529,7 @@ pub async fn gather_witness(
         !wds.is_empty(),
     );
 
-    add_withdrawals_to_txns(
-        &mut proof_gen_ir,
-        &b_data,
-        &final_extra_data,
-        &mut final_tries,
-        wds,
-        dummies_added,
-    );
+    add_withdrawals_to_txns(&mut proof_gen_ir, &mut final_tries, wds);
 
     Ok(proof_gen_ir)
 }

leader/src/main.rs

@@ -6,14 +6,14 @@ use clap::Parser;
 use common::prover_state::set_prover_state_from_config;
 use dotenvy::dotenv;
 use ethers::prelude::*;
+use evm_arithmetization::GenerationInputs;
 use leader::gather_witness;
 
 mod cli;
 mod prover;
 use cli::Command;
 use ops::register;
 use paladin::runtime::Runtime;
-use trace_decoder::types::TxnProofGenIR;
 mod init;
 
 #[tokio::main]
@@ -53,7 +53,7 @@ async fn main() -> Result<()> {
             let mut file = std::fs::File::open(&input_witness)?;
             let mut buffer = String::new();
             file.read_to_string(&mut buffer)?;
-            let proof_gen_ir: Vec<TxnProofGenIR> = serde_json::from_str(&buffer)?;
+            let proof_gen_ir: Vec<GenerationInputs> = serde_json::from_str(&buffer)?;
             let prover_input = prover::ProverInput { proof_gen_ir };
             let runtime = Runtime::from_config(&paladin, register()).await?;
             let proof = prover_input.prove(&runtime, None).await?;

leader/src/mpt.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use ethers::prelude::*;
 use ethers::utils::rlp;
 use evm_arithmetization::generation::mpt::AccountRlp;
-use mpt_trie::nibbles::Nibbles;
+use mpt_trie::nibbles::{Nibbles, NibblesIntern};
 use mpt_trie::partial_trie::PartialTrie;
 use mpt_trie::partial_trie::{HashedPartialTrie, Node};
 use mpt_trie::trie_subsets::create_trie_subset;
@@ -88,7 +88,7 @@ impl Mpt {
                     } else {
                         Nibbles {
                             count: 0,
-                            packed: U512::zero(),
+                            packed: NibblesIntern::zero(),
                         }
                     };
                     ext_prefix.push_nibble_front(b);
@@ -165,7 +165,7 @@ fn nibbles_from_hex_prefix_encoding(b: &[u8]) -> Nibbles {
             Nibbles {
                 count: 2 * b.len() - 1,
                 // packed: U512::from_bytes_be(&bs).unwrap(),
-                packed: U512::from_big_endian(&b),
+                packed: NibblesIntern::from_big_endian(&b),
             }
             // Nibbles::from_bytes_be(&b).unwrap()
         }

leader/src/padding_and_withdrawals.rs

@@ -1,7 +1,4 @@
-use std::{
-    collections::HashMap,
-    iter::{self, empty},
-};
+use std::{collections::HashMap, iter::empty};
 
 use ethers::{
     types::{Address, H256, U256},
@@ -16,7 +13,7 @@ use mpt_trie::partial_trie::PartialTrie;
 use mpt_trie::{
     nibbles::Nibbles, partial_trie::HashedPartialTrie, trie_subsets::create_trie_subset,
 };
-use trace_decoder::types::{HashedAccountAddr, HashedNodeAddr, TxnProofGenIR};
+use trace_decoder::types::HashedAccountAddr;
 
 use crate::{utils::keccak, PartialTrieState};
 
@@ -39,14 +36,14 @@ pub(crate) struct BlockMetaAndHashes {
 /// `[add_withdrawals_to_txns]`), where the final one will mutate the
 /// state trie.
 pub(crate) fn pad_gen_inputs_with_dummy_inputs_if_needed(
-    gen_inputs: &mut Vec<TxnProofGenIR>,
+    gen_inputs: &mut Vec<GenerationInputs>,
     other_data: &BlockMetaAndHashes,
     final_extra_data: &ExtraBlockData,
     initial_extra_data: &ExtraBlockData,
     initial_tries: &PartialTrieState,
     final_tries: &PartialTrieState,
     has_withdrawals: bool,
-) -> bool {
+) {
     match gen_inputs.len() {
         0 => {
             debug_assert!(initial_tries.state == final_tries.state);
@@ -57,8 +54,6 @@ pub(crate) fn pad_gen_inputs_with_dummy_inputs_if_needed(
                 final_extra_data,
                 initial_tries,
             ));
-
-            true
         }
         1 => {
             // We just need one dummy entry.
@@ -79,10 +74,8 @@ pub(crate) fn pad_gen_inputs_with_dummy_inputs_if_needed(
                     gen_inputs.push(dummy_txn)
                 }
             };
-
-            true
         }
-        _ => false,
+        _ => (),
     }
 }
 
@@ -94,12 +87,9 @@ pub(crate) fn pad_gen_inputs_with_dummy_inputs_if_needed(
 /// - If no dummy proofs are already present, then a dummy proof that just
 ///   contains the withdrawals is appended to the end of the IR vec.
 pub(crate) fn add_withdrawals_to_txns(
-    txn_ir: &mut Vec<TxnProofGenIR>,
-    other_data: &BlockMetaAndHashes,
-    final_extra_data: &ExtraBlockData,
+    txn_ir: &mut [GenerationInputs],
     final_trie_state: &mut PartialTrieState,
     withdrawals: Vec<(Address, U256)>,
-    dummies_already_added: bool,
 ) {
     if withdrawals.is_empty() {
         return;
@@ -109,50 +99,17 @@ pub(crate) fn add_withdrawals_to_txns(
         .iter()
         .map(|(addr, v)| (*addr, hash(addr.as_bytes()), *v));
 
-    match dummies_already_added {
-        // If we have no actual dummy proofs, then we create one and append it to the
-        // end of the block.
-        false => {
-            let withdrawals_with_hashed_addrs: Vec<_> =
-                withdrawals_with_hashed_addrs_iter.collect();
-
-            // Dummy state will be the state after the final txn. Also need to include the
-            // account nodes that were accessed by the withdrawals.
-            let withdrawal_addrs = withdrawals_with_hashed_addrs
-                .iter()
-                .cloned()
-                .map(|(_, h_addr, _)| h_addr);
-            let mut withdrawal_dummy = create_dummy_gen_input_with_state_addrs_accessed(
-                other_data,
-                final_extra_data,
-                final_trie_state,
-                withdrawal_addrs,
-            );
-
-            update_trie_state_from_withdrawals(
-                withdrawals_with_hashed_addrs,
-                &mut final_trie_state.state,
-            );
-
-            withdrawal_dummy.withdrawals = withdrawals;
-
-            // Only the state root hash needs to be updated from the withdrawals.
-            withdrawal_dummy.trie_roots_after.state_root = final_trie_state.state.hash();
+    update_trie_state_from_withdrawals(
+        withdrawals_with_hashed_addrs_iter,
+        &mut final_trie_state.state,
+    );
 
-            txn_ir.push(withdrawal_dummy);
-        }
-        true => {
-            update_trie_state_from_withdrawals(
-                withdrawals_with_hashed_addrs_iter,
-                &mut final_trie_state.state,
-            );
+    let last_inputs = txn_ir
+        .last_mut()
+        .expect("We cannot have an empty list of payloads.");
 
-            // If we have dummy proofs (note: `txn_ir[1]` is always a dummy txn in this
-            // case), then this dummy will get the withdrawals.
-            txn_ir[1].withdrawals = withdrawals;
-            txn_ir[1].trie_roots_after.state_root = final_trie_state.state.hash();
-        }
-    }
+    last_inputs.withdrawals = withdrawals;
+    last_inputs.trie_roots_after.state_root = final_trie_state.state.hash();
 }
 
 /// Withdrawals update balances in the account trie, so we need to update
@@ -199,23 +156,6 @@ fn create_dummy_gen_input(
     create_dummy_gen_input_common(other_data, extra_data, sub_tries)
 }
 
-fn create_dummy_gen_input_with_state_addrs_accessed(
-    other_data: &BlockMetaAndHashes,
-    extra_data: &ExtraBlockData,
-    tries: &PartialTrieState,
-    withdrawal_account_addrs_accessed: impl Iterator<Item = HashedAccountAddr>,
-) -> GenerationInputs {
-    let state_trie_hashed_for_withdrawals = create_minimal_state_partial_trie(
-        &tries.state,
-        withdrawal_account_addrs_accessed,
-        iter::empty(),
-    );
-
-    let sub_tries = create_dummy_proof_trie_inputs(tries, state_trie_hashed_for_withdrawals);
-
-    create_dummy_gen_input_common(other_data, extra_data, sub_tries)
-}
-
 fn create_dummy_gen_input_common(
     other_data: &BlockMetaAndHashes,
     extra_data: &ExtraBlockData,
@@ -279,20 +219,6 @@ fn create_dummy_proof_trie_inputs(
     }
 }
 
-fn create_minimal_state_partial_trie(
-    state_trie: &HashedPartialTrie,
-    state_accesses: impl Iterator<Item = HashedNodeAddr>,
-    additional_state_trie_paths_to_not_hash: impl Iterator<Item = Nibbles>,
-) -> HashedPartialTrie {
-    let accesses = state_accesses
-        .into_iter()
-        .map(Nibbles::from_h256_be)
-        .chain(additional_state_trie_paths_to_not_hash);
-
-    create_trie_subset(state_trie, accesses)
-        .expect("Encountered a hash node when creating a subset of the state trie")
-}
-
 // We really want to get a trie with just a hash node here, and this is an easy
 // way to do it.
 fn create_fully_hashed_out_sub_partial_trie(trie: &HashedPartialTrie) -> HashedPartialTrie {

leader/src/prover.rs

@@ -1,4 +1,5 @@
 use anyhow::{bail, Result};
+use evm_arithmetization::GenerationInputs;
 use ops::{AggProof, BlockProof, TxProof};
 use paladin::{
     directive::{Directive, IndexedStream, Literal},
@@ -9,11 +10,10 @@ use proof_gen::{
     types::PlonkyProofIntern,
 };
 use serde::{Deserialize, Serialize};
-use trace_decoder::types::TxnProofGenIR;
 
 #[derive(Debug, Deserialize, Serialize)]
 pub struct ProverInput {
-    pub proof_gen_ir: Vec<TxnProofGenIR>,
+    pub proof_gen_ir: Vec<GenerationInputs>,
 }
 
 impl ProverInput {

ops/Cargo.toml

@@ -11,6 +11,7 @@ proof_gen = { workspace = true }
 trace_decoder = { workspace = true }
 tracing = { workspace = true }
 ethers = { workspace = true }
+evm_arithmetization = { workspace = true }
 
 rlp = "0.5.2"
 

ops/src/lib.rs

@@ -1,5 +1,6 @@
 use common::prover_state::P_STATE;
 use ethers::types::Transaction;
+use evm_arithmetization::GenerationInputs;
 use paladin::{
     operation::{FatalError, Monoid, Operation, Result},
     registry, RemoteExecute,
@@ -10,7 +11,6 @@ use proof_gen::{
     prover_state::ProverState,
 };
 use serde::{Deserialize, Serialize};
-use trace_decoder::types::TxnProofGenIR;
 use tracing::{info_span, Level};
 
 fn p_state() -> &'static ProverState {
@@ -23,7 +23,7 @@ registry!();
 pub struct TxProof;
 
 impl Operation for TxProof {
-    type Input = TxnProofGenIR;
+    type Input = GenerationInputs;
     type Output = AggregatableProof;
 
     fn execute(&self, input: Self::Input) -> Result<Self::Output> {