refactor

crates/core/src/mpt.rs

@@ -62,7 +62,7 @@ pub struct MptNode {
 pub enum Error {
     /// Triggered when an operation reaches an unresolved node. The associated `B256`
     /// value provides details about the unresolved node.
-    #[error("reached an unresolved node: {0:#}")]
+    #[error("reached an unresolved node: {0:?}")]
     NodeNotResolved(B256),
     /// Occurs when a value is unexpectedly found in a branch node.
     #[error("branch node with value")]
@@ -526,13 +526,16 @@ impl MptNode {
                                 mem::take(orphan_child),
                             );
                         }
-                        // if the orphan is a branch or digest, convert to an extension
-                        MptNodeData::Branch(_) | MptNodeData::Digest(_) => {
+                        // if the orphan is a branch, convert to an extension
+                        MptNodeData::Branch(_) => {
                             self.data = MptNodeData::Extension(
                                 to_encoded_path(&[index as u8], false),
                                 orphan,
                             );
                         }
+                        MptNodeData::Digest(digest) => {
+                            return Err(Error::NodeNotResolved(*digest));
+                        }
                         MptNodeData::Null => unreachable!(),
                     }
                 }
@@ -574,8 +577,10 @@ impl MptNode {
                             mem::take(node),
                         );
                     }
-                    // for a branch or digest, the extension is still correct
-                    MptNodeData::Branch(_) | MptNodeData::Digest(_) => {}
+                    // for a branch, the extension is still correct
+                    MptNodeData::Branch(_) => {}
+                    // if the child were a digest an early return should have been hit
+                    MptNodeData::Digest(_) => unreachable!(),
                 }
             }
             MptNodeData::Digest(digest) => return Err(Error::NodeNotResolved(*digest)),
@@ -993,7 +998,7 @@ pub fn shorten_node_path(node: &MptNode) -> Vec<MptNode> {
     let mut res = Vec::new();
     let nibs = node.nibs();
     match node.as_data() {
-        MptNodeData::Null | MptNodeData::Branch(_) | MptNodeData::Digest(_) => {}
+        MptNodeData::Null | MptNodeData::Branch(_) => {}
         MptNodeData::Leaf(_, value) => {
             for i in 0..=nibs.len() {
                 res.push(MptNodeData::Leaf(to_encoded_path(&nibs[i..], true), value.clone()).into())
@@ -1007,6 +1012,7 @@ pub fn shorten_node_path(node: &MptNode) -> Vec<MptNode> {
                 )
             }
         }
+        MptNodeData::Digest(_) => unreachable!(),
     };
     res
 }

crates/core/src/stateless/finalize.rs

@@ -65,26 +65,30 @@ impl<Driver: CoreDriver> FinalizationStrategy<Driver, MemoryDB> for RethFinaliza
                 storage_trie.clear();
             }
             // apply all new storage entries for the current account (address)
+            let mut deletions = Vec::with_capacity(storage_change.storage.len());
             for (key, value) in &storage_change.storage {
                 let storage_trie_index = keccak(key.to_be_bytes::<32>());
                 if value.is_zero() {
-                    storage_trie
-                        .delete(&storage_trie_index)
-                        .context("storage_trie.delete")?;
+                    deletions.push(storage_trie_index);
                 } else {
                     storage_trie
                         .insert_rlp(&storage_trie_index, value)
                         .context("storage_trie.insert_rlp")?;
                 }
             }
+            // Apply deferred storage trie deletions
+            for storage_trie_index in deletions {
+                storage_trie
+                    .delete(&storage_trie_index)
+                    .context("storage_trie.delete")?;
+            }
         }
         // Apply account info + storage changes
+        let mut deletions = Vec::with_capacity(accounts.len());
         for (address, account_info) in accounts {
             let state_trie_index = keccak(address);
             if account_info.is_none() {
-                state_trie
-                    .delete(&state_trie_index)
-                    .context("state_trie.delete")?;
+                deletions.push(state_trie_index);
                 continue;
             }
             let storage_root = {
@@ -121,6 +125,12 @@ impl<Driver: CoreDriver> FinalizationStrategy<Driver, MemoryDB> for RethFinaliza
                     .context("state_trie.insert_rlp (2)")?;
             }
         }
