From 24439f7ed578edeedec914ca29c53b7babad7597 Mon Sep 17 00:00:00 2001
From: Yuwen Zhang <yuwen01@gmail.com>
Date: Wed, 11 Sep 2024 17:26:35 -0700
Subject: [PATCH 1/2] make mpt mod public

---
 crates/mpt/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/mpt/src/lib.rs b/crates/mpt/src/lib.rs
index d7977a2..4cbdfeb 100644
--- a/crates/mpt/src/lib.rs
+++ b/crates/mpt/src/lib.rs
@@ -4,7 +4,7 @@ use revm::primitives::{Address, HashMap, B256};
 use serde::{Deserialize, Serialize};
 
 /// Module containing MPT code adapted from `zeth`.
-mod mpt;
+pub mod mpt;
 use mpt::{proofs_to_tries, MptNode};
 
 /// Ethereum state trie and account storage tries.

From 5703b08ef7cc5edf12b735564ce4e1476e5b3e07 Mon Sep 17 00:00:00 2001
From: Yuwen Zhang <yuwen01@gmail.com>
Date: Wed, 11 Sep 2024 17:47:09 -0700
Subject: [PATCH 2/2] new method to EthereumState

---
 crates/executor/host/src/lib.rs |  2 +-
 crates/mpt/src/lib.rs           | 14 +++++--
 crates/mpt/src/mpt.rs           | 65 +++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/crates/executor/host/src/lib.rs b/crates/executor/host/src/lib.rs
index 7e190f5..43b8568 100644
--- a/crates/executor/host/src/lib.rs
+++ b/crates/executor/host/src/lib.rs
@@ -150,7 +150,7 @@ impl<T: Transport + Clone, P: Provider<T, AnyNetwork> + Clone> HostExecutor<T, P
             after_storage_proofs.push(eip1186_proof_to_account_proof(storage_proof));
         }
 
-        let state = EthereumState::from_proofs(
+        let state = EthereumState::from_transition_proofs(
             previous_block.state_root,
             &before_storage_proofs.iter().map(|item| (item.address, item.clone())).collect(),
             &after_storage_proofs.iter().map(|item| (item.address, item.clone())).collect(),
diff --git a/crates/mpt/src/lib.rs b/crates/mpt/src/lib.rs
index 4cbdfeb..b9e3828 100644
--- a/crates/mpt/src/lib.rs
+++ b/crates/mpt/src/lib.rs
@@ -4,8 +4,8 @@ use revm::primitives::{Address, HashMap, B256};
 use serde::{Deserialize, Serialize};
 
 /// Module containing MPT code adapted from `zeth`.
-pub mod mpt;
-use mpt::{proofs_to_tries, MptNode};
+mod mpt;
+use mpt::{proofs_to_tries, transition_proofs_to_tries, MptNode};
 
 /// Ethereum state trie and account storage tries.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -16,12 +16,18 @@ pub struct EthereumState {
 
 impl EthereumState {
     /// Builds Ethereum state tries from relevant proofs before and after a state transition.
-    pub fn from_proofs(
+    pub fn from_transition_proofs(
         state_root: B256,
         parent_proofs: &HashMap<Address, AccountProof>,
         proofs: &HashMap<Address, AccountProof>,
     ) -> Result<Self> {
-        proofs_to_tries(state_root, parent_proofs, proofs).map_err(|err| eyre::eyre!("{}", err))
+        transition_proofs_to_tries(state_root, parent_proofs, proofs)
+            .map_err(|err| eyre::eyre!("{}", err))
+    }
+
+    /// Builds Ethereum state tries from relevant proofs from a given state.
+    pub fn from_proofs(state_root: B256, proofs: &HashMap<Address, AccountProof>) -> Result<Self> {
+        proofs_to_tries(state_root, proofs).map_err(|err| eyre::eyre!("{}", err))
     }
 
     /// Mutates state based on diffs provided in [`HashedPostState`].
diff --git a/crates/mpt/src/mpt.rs b/crates/mpt/src/mpt.rs
index fbddb4e..8790edd 100644
--- a/crates/mpt/src/mpt.rs
+++ b/crates/mpt/src/mpt.rs
@@ -962,6 +962,71 @@ pub fn shorten_node_path(node: &MptNode) -> Vec<MptNode> {
 }
 
 pub fn proofs_to_tries(
+    state_root: B256,
+    proofs: &HashMap<Address, AccountProof>,
+) -> Result<EthereumState> {
+    // if no addresses are provided, return the trie only consisting of the state root
+    if proofs.is_empty() {
+        return Ok(EthereumState {
+            state_trie: node_from_digest(state_root),
+            storage_tries: HashMap::new(),
+        });
+    }
+
+    let mut storage: HashMap<B256, MptNode> = HashMap::with_capacity(proofs.len());
+
+    let mut state_nodes = HashMap::new();
+    let mut state_root_node = MptNode::default();
+    for (address, proof) in proofs {
+        let proof_nodes = parse_proof(&proof.proof).unwrap();
+        mpt_from_proof(&proof_nodes).unwrap();
+
+        // the first node in the proof is the root
+        if let Some(node) = proof_nodes.first() {
+            state_root_node = node.clone();
+        }
+
+        proof_nodes.into_iter().for_each(|node| {
+            state_nodes.insert(node.reference(), node);
+        });
+
+        // if no slots are provided, return the trie only consisting of the storage root
+        let storage_root = proof.storage_root;
+        if proof.storage_proofs.is_empty() {
+            let storage_root_node = node_from_digest(storage_root);
+            storage.insert(B256::from(&keccak(address)), storage_root_node);
+            continue;
+        }
+
+        let mut storage_nodes = HashMap::new();
+        let mut storage_root_node = MptNode::default();
+        for storage_proof in &proof.storage_proofs {
+            let proof_nodes = parse_proof(&storage_proof.proof).unwrap();
+            mpt_from_proof(&proof_nodes).unwrap();
+
+            // the first node in the proof is the root
+            if let Some(node) = proof_nodes.first() {
+                storage_root_node = node.clone();
+            }
+
+            proof_nodes.into_iter().for_each(|node| {
+                storage_nodes.insert(node.reference(), node);
+            });
+        }
+
+        // create the storage trie, from all the relevant nodes
+        let storage_trie = resolve_nodes(&storage_root_node, &storage_nodes);
+        assert_eq!(storage_trie.hash(), storage_root);
+
+        storage.insert(B256::from(&keccak(address)), storage_trie);
+    }
+    let state_trie = resolve_nodes(&state_root_node, &state_nodes);
+    assert_eq!(state_trie.hash(), state_root);
+
+    Ok(EthereumState { state_trie, storage_tries: storage })
+}
+
+pub fn transition_proofs_to_tries(
     state_root: B256,
     parent_proofs: &HashMap<Address, AccountProof>,
     proofs: &HashMap<Address, AccountProof>,