feat(engine): use SparseStateTrie::update_account in state root task

crates/engine/tree/src/tree/root.rs

@@ -489,11 +489,9 @@ fn update_sparse_trie(
     trie.reveal_multiproof(targets, multiproof)?;
 
     // Update storage slots with new values and calculate storage roots.
-    let mut storage_roots = FbHashMap::default();
     for (address, storage) in state.storages {
         if storage.wiped {
             trie.wipe_storage(address)?;
-            storage_roots.insert(address, EMPTY_ROOT_HASH);
         }
 
         for (slot, value) in storage.storage {
@@ -509,28 +507,11 @@ fn update_sparse_trie(
                 )?;
             }
         }
-
-        storage_roots.insert(address, trie.storage_root(address).unwrap());
     }
 
-    // Update accounts with new values and include updated storage roots
+    // Update accounts with new values
     for (address, account) in state.accounts {
-        let account_nibbles = Nibbles::unpack(address);
-
-        if let Some(account) = account {
-            let storage_root = storage_roots
-                .remove(&address)
-                .map(Some)
-                .unwrap_or_else(|| trie.storage_root(address))
-                .unwrap_or(EMPTY_ROOT_HASH);
-
-            let mut encoded = Vec::with_capacity(128);
-            TrieAccount::from((account, storage_root)).encode(&mut encoded as &mut dyn BufMut);
-            trie.update_account_leaf(account_nibbles, encoded)?;
-        } else {
-            // TODO: handle blinded node error
-            trie.remove_account_leaf(&account_nibbles)?;
-        }
+        trie.update_account(address, account)?;
     }
 
     trie.calculate_below_level(SPARSE_TRIE_INCREMENTAL_LEVEL);

crates/trie/sparse/src/state.rs

@@ -209,7 +209,11 @@ impl SparseStateTrie {
     /// the storage root based on update storage trie or look it up from existing leaf value.
     ///
     /// If the new account info and storage trie are empty, the account leaf will be removed.
-    pub fn update_account(&mut self, address: B256, account: Account) -> SparseStateTrieResult<()> {
+    pub fn update_account(
+        &mut self,
+        address: B256,
+        account: Option<Account>,
+    ) -> SparseStateTrieResult<()> {
         let nibbles = Nibbles::unpack(address);
         let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
             storage_trie.root().ok_or(SparseTrieError::Blind)?
@@ -227,12 +231,13 @@ impl SparseStateTrie {
             return Err(SparseTrieError::Blind.into())
         };
 
-        if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
-            self.remove_account_leaf(&nibbles)
-        } else {
-            self.account_rlp_buf.clear();
-            TrieAccount::from((account, storage_root)).encode(&mut self.account_rlp_buf);
-            self.update_account_leaf(nibbles, self.account_rlp_buf.clone())
+        match account {
+            Some(account) if !account.is_empty() || storage_root != EMPTY_ROOT_HASH => {
+                self.account_rlp_buf.clear();
+                TrieAccount::from((account, storage_root)).encode(&mut self.account_rlp_buf);
+                self.update_account_leaf(nibbles, self.account_rlp_buf.clone())
+            }
+            _ => self.remove_account_leaf(&nibbles),
         }
     }