feat(trie): update sparse trie storage roots independently (#14874)

crates/engine/tree/src/tree/payload_processor/sparse_trie.rs

@@ -130,7 +130,7 @@ pub(super) type SparseTrieEvent = SparseTrieUpdate;
 /// Updates the sparse trie with the given proofs and state, and returns the elapsed time.
 pub(crate) fn update_sparse_trie<BPF>(
     trie: &mut SparseStateTrie<BPF>,
-    SparseTrieUpdate { state, multiproof }: SparseTrieUpdate,
+    SparseTrieUpdate { mut state, multiproof }: SparseTrieUpdate,
 ) -> SparseStateTrieResult<Duration>
 where
     BPF: BlindedProviderFactory + Send + Sync,
@@ -178,12 +178,25 @@ where
         })
         .for_each_init(|| tx.clone(), |tx, result| tx.send(result).unwrap());
     drop(tx);
+
+    // Update account storage roots
     for result in rx {
         let (address, storage_trie) = result?;
         trie.insert_storage_trie(address, storage_trie);
+
+        if let Some(account) = state.accounts.remove(&address) {
+            // If the account itself has an update, remove it from the state update and update in
+            // one go instead of doing it down below.
+            trace!(target: "engine::root::sparse", ?address, "Updating account and its storage root");
+            trie.update_account(address, account.unwrap_or_default())?;
+        } else if trie.is_account_revealed(address) {
+            // Otherwise, if the account is revealed, only update its storage root.
+            trace!(target: "engine::root::sparse", ?address, "Updating account storage root");
+            trie.update_account_storage_root(address)?;
+        }
     }
 
-    // Update accounts with new values
+    // Update accounts
     for (address, account) in state.accounts {
         trace!(target: "engine::root::sparse", ?address, "Updating account");
         trie.update_account(address, account.unwrap_or_default())?;

crates/trie/sparse/src/state.rs

@@ -602,14 +602,14 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
     /// 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<()> {
         let nibbles = Nibbles::unpack(address);
+
         let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
             trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
             storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
         } else if self.is_account_revealed(address) {
             trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account");
-            let state = self.state.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?;
             // The account was revealed, either...
-            if let Some(value) = state.get_leaf_value(&nibbles) {
+            if let Some(value) = self.get_account_value(&address) {
                 // ..it exists and we should take it's current storage root or...
                 TrieAccount::decode(&mut &value[..])?.storage_root
             } else {
@@ -631,6 +631,54 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         }
     }
 
+    /// Update the storage root of a revealed account.
+    ///
+    /// If the account doesn't exist in the trie, the function is a no-op.
+    ///
+    /// If the new storage root is empty, and the account info was already empty, the account leaf
+    /// will be removed.
+    pub fn update_account_storage_root(&mut self, address: B256) -> SparseStateTrieResult<()> {
+        if !self.is_account_revealed(address) {
+            return Err(SparseTrieErrorKind::Blind.into())
+        }
+
+        // Nothing to update if the account doesn't exist in the trie.
+        let Some(mut trie_account) = self
+            .get_account_value(&address)
+            .map(|v| TrieAccount::decode(&mut &v[..]))
+            .transpose()?
+        else {
+            return Ok(())
+        };
+
+        // Calculate the new storage root. If the storage trie doesn't exist, the storage root will
+        // be empty.
+        let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
+            trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
+            storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
+        } else {
+            EMPTY_ROOT_HASH
+        };
+
+        // Update the account with the new storage root.
+        trie_account.storage_root = storage_root;
+
+        let nibbles = Nibbles::unpack(address);
+        if trie_account == TrieAccount::default() {
+            // If the account is empty, remove it.
+            trace!(target: "trie::sparse", ?address, "Removing account because the storage root is empty");
+            self.remove_account_leaf(&nibbles)?;
+        } else {
+            // Otherwise, update the account leaf.
+            trace!(target: "trie::sparse", ?address, "Updating account with the new storage root");
+            self.account_rlp_buf.clear();
+            trie_account.encode(&mut self.account_rlp_buf);
+            self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?;
+        }
+
+        Ok(())
+    }
+
     /// Remove the account leaf node.
     pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> {
         self.state.remove_leaf(path)?;