From 4f4db67bc1464ca56454af6f7feacd92df1c6389 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com> Date: Fri, 7 Mar 2025 12:00:49 +0000 Subject: [PATCH] feat(trie): update sparse trie storage roots independently (#14874) --- .../src/tree/payload_processor/sparse_trie.rs | 17 +++++- crates/trie/sparse/src/state.rs | 52 ++++++++++++++++++- 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs index 6b3c45eb34..704ffb1b1c 100644 --- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs +++ b/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( trie: &mut SparseStateTrie, - SparseTrieUpdate { state, multiproof }: SparseTrieUpdate, + SparseTrieUpdate { mut state, multiproof }: SparseTrieUpdate, ) -> SparseStateTrieResult 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())?; diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 2d2f162711..3b68e36e2f 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -602,14 +602,14 @@ impl SparseStateTrie { /// 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 SparseStateTrie { } } + /// 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)?;