+        // Apply deferred state trie deletions
+        for state_trie_index in deletions {
+            state_trie
+                .delete(&state_trie_index)
+                .context("state_trie.delete")?;
+        }
 
         // Update the database
         if let Some(db) = db {

crates/preflight/src/client.rs

@@ -145,22 +145,6 @@ where
             let mut preflight_db = engine.db.take().unwrap().unwrap();
             preflight_db.apply_changeset(state_changeset)?;
 
-            // storage sanity check
-            // {
-            //     let init_db = preflight_db.db.db.db.borrow_mut();
-            //     let mut provider_db = init_db.db.clone();
-            //     provider_db.block_no += 1;
-            //     for (address, db_account) in &preflight_db.db.accounts {
-            //         use reth_revm::Database;
-            //         let provider_info = provider_db.basic(*address)?.unwrap();
-            //         if db_account.info != provider_info {
-            //             dbg!(&address);
-            //             dbg!(&db_account.info);
-            //             dbg!(&provider_info);
-            //         }
-            //     }
-            // }
-
             // Save the provider cache
             info!("Saving provider cache ...");
             preflight_db.save_provider()?;

crates/preflight/src/db.rs

@@ -172,13 +172,13 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
     ) -> anyhow::Result<HashMap<Address, EIP1186AccountProofResponse>> {
         // get initial keys
         let initial_db = &self.inner.db;
-        let mut initial_storage_keys = enumerate_storage_keys(&initial_db.db.borrow());
+        let mut storage_keys = enumerate_storage_keys(&initial_db.db.borrow());
         // merge initial keys with latest db storage keys
         for (address, mut indices) in enumerate_storage_keys(&self.inner) {
-            match initial_storage_keys.get_mut(&address) {
+            match storage_keys.get_mut(&address) {
                 Some(initial_indices) => initial_indices.append(&mut indices),
                 None => {
-                    initial_storage_keys.insert(address, indices);
+                    storage_keys.insert(address, indices);
                 }
             }
         }
@@ -188,7 +188,7 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
         let res = get_proofs(
             initial_db.db.provider.borrow_mut().deref_mut(),
             block_no,
-            initial_storage_keys,
+            storage_keys,
         )?;
         Ok(res)
     }

crates/preflight/src/trie.rs

@@ -46,7 +46,7 @@ pub fn extend_proof_tries(
         let finalization_proof = finalization_proofs
             .get(&address)
             .with_context(|| format!("missing finalization proof for address {:#}", &address))?;
-        add_orphaned_leafs(address, &finalization_proof.account_proof, &mut state_nodes)?;
+        add_orphaned_nodes(address, &finalization_proof.account_proof, &mut state_nodes)?;
         // insert inaccessible storage trie
         if let alloy::primitives::map::Entry::Vacant(e) = storage_tries.entry(address) {
             e.insert((initialization_proof.storage_hash.into(), vec![]));
@@ -72,7 +72,7 @@ pub fn extend_proof_tries(
 
         // assure that slots can be deleted from the storage trie
         for storage_proof in &finalization_proof.storage_proof {
-            add_orphaned_leafs(
+            add_orphaned_nodes(
                 storage_proof.key.0,
                 &storage_proof.proof,
                 &mut storage_nodes,
@@ -123,7 +123,7 @@ pub fn proofs_to_tries(
             .with_context(|| format!("missing finalization proof for address {:#}", &address))?;
 
         // assure that addresses can be deleted from the state trie
-        add_orphaned_leafs(address, &finalization_proof.account_proof, &mut state_nodes)?;
+        add_orphaned_nodes(address, &finalization_proof.account_proof, &mut state_nodes)?;
 
         // if no slots are provided, return the trie only consisting of the storage root
         if initialization_proof.storage_proof.is_empty() {
@@ -150,7 +150,7 @@ pub fn proofs_to_tries(
 
         // assure that slots can be deleted from the storage trie
         for storage_proof in &finalization_proof.storage_proof {
-            add_orphaned_leafs(
+            add_orphaned_nodes(
                 storage_proof.key.0,
                 &storage_proof.proof,
                 &mut storage_nodes,
@@ -174,20 +174,22 @@ pub fn proofs_to_tries(
     Ok((state_trie, storage))
 }
 
-/// Adds all the leaf nodes of non-inclusion proofs to the nodes.
-pub fn add_orphaned_leafs(
+/// Adds all the nodes of non-inclusion proofs to the nodes.
+pub fn add_orphaned_nodes(
     key: impl AsRef<[u8]>,
     proof: &[impl AsRef<[u8]>],
     nodes_by_reference: &mut HashMap<MptNodeReference, MptNode>,
 ) -> anyhow::Result<()> {
     if !proof.is_empty() {
         let proof_nodes = parse_proof(proof).context("invalid proof encoding")?;
         if is_not_included(&keccak(key), &proof_nodes)? {
-            // add the leaf node to the nodes
-            let leaf = proof_nodes.last().unwrap();
-            shorten_node_path(leaf).into_iter().for_each(|node| {
-                nodes_by_reference.insert(node.reference(), node);
-            });
+            // extract possible orphan
+            for node in &proof_nodes {
+                nodes_by_reference.insert(node.reference(), node.clone());
+                shorten_node_path(node).into_iter().for_each(|node| {
+                    nodes_by_reference.insert(node.reference(), node);
+                });
+            }
         }
     }