feat(trie): SparseStateTrie::update_account

crates/trie/sparse/Cargo.toml

@@ -14,8 +14,9 @@ workspace = true
 
 [dependencies]
 # reth
-reth-tracing.workspace = true
+reth-primitives-traits.workspace = true
 reth-trie-common.workspace = true
+reth-tracing.workspace = true
 
 # alloy
 alloy-primitives.workspace = true

crates/trie/sparse/src/state.rs

@@ -1,19 +1,22 @@
-use crate::{RevealedSparseTrie, SparseStateTrieError, SparseStateTrieResult, SparseTrie};
+use crate::{
+    RevealedSparseTrie, SparseStateTrieError, SparseStateTrieResult, SparseTrie, SparseTrieError,
+};
 use alloy_primitives::{
     map::{HashMap, HashSet},
     Bytes, B256,
 };
-use alloy_rlp::Decodable;
+use alloy_rlp::{Decodable, Encodable};
+use reth_primitives_traits::Account;
+use reth_trie::TRIE_ACCOUNT_RLP_MAX_SIZE;
 use reth_trie_common::{
     updates::{StorageTrieUpdates, TrieUpdates},
-    MultiProof, Nibbles, TrieNode,
+    MultiProof, Nibbles, TrieAccount, TrieNode, EMPTY_ROOT_HASH,
 };
 use std::iter::Peekable;
 
 /// Sparse state trie representing lazy-loaded Ethereum state trie.
 #[derive(Default, Debug)]
 pub struct SparseStateTrie {
-    retain_updates: bool,
     /// Sparse account trie.
     state: SparseTrie,
     /// Sparse storage tries.
@@ -22,6 +25,23 @@
     revealed: HashMap<B256, HashSet<B256>>,
     /// Collection of addresses that had their storage tries wiped.
     wiped_storages: HashSet<B256>,
+    /// Flag indicating whether trie updates should be retained.
+    retain_updates: bool,
+    /// Re-usable buffer for RLP encoding of trie accounts.
+    account_rlp_buf: Vec<u8>,
+}
+
+impl Default for SparseStateTrie {
+    fn default() -> Self {
+        Self {
+            state: Default::default(),
+            storages: Default::default(),
+            revealed: Default::default(),
+            wiped_storages: Default::default(),
+            retain_updates: false,
+            account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
+        }
+    }
 }
 
 impl SparseStateTrie {
@@ -186,6 +206,31 @@
         Ok(Some(root_node))
     }
 
+    /// Update trie account with new account info. This method will either recompute the storage
+    /// root based on update storage trie or look it up from existing leaf value.
+    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) {
+            storage_trie.root().ok_or(SparseTrieError::Blind)?
+        } else if self.revealed.contains_key(&address) {
+            let state = self.state.as_revealed_mut().ok_or(SparseTrieError::Blind)?;
+            // The account was revealed, either...
+            if let Some(value) = state.get_leaf_value(&nibbles) {
+                // ..it exists and we should take it's current storage root or...
+                TrieAccount::decode(&mut &value[..])?.storage_root
+            } else {
+                // ...the account is newly created and the storage trie is empty.
+                EMPTY_ROOT_HASH
+            }
+        } else {
+            return Err(SparseTrieError::Blind.into())
+        };
+
+        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())
+    }
+
     /// Update the account leaf node.
     pub fn update_account_leaf(
         &mut self,

crates/trie/sparse/src/trie.rs

@@ -166,6 +166,11 @@ impl RevealedSparseTrie {
         self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed)
     }
 
+    /// Returns a reference to the leaf value if present.
+    pub fn get_leaf_value(&self, path: &Nibbles) -> Option<&Vec<u8>> {
+        self.values.get(path)
+    }
+
     /// Takes and returns the retained sparse node updates
     pub fn take_updates(&mut self) -> SparseTrieUpdates {
         self.updates.take().unwrap_or_default()