perf(trie): post state cursors (#3588)

crates/trie/src/hashed_cursor/post_state.rs

@@ -6,27 +6,113 @@ use reth_db::{
     transaction::{DbTx, DbTxGAT},
 };
 use reth_primitives::{trie::Nibbles, Account, StorageEntry, H256, U256};
-use std::collections::{BTreeMap, HashMap};
+use std::collections::{HashMap, HashSet};
 
 /// The post state account storage with hashed slots.
-#[derive(Debug, Default, Clone, Eq, PartialEq)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub struct HashedStorage {
+    /// Hashed storage slots with non-zero.
+    non_zero_valued_storage: Vec<(H256, U256)>,
+    /// Slots that have been zero valued.
+    zero_valued_slots: HashSet<H256>,
     /// Whether the storage was wiped or not.
-    pub wiped: bool,
-    /// Hashed storage slots.
-    pub storage: BTreeMap<H256, U256>,
+    wiped: bool,
+    /// Whether the storage entries were sorted or not.
+    sorted: bool,
+}
+
+impl HashedStorage {
+    /// Create new instance of [HashedStorage].
+    pub fn new(wiped: bool) -> Self {
+        Self {
+            non_zero_valued_storage: Vec::new(),
+            zero_valued_slots: HashSet::new(),
+            wiped,
+            sorted: true, // empty is sorted
+        }
+    }
+
+    /// Sorts the non zero value storage entries.
+    pub fn sort_storage(&mut self) {
+        if !self.sorted {
+            self.non_zero_valued_storage.sort_unstable_by_key(|(slot, _)| *slot);
+            self.sorted = true;
+        }
+    }
+
+    /// Insert non zero-valued storage entry.
+    pub fn insert_non_zero_valued_storage(&mut self, slot: H256, value: U256) {
+        debug_assert!(value != U256::ZERO, "value cannot be zero");
+        self.non_zero_valued_storage.push((slot, value));
+        self.sorted = false;
+    }
+
+    /// Insert zero-valued storage slot.
+    pub fn insert_zero_valued_slot(&mut self, slot: H256) {
+        self.zero_valued_slots.insert(slot);
+    }
 }
 
 /// The post state with hashed addresses as keys.
-#[derive(Debug, Default, Clone, Eq, PartialEq)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub struct HashedPostState {
     /// Map of hashed addresses to account info.
-    pub accounts: BTreeMap<H256, Option<Account>>,
+    accounts: Vec<(H256, Account)>,
+    /// Set of cleared accounts.
+    cleared_accounts: HashSet<H256>,
     /// Map of hashed addresses to hashed storage.
