feat(trie): construct HashedPostState from revert range (#6072)

crates/storage/provider/src/bundle_state/hashed_state_changes.rs

@@ -104,7 +104,7 @@ mod tests {
         }
 
         let mut hashed_state = HashedPostState::default();
-        hashed_state.insert_destroyed_account(destroyed_address_hashed);
+        hashed_state.insert_account(destroyed_address_hashed, None);
         hashed_state.insert_hashed_storage(destroyed_address_hashed, HashedStorage::new(true));
 
         let provider_rw = provider_factory.provider_rw().unwrap();

crates/trie/src/hashed_cursor/post_state.rs

@@ -453,7 +453,7 @@ mod tests {
 
         let mut hashed_post_state = HashedPostState::default();
         for (hashed_address, account) in &accounts {
-            hashed_post_state.insert_account(*hashed_address, *account);
+            hashed_post_state.insert_account(*hashed_address, Some(*account));
         }
         hashed_post_state.sort();
 
@@ -499,7 +499,7 @@ mod tests {
 
         let mut hashed_post_state = HashedPostState::default();
         for (hashed_address, account) in accounts.iter().filter(|x| x.0[31] % 2 != 0) {
-            hashed_post_state.insert_account(*hashed_address, *account);
+            hashed_post_state.insert_account(*hashed_address, Some(*account));
         }
         hashed_post_state.sort();
 
@@ -526,11 +526,9 @@ mod tests {
 
         let mut hashed_post_state = HashedPostState::default();
         for (hashed_address, account) in accounts.iter().filter(|x| x.0[31] % 2 != 0) {
-            if removed_keys.contains(hashed_address) {
-                hashed_post_state.insert_destroyed_account(*hashed_address);
-            } else {
-                hashed_post_state.insert_account(*hashed_address, *account);
-            }
+            let account_info =
+                if removed_keys.contains(hashed_address) { None } else { Some(*account) };
+            hashed_post_state.insert_account(*hashed_address, account_info)
         }
         hashed_post_state.sort();
 
@@ -557,7 +555,7 @@ mod tests {
 
         let mut hashed_post_state = HashedPostState::default();
         for (hashed_address, account) in &accounts {
-            hashed_post_state.insert_account(*hashed_address, *account);
+            hashed_post_state.insert_account(*hashed_address, Some(*account));
         }
         hashed_post_state.sort();
 
@@ -579,11 +577,7 @@ mod tests {
 
                 let mut hashed_post_state = HashedPostState::default();
                 for (hashed_address, account) in &post_state_accounts {
-                    if let Some(account) = account {
-                        hashed_post_state.insert_account(*hashed_address, *account);
-                    } else {
-                        hashed_post_state.insert_destroyed_account(*hashed_address);
-                    }
+                    hashed_post_state.insert_account(*hashed_address, *account);
                 }
                 hashed_post_state.sort();
 
@@ -659,7 +653,7 @@ mod tests {
         {
             let wiped = true;
             let mut hashed_storage = HashedStorage::new(wiped);
-            hashed_storage.insert_storage(B256::random(), U256::ZERO);
+            hashed_storage.insert_slot(B256::random(), U256::ZERO);
 
             let mut hashed_post_state = HashedPostState::default();
             hashed_post_state.insert_hashed_storage(address, hashed_storage);
@@ -674,7 +668,7 @@ mod tests {
         {
             let wiped = true;
             let mut hashed_storage = HashedStorage::new(wiped);
-            hashed_storage.insert_storage(B256::random(), U256::from(1));
+            hashed_storage.insert_slot(B256::random(), U256::from(1));
 
             let mut hashed_post_state = HashedPostState::default();
             hashed_post_state.insert_hashed_storage(address, hashed_storage);
@@ -710,7 +704,7 @@ mod tests {
         let wiped = false;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in post_state_storage.iter() {
-            hashed_storage.insert_storage(*slot, *value);
+            hashed_storage.insert_slot(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -746,7 +740,7 @@ mod tests {
         let wiped = false;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in post_state_storage.iter() {
-            hashed_storage.insert_storage(*slot, *value);
+            hashed_storage.insert_slot(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -784,7 +778,7 @@ mod tests {
         let wiped = true;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in post_state_storage.iter() {
-            hashed_storage.insert_storage(*slot, *value);
+            hashed_storage.insert_slot(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -819,7 +813,7 @@ mod tests {
         let wiped = false;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in storage.iter() {
-            hashed_storage.insert_storage(*slot, *value);
+            hashed_storage.insert_slot(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -856,7 +850,7 @@ mod tests {
             for (address, (wiped, storage)) in &post_state_storages {
                 let mut hashed_storage = HashedStorage::new(*wiped);
                 for (slot, value) in storage {
-                    hashed_storage.insert_storage(*slot, *value);
+                    hashed_storage.insert_slot(*slot, *value);
                 }
                 hashed_post_state.insert_hashed_storage(*address, hashed_storage);
             }

crates/trie/src/state.rs

@@ -5,11 +5,22 @@ use crate::{
     StateRoot, StateRootError,
 };
 use ahash::{AHashMap, AHashSet};
-use reth_db::transaction::DbTx;
+use reth_db::{
+    cursor::DbCursorRO,
+    models::{AccountBeforeTx, BlockNumberAddress},
+    tables,
+    transaction::DbTx,
+    DatabaseError,
+};
 use reth_primitives::{
-    keccak256, revm::compat::into_reth_acc, trie::Nibbles, Account, Address, B256, U256,
+    keccak256, revm::compat::into_reth_acc, trie::Nibbles, Account, Address, BlockNumber, B256,
+    U256,
 };
 use revm::db::BundleAccount;
+use std::{
+    collections::{hash_map, HashMap},
+    ops::RangeInclusive,
+};
 
 /// The post state with hashed addresses as keys.
 #[derive(Debug, Clone, Eq, PartialEq)]
@@ -42,26 +53,83 @@ impl HashedPostState {
     pub fn from_bundle_state<'a>(
         state: impl IntoIterator<Item = (&'a Address, &'a BundleAccount)>,
     ) -> Self {
-        let mut hashed_state = Self::default();
+        let mut this = Self::default();
 
         for (address, account) in state {
             let hashed_address = keccak256(address);
-            if let Some(account) = &account.info {
-                hashed_state.insert_account(hashed_address, into_reth_acc(account.clone()))
-            } else {
-                hashed_state.insert_destroyed_account(hashed_address);
-            }
+            this.insert_account(hashed_address, account.info.clone().map(into_reth_acc));
 
             // insert storage.
             let mut hashed_storage = HashedStorage::new(account.status.was_destroyed());
 
             for (key, value) in account.storage.iter() {
                 let hashed_key = keccak256(B256::new(key.to_be_bytes()));
-                hashed_storage.insert_storage(hashed_key, value.present_value);
+                hashed_storage.insert_slot(hashed_key, value.present_value);
             }
-            hashed_state.insert_hashed_storage(hashed_address, hashed_storage)
+            this.insert_hashed_storage(hashed_address, hashed_storage)
         }
-        hashed_state.sorted()
+        this.sorted()
+    }
+
+    /// Initialize [HashedPostState] from revert range.
+    /// Iterate over state reverts in the specified block range and
+    /// apply them to hashed state in reverse.
+    ///
+    /// NOTE: In order to have the resulting [HashedPostState] be a correct
+    /// overlay of the plain state, the end of the range must be the current tip.
+    pub fn from_revert_range<TX: DbTx>(
+        tx: TX,
+        range: RangeInclusive<BlockNumber>,
+    ) -> Result<Self, DatabaseError> {
+        // A single map for aggregating state changes where each map value is a tuple
+        // `(maybe_account_change, storage_changes)`.
+        // If `maybe_account_change` is `None`, no account info change had occurred.
+        // If `maybe_account_change` is `Some(None)`, the account had previously been destroyed
+        // or non-existent.
+        // If `maybe_account_change` is `Some(Some(info))`, the contained value is the previous
+        // account state.
+        let mut state =
+            HashMap::<Address, (Option<Option<Account>>, HashMap<B256, U256>)>::default();
+
+        // Iterate over account changesets in reverse.
+        let mut account_changesets_cursor = tx.cursor_read::<tables::AccountChangeSet>()?;
+        for entry in account_changesets_cursor.walk_range(range.clone())? {
+            let (_, AccountBeforeTx { address, info }) = entry?;
+            let account_entry = state.entry(address).or_default();
+            if account_entry.0.is_none() {
+                account_entry.0 = Some(info);
+            }
+        }
+
+        // Iterate over storage changesets in reverse.
+        let mut storage_changesets_cursor = tx.cursor_read::<tables::StorageChangeSet>()?;
+        for entry in storage_changesets_cursor.walk_range(BlockNumberAddress::range(range))? {
+            let (BlockNumberAddress((_, address)), storage) = entry?;
+            let account_entry = state.entry(address).or_default();
+            if let hash_map::Entry::Vacant(entry) = account_entry.1.entry(storage.key) {
+                entry.insert(storage.value);
+            }
+        }
+
+        let mut this = Self::default();
+        for (address, (maybe_account_change, storage)) in state {
+            let hashed_address = keccak256(address);
+
+            if let Some(account_change) = maybe_account_change {
+                this.insert_account(hashed_address, account_change);
+            }
+
+            // The `wiped`` flag indicates only  whether previous storage entries should be looked
+            // up in db or not. For reverts it's a noop since all wiped changes had been written as
+            // storage reverts.
+            let mut hashed_storage = HashedStorage::new(false);
+            for (slot, value) in storage {
+                hashed_storage.insert_slot(keccak256(slot), value);
+            }
+            this.insert_hashed_storage(hashed_address, hashed_storage);
+        }
+
+        Ok(this.sorted())
     }
 
     /// Sort and return self.
@@ -94,15 +162,14 @@ impl HashedPostState {
         }
     }
 
-    /// Insert non-empty account info.
-    pub fn insert_account(&mut self, hashed_address: B256, account: Account) {
-        self.accounts.push((hashed_address, account));
-        self.sorted = false;
-    }
-
-    /// Insert destroyed hashed account key.
-    pub fn insert_destroyed_account(&mut self, hashed_address: B256) {
-        self.destroyed_accounts.insert(hashed_address);
+    /// Insert account. If `account` is `None`, the account had previously been destroyed.
+    pub fn insert_account(&mut self, hashed_address: B256, account: Option<Account>) {
+        if let Some(account) = account {
+            self.accounts.push((hashed_address, account));
+            self.sorted = false;
+        } else {
+            self.destroyed_accounts.insert(hashed_address);
+        }
     }
 
     /// Insert hashed storage entry.
@@ -183,7 +250,7 @@ impl HashedPostState {
     /// let mut hashed_state = HashedPostState::default();
     /// hashed_state.insert_account(
     ///     [0x11; 32].into(),
-    ///     Account { nonce: 1, balance: U256::from(10), bytecode_hash: None },
+    ///     Some(Account { nonce: 1, balance: U256::from(10), bytecode_hash: None }),
     /// );
     /// hashed_state.sort();
     ///
@@ -256,7 +323,7 @@ impl HashedStorage {
 
     /// Insert storage entry.
     #[inline]
-    pub fn insert_storage(&mut self, slot: B256, value: U256) {
+    pub fn insert_slot(&mut self, slot: B256, value: U256) {
         if value.is_zero() {
             self.zero_valued_slots.insert(slot);
         } else {