-    pub storages: BTreeMap<H256, HashedStorage>,
+    storages: HashMap<H256, HashedStorage>,
+    /// Whether the account and storage entries were sorted or not.
+    sorted: bool,
+}
+
+impl Default for HashedPostState {
+    fn default() -> Self {
+        Self {
+            accounts: Vec::new(),
+            cleared_accounts: HashSet::new(),
+            storages: HashMap::new(),
+            sorted: true, // empty is sorted
+        }
+    }
 }
 
 impl HashedPostState {
+    /// Sort and return self.
+    pub fn sorted(mut self) -> Self {
+        self.sort();
+        self
+    }
+
+    /// Sort account and storage entries.
+    pub fn sort(&mut self) {
+        if !self.sorted {
+            for (_, storage) in self.storages.iter_mut() {
+                storage.sort_storage();
+            }
+
+            self.accounts.sort_unstable_by_key(|(address, _)| *address);
+            self.sorted = true;
+        }
+    }
+
+    /// Insert non-empty account info.
+    pub fn insert_account(&mut self, hashed_address: H256, account: Account) {
+        self.accounts.push((hashed_address, account));
+        self.sorted = false;
+    }
+
+    /// Insert cleared hashed account key.
+    pub fn insert_cleared_account(&mut self, hashed_address: H256) {
+        self.cleared_accounts.insert(hashed_address);
+    }
+
+    /// Insert hashed storage entry.
+    pub fn insert_hashed_storage(&mut self, hashed_address: H256, hashed_storage: HashedStorage) {
+        self.sorted &= hashed_storage.sorted;
+        self.storages.insert(hashed_address, hashed_storage);
+    }
+
     /// Construct (PrefixSet)[PrefixSet] from hashed post state.
     /// The prefix sets contain the hashed account and storage keys that have been changed in the
     /// post state.
@@ -35,17 +121,24 @@ impl HashedPostState {
         let mut account_prefix_set = PrefixSetMut::default();
         let mut storage_prefix_set: HashMap<H256, PrefixSetMut> = HashMap::default();
 
-        for hashed_address in self.accounts.keys() {
+        // Populate account prefix set.
+        for (hashed_address, _) in &self.accounts {
+            account_prefix_set.insert(Nibbles::unpack(hashed_address));
+        }
+        for hashed_address in &self.cleared_accounts {
             account_prefix_set.insert(Nibbles::unpack(hashed_address));
         }
 
+        // Populate storage prefix sets.
         for (hashed_address, hashed_storage) in self.storages.iter() {
             account_prefix_set.insert(Nibbles::unpack(hashed_address));
-            for hashed_slot in hashed_storage.storage.keys() {
-                storage_prefix_set
-                    .entry(*hashed_address)
-                    .or_default()
-                    .insert(Nibbles::unpack(hashed_slot));
+
+            let storage_prefix_set_entry = storage_prefix_set.entry(*hashed_address).or_default();
+            for (hashed_slot, _) in &hashed_storage.non_zero_valued_storage {
+                storage_prefix_set_entry.insert(Nibbles::unpack(hashed_slot));
+            }
+            for hashed_slot in &hashed_storage.zero_valued_slots {
+                storage_prefix_set_entry.insert(Nibbles::unpack(hashed_slot));
             }
         }
 
@@ -74,22 +167,17 @@ impl<'a, 'b, 'tx, TX: DbTx<'tx>> HashedCursorFactory<'a>
 where
     'a: 'b,
 {
-    type AccountCursor = HashedPostStateAccountCursor<'b, <TX as DbTxGAT<'a>>::Cursor<tables::HashedAccount>> where Self: 'a ;
+    type AccountCursor = HashedPostStateAccountCursor<'b, <TX as DbTxGAT<'a>>::Cursor<tables::HashedAccount>> where Self: 'a;
     type StorageCursor = HashedPostStateStorageCursor<'b, <TX as DbTxGAT<'a>>::DupCursor<tables::HashedStorage>> where Self: 'a;
 
     fn hashed_account_cursor(&'a self) -> Result<Self::AccountCursor, reth_db::DatabaseError> {
         let cursor = self.tx.cursor_read::<tables::HashedAccount>()?;
-        Ok(HashedPostStateAccountCursor { post_state: self.post_state, cursor, last_account: None })
+        Ok(HashedPostStateAccountCursor::new(cursor, self.post_state))
     }
 
     fn hashed_storage_cursor(&'a self) -> Result<Self::StorageCursor, reth_db::DatabaseError> {
         let cursor = self.tx.cursor_dup_read::<tables::HashedStorage>()?;
-        Ok(HashedPostStateStorageCursor {
-            post_state: self.post_state,
-            cursor,
-            account: None,
-            last_slot: None,
-        })
+        Ok(HashedPostStateStorageCursor::new(cursor, self.post_state))
     }
 }
 
@@ -101,22 +189,26 @@ pub struct HashedPostStateAccountCursor<'b, C> {
     cursor: C,
     /// The reference to the in-memory [HashedPostState].
     post_state: &'b HashedPostState,
+    /// The post state account index where the cursor is currently at.
+    post_state_account_index: usize,
     /// The last hashed account key that was returned by the cursor.
     /// De facto, this is a current cursor position.
     last_account: Option<H256>,
 }
 
-impl<'b, 'tx, C> HashedPostStateAccountCursor<'b, C>
-where
-    C: DbCursorRO<'tx, tables::HashedAccount>,
-{
+impl<'b, C> HashedPostStateAccountCursor<'b, C> {
+    /// Create new instance of [HashedPostStateAccountCursor].
+    pub fn new(cursor: C, post_state: &'b HashedPostState) -> Self {
+        Self { cursor, post_state, last_account: None, post_state_account_index: 0 }
+    }
+
     /// Returns `true` if the account has been destroyed.
     /// This check is used for evicting account keys from the state trie.
     ///
     /// This function only checks the post state, not the database, because the latter does not
     /// store destroyed accounts.
     fn is_account_cleared(&self, account: &H256) -> bool {
-        matches!(self.post_state.accounts.get(account), Some(None))
+        self.post_state.cleared_accounts.contains(account)
     }
 
     /// Return the account with the lowest hashed account key.
@@ -124,30 +216,28 @@ where
     /// Given the next post state and database entries, return the smallest of the two.
     /// If the account keys are the same, the post state entry is given precedence.
     fn next_account(
-        &self,
-        post_state_item: Option<(H256, Account)>,
+        post_state_item: Option<&(H256, Account)>,
         db_item: Option<(H256, Account)>,
-    ) -> Result<Option<(H256, Account)>, reth_db::DatabaseError> {
-        let result = match (post_state_item, db_item) {
+    ) -> Option<(H256, Account)> {
+        match (post_state_item, db_item) {
             // If both are not empty, return the smallest of the two
             // Post state is given precedence if keys are equal
             (Some((post_state_address, post_state_account)), Some((db_address, db_account))) => {
-                if post_state_address <= db_address {
-                    Some((post_state_address, post_state_account))
+                if post_state_address <= &db_address {
+                    Some((*post_state_address, *post_state_account))
                 } else {
                     Some((db_address, db_account))
                 }
             }
             // If the database is empty, return the post state entry
             (Some((post_state_address, post_state_account)), None) => {
-                Some((post_state_address, post_state_account))
+                Some((*post_state_address, *post_state_account))
             }
             // If the post state is empty, return the database entry
             (None, Some((db_address, db_account))) => Some((db_address, db_account)),
             // If both are empty, return None
             (None, None) => None,
-        };
-        Ok(result)
+        }
     }
 }
 
@@ -164,35 +254,45 @@ where
     /// The returned account key is memoized and the cursor remains positioned at that key until
     /// [HashedAccountCursor::seek] or [HashedAccountCursor::next] are called.
     fn seek(&mut self, key: H256) -> Result<Option<(H256, Account)>, reth_db::DatabaseError> {
+        debug_assert!(self.post_state.sorted, "`HashedPostState` must be pre-sorted");
+
         self.last_account = None;
 
-        // Attempt to find the account in poststate.
-        let post_state_item = self
-            .post_state
-            .accounts
-            .iter()
-            .find_map(|(k, v)| v.filter(|_| k >= &key).map(|v| (*k, v)));
-        if let Some((address, account)) = post_state_item {
-            // It's an exact match, return the account from post state without looking up in the
-            // database.
-            if address == key {
-                self.last_account = Some(address);
-                return Ok(Some((address, account)))
+        // Take the next account from the post state with the key greater than or equal to the
+        // sought key.
+        let mut post_state_entry = self.post_state.accounts.get(self.post_state_account_index);
+        while let Some((k, _)) = post_state_entry {
+            if k >= &key {
+                // Found the next entry that is equal or greater than the key.
+                break
+            }
+
+            self.post_state_account_index += 1;
+            post_state_entry = self.post_state.accounts.get(self.post_state_account_index);
+        }
+
+        // It's an exact match, return the account from post state without looking up in the
+        // database.
+        if let Some((address, account)) = post_state_entry {
+            if address == &key {
+                self.last_account = Some(*address);
+                return Ok(Some((*address, *account)))
             }
         }
 
         // It's not an exact match, reposition to the first greater or equal account that wasn't
         // cleared.
-        let mut db_item = self.cursor.seek(key)?;
-        while db_item
+        let mut db_entry = self.cursor.seek(key)?;
+        while db_entry
             .as_ref()
             .map(|(address, _)| self.is_account_cleared(address))
             .unwrap_or_default()
         {
-            db_item = self.cursor.next()?;
+            db_entry = self.cursor.next()?;
         }
 
-        let result = self.next_account(post_state_item, db_item)?;
+        // Compare two entries and return the lowest.
+        let result = Self::next_account(post_state_entry, db_entry);
         self.last_account = result.as_ref().map(|(address, _)| *address);
         Ok(result)
     }
@@ -205,28 +305,38 @@ where
     /// NOTE: This function will not return any entry unless [HashedAccountCursor::seek] has been
     /// called.
     fn next(&mut self) -> Result<Option<(H256, Account)>, reth_db::DatabaseError> {
+        debug_assert!(self.post_state.sorted, "`HashedPostState` must be pre-sorted");
+
         let last_account = match self.last_account.as_ref() {
             Some(account) => account,
             None => return Ok(None), // no previous entry was found
         };
 
         // If post state was given precedence, move the cursor forward.
-        let mut db_item = self.cursor.current()?;
-        while db_item
+        let mut db_entry = self.cursor.current()?;
+        while db_entry
             .as_ref()
             .map(|(address, _)| address <= last_account || self.is_account_cleared(address))
             .unwrap_or_default()
         {
-            db_item = self.cursor.next()?;
+            db_entry = self.cursor.next()?;
         }
 
-        let post_state_item = self
-            .post_state
-            .accounts
-            .iter()
-            .find(|(k, v)| k > &last_account && v.is_some())
-            .map(|(address, info)| (*address, info.unwrap()));
-        let result = self.next_account(post_state_item, db_item)?;
+        // Take the next account from the post state with the key greater than or equal to the
+        // sought key.
+        let mut post_state_entry = self.post_state.accounts.get(self.post_state_account_index);
+        while let Some((k, _)) = post_state_entry {
+            if k > last_account {
+                // Found the next entry in the post state.
+                break
+            }
+
+            self.post_state_account_index += 1;
+            post_state_entry = self.post_state.accounts.get(self.post_state_account_index);
+        }
+
+        // Compare two entries and return the lowest.
+        let result = Self::next_account(post_state_entry, db_entry);
         self.last_account = result.as_ref().map(|(address, _)| *address);
         Ok(result)
     }
@@ -240,6 +350,8 @@ pub struct HashedPostStateStorageCursor<'b, C> {
     cursor: C,
     /// The reference to the post state.
     post_state: &'b HashedPostState,
+    /// The post state index where the cursor is currently at.
+    post_state_storage_index: usize,
     /// The current hashed account key.
     account: Option<H256>,
     /// The last slot that has been returned by the cursor.
@@ -248,6 +360,11 @@ pub struct HashedPostStateStorageCursor<'b, C> {
 }
 
 impl<'b, C> HashedPostStateStorageCursor<'b, C> {
+    /// Create new instance of [HashedPostStateStorageCursor].
+    pub fn new(cursor: C, post_state: &'b HashedPostState) -> Self {
+        Self { cursor, post_state, account: None, last_slot: None, post_state_storage_index: 0 }
+    }
+
     /// Returns `true` if the storage for the given
     /// The database is not checked since it already has no wiped storage entries.
     fn is_db_storage_wiped(&self, account: &H256) -> bool {
@@ -259,12 +376,11 @@ impl<'b, C> HashedPostStateStorageCursor<'b, C> {
 
     /// Check if the slot was zeroed out in the post state.
     /// The database is not checked since it already has no zero-valued slots.
-    fn is_touched_slot_value_zero(&self, account: &H256, slot: &H256) -> bool {
+    fn is_slot_zero_valued(&self, account: &H256, slot: &H256) -> bool {
         self.post_state
             .storages
             .get(account)
-            .and_then(|storage| storage.storage.get(slot))
-            .map(|value| *value == U256::ZERO)
+            .map(|storage| storage.zero_valued_slots.contains(slot))
             .unwrap_or_default()
     }
 
@@ -274,10 +390,10 @@ impl<'b, C> HashedPostStateStorageCursor<'b, C> {
     /// If the storage keys are the same, the post state entry is given precedence.
     fn next_slot(
         &self,
-        post_state_item: Option<(&H256, &U256)>,
+        post_state_item: Option<&(H256, U256)>,
         db_item: Option<StorageEntry>,
-    ) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
-        let result = match (post_state_item, db_item) {
+    ) -> Option<StorageEntry> {
+        match (post_state_item, db_item) {
             // If both are not empty, return the smallest of the two
             // Post state is given precedence if keys are equal
             (Some((post_state_slot, post_state_value)), Some(db_entry)) => {
@@ -295,8 +411,7 @@ impl<'b, C> HashedPostStateStorageCursor<'b, C> {
             (None, Some(db_entry)) => Some(db_entry),
             // If both are empty, return None
             (None, None) => None,
-        };
-        Ok(result)
+        }
     }
 }
 
@@ -314,7 +429,7 @@ where
                 // If the storage has been wiped at any point
                 storage.wiped &&
                     // and the current storage does not contain any non-zero values 
-                    storage.storage.iter().all(|(_, value)| *value == U256::ZERO)
+                    storage.non_zero_valued_storage.is_empty()
             }
             None => self.cursor.seek_exact(key)?.is_none(),
         };
@@ -327,24 +442,34 @@ where
         account: H256,
         subkey: H256,
     ) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
-        self.last_slot = None;
-        self.account = Some(account);
+        if self.account.map_or(true, |acc| acc != account) {
+            self.account = Some(account);
+            self.last_slot = None;
+            self.post_state_storage_index = 0;
+        }
 
-        // Attempt to find the account's storage in poststate.
-        let post_state_item = self
-            .post_state
-            .storages
-            .get(&account)
-            .map(|storage| {
-                storage
-                    .storage
-                    .iter()
-                    .skip_while(|(slot, value)| slot < &&subkey || value == &&U256::ZERO)
-            })
-            .and_then(|mut iter| iter.next());
-        if let Some((slot, value)) = post_state_item {
-            // It's an exact match, return the storage slot from post state without looking up in
-            // the database.
+        // Attempt to find the account's storage in post state.
+        let mut post_state_entry = None;
+        if let Some(storage) = self.post_state.storages.get(&account) {
+            debug_assert!(storage.sorted, "`HashStorage` must be pre-sorted");
+
+            post_state_entry = storage.non_zero_valued_storage.get(self.post_state_storage_index);
+
+            while let Some((slot, _)) = post_state_entry {
+                if slot >= &subkey {
+                    // Found the next entry that is equal or greater than the key.
+                    break
+                }
+
+                self.post_state_storage_index += 1;
+                post_state_entry =
+                    storage.non_zero_valued_storage.get(self.post_state_storage_index);
+            }
+        }
+
+        // It's an exact match, return the storage slot from post state without looking up in
+        // the database.
+        if let Some((slot, value)) = post_state_entry {
             if slot == &subkey {
                 self.last_slot = Some(*slot);
                 return Ok(Some(StorageEntry { key: *slot, value: *value }))
@@ -352,23 +477,24 @@ where
         }
 
         // It's not an exact match, reposition to the first greater or equal account.
-        let db_item = if self.is_db_storage_wiped(&account) {
+        let db_entry = if self.is_db_storage_wiped(&account) {
             None
         } else {
-            let mut db_item = self.cursor.seek_by_key_subkey(account, subkey)?;
+            let mut db_entry = self.cursor.seek_by_key_subkey(account, subkey)?;
 
-            while db_item
+            while db_entry
                 .as_ref()
-                .map(|entry| self.is_touched_slot_value_zero(&account, &entry.key))
+                .map(|entry| self.is_slot_zero_valued(&account, &entry.key))
                 .unwrap_or_default()
             {
-                db_item = self.cursor.next_dup_val()?;
+                db_entry = self.cursor.next_dup_val()?;
             }
 
-            db_item
+            db_entry
         };
 
-        let result = self.next_slot(post_state_item, db_item)?;
+        // Compare two entries and return the lowest.
+        let result = self.next_slot(post_state_entry, db_entry);
         self.last_slot = result.as_ref().map(|entry| entry.key);
         Ok(result)
     }
@@ -387,40 +513,49 @@ where
             None => return Ok(None), // no previous entry was found
         };
 
-        let db_item = if self.is_db_storage_wiped(&account) {
+        let db_entry = if self.is_db_storage_wiped(&account) {
             None
         } else {
             // If post state was given precedence, move the cursor forward.
-            let mut db_item = self.cursor.seek_by_key_subkey(account, *last_slot)?;
+            let mut db_entry = self.cursor.seek_by_key_subkey(account, *last_slot)?;
 
             // If the entry was already returned, move to the next.
-            if db_item.as_ref().map(|entry| &entry.key == last_slot).unwrap_or_default() {
-                db_item = self.cursor.next_dup_val()?;
+            if db_entry.as_ref().map(|entry| &entry.key == last_slot).unwrap_or_default() {
+                db_entry = self.cursor.next_dup_val()?;
             }
 
-            while db_item
+            while db_entry
                 .as_ref()
-                .map(|entry| self.is_touched_slot_value_zero(&account, &entry.key))
+                .map(|entry| self.is_slot_zero_valued(&account, &entry.key))
                 .unwrap_or_default()
             {
-                db_item = self.cursor.next_dup_val()?;
+                db_entry = self.cursor.next_dup_val()?;
             }
 
-            db_item
+            db_entry
         };
 
-        let post_state_item = self
-            .post_state
-            .storages
-            .get(&account)
-            .map(|storage| {
-                storage
-                    .storage
-                    .iter()
-                    .skip_while(|(slot, value)| slot <= &last_slot || value == &&U256::ZERO)
-            })
-            .and_then(|mut iter| iter.next());
-        let result = self.next_slot(post_state_item, db_item)?;
+        // Attempt to find the account's storage in post state.
+        let mut post_state_entry = None;
+        if let Some(storage) = self.post_state.storages.get(&account) {
+            debug_assert!(storage.sorted, "`HashStorage` must be pre-sorted");
+
+            post_state_entry = storage.non_zero_valued_storage.get(self.post_state_storage_index);
+
+            while let Some((k, _)) = post_state_entry {
+                if k > last_slot {
+                    // Found the next entry.
+                    break
+                }
+
+                self.post_state_storage_index += 1;
+                post_state_entry =
+                    storage.non_zero_valued_storage.get(self.post_state_storage_index);
+            }
+        }
+
+        // Compare two entries and return the lowest.
+        let result = self.next_slot(post_state_entry, db_entry);
         self.last_slot = result.as_ref().map(|entry| entry.key);
         Ok(result)
     }
@@ -431,6 +566,7 @@ mod tests {
     use super::*;
     use proptest::prelude::*;
     use reth_db::{database::Database, test_utils::create_test_rw_db, transaction::DbTxMut};
+    use std::collections::BTreeMap;
 
     fn assert_account_cursor_order<'a, 'b>(
         factory: &'a impl HashedCursorFactory<'b>,
@@ -478,17 +614,17 @@ mod tests {
     fn post_state_only_accounts() {
         let accounts =
             Vec::from_iter((1..11).map(|key| (H256::from_low_u64_be(key), Account::default())));
-        let post_state = HashedPostState {
-            accounts: BTreeMap::from_iter(
-                accounts.iter().map(|(key, account)| (*key, Some(*account))),
-            ),
-            storages: Default::default(),
-        };
+
+        let mut hashed_post_state = HashedPostState::default();
+        for (hashed_address, account) in &accounts {
+            hashed_post_state.insert_account(*hashed_address, *account);
+        }
+        hashed_post_state.sort();
 
         let db = create_test_rw_db();
         let tx = db.tx().unwrap();
 
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         assert_account_cursor_order(&factory, accounts.into_iter());
     }
 
@@ -525,18 +661,14 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: BTreeMap::from_iter(
-                accounts
-                    .iter()
-                    .filter(|x| x.0.to_low_u64_be() % 2 != 0)
-                    .map(|(key, account)| (*key, Some(*account))),
-            ),
-            storages: Default::default(),
-        };
+        let mut hashed_post_state = HashedPostState::default();
+        for (hashed_address, account) in accounts.iter().filter(|x| x.0.to_low_u64_be() % 2 != 0) {
+            hashed_post_state.insert_account(*hashed_address, *account);
+        }
+        hashed_post_state.sort();
 
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         assert_account_cursor_order(&factory, accounts.into_iter());
     }
 
@@ -556,17 +688,18 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: BTreeMap::from_iter(
-                accounts.iter().filter(|x| x.0.to_low_u64_be() % 2 != 0).map(|(key, account)| {
-                    (*key, if removed_keys.contains(key) { None } else { Some(*account) })
-                }),
-            ),
-            storages: Default::default(),
-        };
+        let mut hashed_post_state = HashedPostState::default();
+        for (hashed_address, account) in accounts.iter().filter(|x| x.0.to_low_u64_be() % 2 != 0) {
+            if removed_keys.contains(hashed_address) {
+                hashed_post_state.insert_cleared_account(*hashed_address);
+            } else {
+                hashed_post_state.insert_account(*hashed_address, *account);
+            }
+        }
+        hashed_post_state.sort();
 
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         let expected = accounts.into_iter().filter(|x| !removed_keys.contains(&x.0));
         assert_account_cursor_order(&factory, expected);
     }
@@ -587,15 +720,14 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: BTreeMap::from_iter(
-                accounts.iter().map(|(key, account)| (*key, Some(*account))),
-            ),
-            storages: Default::default(),
-        };
+        let mut hashed_post_state = HashedPostState::default();
+        for (hashed_address, account) in &accounts {
+            hashed_post_state.insert_account(*hashed_address, *account);
+        }
+        hashed_post_state.sort();
 
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         assert_account_cursor_order(&factory, accounts.into_iter());
     }
 
@@ -610,12 +742,15 @@ mod tests {
                 })
                 .unwrap();
 
-                let post_state = HashedPostState {
-                    accounts: BTreeMap::from_iter(
-                        post_state_accounts.iter().map(|(key, account)| (*key, *account)),
-                    ),
-                    storages: Default::default(),
-                };
+                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_cleared_account(*hashed_address);
+                    }
+                }
+                hashed_post_state.sort();
 
                 let mut expected = db_accounts;
                 // overwrite or remove accounts from the expected result
@@ -628,7 +763,7 @@ mod tests {
                 }
 
                 let tx = db.tx().unwrap();
-                let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+                let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
                 assert_account_cursor_order(&factory, expected.into_iter());
             }
         );
@@ -673,51 +808,44 @@ mod tests {
 
         // wiped storage, must be empty
         {
-            let post_state = HashedPostState {
-                accounts: BTreeMap::default(),
-                storages: BTreeMap::from_iter([(
-                    address,
-                    HashedStorage { wiped: true, ..Default::default() },
-                )]),
-            };
+            let wiped = true;
+            let hashed_storage = HashedStorage::new(wiped);
+
+            let mut hashed_post_state = HashedPostState::default();
+            hashed_post_state.insert_hashed_storage(address, hashed_storage);
+
             let tx = db.tx().unwrap();
-            let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+            let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
             let mut cursor = factory.hashed_storage_cursor().unwrap();
             assert!(cursor.is_storage_empty(address).unwrap());
         }
 
         // wiped storage, but post state has zero-value entries
         {
-            let post_state = HashedPostState {
-                accounts: BTreeMap::default(),
-                storages: BTreeMap::from_iter([(
-                    address,
-                    HashedStorage {
-                        wiped: true,
-                        storage: BTreeMap::from_iter([(H256::random(), U256::ZERO)]),
-                    },
-                )]),
-            };
+            let wiped = true;
+            let mut hashed_storage = HashedStorage::new(wiped);
+            hashed_storage.insert_zero_valued_slot(H256::random());
+
+            let mut hashed_post_state = HashedPostState::default();
+            hashed_post_state.insert_hashed_storage(address, hashed_storage);
+
             let tx = db.tx().unwrap();
-            let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+            let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
             let mut cursor = factory.hashed_storage_cursor().unwrap();
             assert!(cursor.is_storage_empty(address).unwrap());
         }
 
         // wiped storage, but post state has non-zero entries
         {
-            let post_state = HashedPostState {
-                accounts: BTreeMap::default(),
-                storages: BTreeMap::from_iter([(
-                    address,
-                    HashedStorage {
-                        wiped: true,
-                        storage: BTreeMap::from_iter([(H256::random(), U256::from(1))]),
-                    },
-                )]),
-            };
+            let wiped = true;
+            let mut hashed_storage = HashedStorage::new(wiped);
+            hashed_storage.insert_non_zero_valued_storage(H256::random(), U256::from(1));
+
+            let mut hashed_post_state = HashedPostState::default();
+            hashed_post_state.insert_hashed_storage(address, hashed_storage);
+
             let tx = db.tx().unwrap();
-            let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+            let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
             let mut cursor = factory.hashed_storage_cursor().unwrap();
             assert!(!cursor.is_storage_empty(address).unwrap());
         }
@@ -727,9 +855,9 @@ mod tests {
     fn storage_cursor_correct_order() {
         let address = H256::random();
         let db_storage =
-            BTreeMap::from_iter((0..10).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
+            BTreeMap::from_iter((1..11).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
         let post_state_storage =
-            BTreeMap::from_iter((10..20).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
+            BTreeMap::from_iter((11..21).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
 
         let db = create_test_rw_db();
         db.update(|tx| {
@@ -744,16 +872,18 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: Default::default(),
-            storages: BTreeMap::from([(
-                address,
-                HashedStorage { wiped: false, storage: post_state_storage.clone() },
-            )]),
-        };
+        let wiped = false;
+        let mut hashed_storage = HashedStorage::new(wiped);
+        for (slot, value) in post_state_storage.iter() {
+            hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+        }
+
+        let mut hashed_post_state = HashedPostState::default();
+        hashed_post_state.insert_hashed_storage(address, hashed_storage);
+        hashed_post_state.sort();
 
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         let expected =
             [(address, db_storage.into_iter().chain(post_state_storage.into_iter()).collect())]
                 .into_iter();
@@ -764,8 +894,7 @@ mod tests {
     fn zero_value_storage_entries_are_discarded() {
         let address = H256::random();
         let db_storage =
-            BTreeMap::from_iter((0..10).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
-        // every even number is changed to zero value
+            BTreeMap::from_iter((0..10).map(|key| (H256::from_low_u64_be(key), U256::from(key)))); // every even number is changed to zero value
         let post_state_storage = BTreeMap::from_iter((0..10).map(|key| {
             (H256::from_low_u64_be(key), if key % 2 == 0 { U256::ZERO } else { U256::from(key) })
         }));
@@ -780,15 +909,22 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: Default::default(),
-            storages: BTreeMap::from([(
-                address,
-                HashedStorage { wiped: false, storage: post_state_storage.clone() },
-            )]),
-        };
+        let wiped = false;
+        let mut hashed_storage = HashedStorage::new(wiped);
+        for (slot, value) in post_state_storage.iter() {
+            if *value == U256::ZERO {
+                hashed_storage.insert_zero_valued_slot(*slot);
+            } else {
+                hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+            }
+        }
+
+        let mut hashed_post_state = HashedPostState::default();
+        hashed_post_state.insert_hashed_storage(address, hashed_storage);
+        hashed_post_state.sort();
+
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         let expected = [(
             address,
             post_state_storage.into_iter().filter(|(_, value)| *value > U256::ZERO).collect(),
@@ -801,9 +937,9 @@ mod tests {
     fn wiped_storage_is_discarded() {
         let address = H256::random();
         let db_storage =
-            BTreeMap::from_iter((0..10).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
+            BTreeMap::from_iter((1..11).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
         let post_state_storage =
-            BTreeMap::from_iter((10..20).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
+            BTreeMap::from_iter((11..21).map(|key| (H256::from_low_u64_be(key), U256::from(key))));
 
         let db = create_test_rw_db();
         db.update(|tx| {
@@ -815,16 +951,18 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: Default::default(),
-            storages: BTreeMap::from([(
-                address,
-                HashedStorage { wiped: true, storage: post_state_storage.clone() },
-            )]),
-        };
+        let wiped = true;
+        let mut hashed_storage = HashedStorage::new(wiped);
+        for (slot, value) in post_state_storage.iter() {
+            hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+        }
+
+        let mut hashed_post_state = HashedPostState::default();
+        hashed_post_state.insert_hashed_storage(address, hashed_storage);
+        hashed_post_state.sort();
 
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         let expected = [(address, post_state_storage)].into_iter();
         assert_storage_cursor_order(&factory, expected);
     }
@@ -848,16 +986,18 @@ mod tests {
         })
         .unwrap();
 
-        let post_state = HashedPostState {
-            accounts: Default::default(),
-            storages: BTreeMap::from([(
-                address,
-                HashedStorage { wiped: false, storage: storage.clone() },
-            )]),
-        };
+        let wiped = false;
+        let mut hashed_storage = HashedStorage::new(wiped);
+        for (slot, value) in storage.iter() {
+            hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+        }
+
+        let mut hashed_post_state = HashedPostState::default();
+        hashed_post_state.insert_hashed_storage(address, hashed_storage);
+        hashed_post_state.sort();
 
         let tx = db.tx().unwrap();
-        let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+        let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
         let expected = [(address, storage)].into_iter();
         assert_storage_cursor_order(&factory, expected);
     }
@@ -881,14 +1021,21 @@ mod tests {
             })
             .unwrap();
 
-            let post_state = HashedPostState {
-                accounts: Default::default(),
-                storages: BTreeMap::from_iter(post_state_storages.iter().map(
-                    |(address, (wiped, storage))| {
-                        (*address, HashedStorage { wiped: *wiped, storage: storage.clone() })
-                    },
-                )),
-            };
+            let mut hashed_post_state = HashedPostState::default();
+
+            for (address, (wiped, storage)) in &post_state_storages {
+                let mut hashed_storage = HashedStorage::new(*wiped);
+                for (slot, value) in storage {
+                    if *value == U256::ZERO {
+                        hashed_storage.insert_zero_valued_slot(*slot);
+                    } else {
+                        hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+                    }
+                }
+                hashed_post_state.insert_hashed_storage(*address, hashed_storage);
+            }
+
+            hashed_post_state.sort();
 
             let mut expected = db_storages;
             // overwrite or remove accounts from the expected result
@@ -901,7 +1048,7 @@ mod tests {
             }
 
             let tx = db.tx().unwrap();
-            let factory = HashedPostStateCursorFactory::new(&tx, &post_state);
+            let factory = HashedPostStateCursorFactory::new(&tx, &hashed_post_state);
             assert_storage_cursor_order(&factory, expected.into_iter());
         });
     }