diff --git a/crates/trie/sparse/src/blinded.rs b/crates/trie/sparse/src/blinded.rs new file mode 100644 index 000000000..3875e819d --- /dev/null +++ b/crates/trie/sparse/src/blinded.rs @@ -0,0 +1,58 @@ +//! Traits and default implementations related to retrieval of blinded trie nodes. + +use crate::SparseTrieError; +use alloy_primitives::Bytes; +use reth_trie_common::Nibbles; +use std::convert::Infallible; + +/// Factory for instantiating blinded node providers. +pub trait BlindedProviderFactory { + /// Type capable of fetching blinded account nodes. + type AccountNodeProvider: BlindedProvider; + /// Type capable of fetching blinded storage nodes. + type StorageNodeProvider: BlindedProvider; + + /// Returns blinded account node provider. + fn account_node_provider(&self) -> Self::AccountNodeProvider; + + /// Returns blinded storage node provider. + fn storage_node_provider(&self) -> Self::StorageNodeProvider; +} + +/// Trie node provider for retrieving blinded nodes. +pub trait BlindedProvider { + /// The error type for the provider. + type Error: Into; + + /// Retrieve blinded node by path. + fn blinded_node(&mut self, path: Nibbles) -> Result, Self::Error>; +} + +/// Default blinded node provider factory that creates [`DefaultBlindedProvider`]. +#[derive(PartialEq, Eq, Clone, Default, Debug)] +pub struct DefaultBlindedProviderFactory; + +impl BlindedProviderFactory for DefaultBlindedProviderFactory { + type AccountNodeProvider = DefaultBlindedProvider; + type StorageNodeProvider = DefaultBlindedProvider; + + fn account_node_provider(&self) -> Self::AccountNodeProvider { + DefaultBlindedProvider + } + + fn storage_node_provider(&self) -> Self::StorageNodeProvider { + DefaultBlindedProvider + } +} + +/// Default blinded node provider that always returns `Ok(None)`. +#[derive(PartialEq, Eq, Clone, Default, Debug)] +pub struct DefaultBlindedProvider; + +impl BlindedProvider for DefaultBlindedProvider { + type Error = Infallible; + + fn blinded_node(&mut self, _path: Nibbles) -> Result, Self::Error> { + Ok(None) + } +} diff --git a/crates/trie/sparse/src/errors.rs b/crates/trie/sparse/src/errors.rs index a38a92395..ca3b279ce 100644 --- a/crates/trie/sparse/src/errors.rs +++ b/crates/trie/sparse/src/errors.rs @@ -1,5 +1,7 @@ //! Errors for sparse trie. +use std::convert::Infallible; + use alloy_primitives::{Bytes, B256}; use reth_trie_common::Nibbles; use thiserror::Error; @@ -56,4 +58,7 @@ pub enum SparseTrieError { /// RLP error. #[error(transparent)] Rlp(#[from] alloy_rlp::Error), + /// Infallible. + #[error(transparent)] + Infallible(#[from] Infallible), } diff --git a/crates/trie/sparse/src/lib.rs b/crates/trie/sparse/src/lib.rs index b3cb2c5fd..ec5117fdb 100644 --- a/crates/trie/sparse/src/lib.rs +++ b/crates/trie/sparse/src/lib.rs @@ -8,3 +8,5 @@ pub use trie::*; mod errors; pub use errors::*; + +pub mod blinded; diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 6444c7cd2..877744954 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -1,4 +1,5 @@ use crate::{ + blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory}, RevealedSparseTrie, SparseStateTrieError, SparseStateTrieResult, SparseTrie, SparseTrieError, }; use alloy_primitives::{ @@ -16,11 +17,13 @@ use std::iter::Peekable; /// Sparse state trie representing lazy-loaded Ethereum state trie. #[derive(Debug)] -pub struct SparseStateTrie { +pub struct SparseStateTrie { + /// Blinded node provider factory. + provider_factory: F, /// Sparse account trie. - state: SparseTrie, + state: SparseTrie, /// Sparse storage tries. - storages: HashMap, + storages: HashMap>, /// Collection of revealed account and storage keys. revealed: HashMap>, /// Flag indicating whether trie updates should be retained. @@ -35,6 +38,7 @@ impl Default for SparseStateTrie { state: Default::default(), storages: Default::default(), revealed: Default::default(), + provider_factory: Default::default(), retain_updates: false, account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE), } @@ -46,7 +50,9 @@ impl SparseStateTrie { pub fn from_state(state: SparseTrie) -> Self { Self { state, ..Default::default() } } +} +impl SparseStateTrie { /// Set the retention of branch node updates and deletions. pub const fn with_updates(mut self, retain_updates: bool) -> Self { self.retain_updates = retain_updates; @@ -64,7 +70,10 @@ impl SparseStateTrie { } /// Returns mutable reference to storage sparse trie if it was revealed. - pub fn storage_trie_mut(&mut self, account: &B256) -> Option<&mut RevealedSparseTrie> { + pub fn storage_trie_mut( + &mut self, + account: &B256, + ) -> Option<&mut RevealedSparseTrie> { self.storages.get_mut(account).and_then(|e| e.as_revealed_mut()) } @@ -84,7 +93,11 @@ impl SparseStateTrie { let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) }; // Reveal root node if it wasn't already. - let trie = self.state.reveal_root(root_node, self.retain_updates)?; + let trie = self.state.reveal_root_with_provider( + self.provider_factory.account_node_provider(), + root_node, + self.retain_updates, + )?; // Reveal the remaining proof nodes. for (path, bytes) in proof { @@ -115,11 +128,11 @@ impl SparseStateTrie { let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) }; // Reveal root node if it wasn't already. - let trie = self - .storages - .entry(account) - .or_default() - .reveal_root(root_node, self.retain_updates)?; + let trie = self.storages.entry(account).or_default().reveal_root_with_provider( + self.provider_factory.storage_node_provider(), + root_node, + self.retain_updates, + )?; // Reveal the remaining proof nodes. for (path, bytes) in proof { @@ -145,7 +158,11 @@ impl SparseStateTrie { if let Some(root_node) = self.validate_root_node(&mut account_nodes)? { // Reveal root node if it wasn't already. - let trie = self.state.reveal_root(root_node, self.retain_updates)?; + let trie = self.state.reveal_root_with_provider( + self.provider_factory.account_node_provider(), + root_node, + self.retain_updates, + )?; // Reveal the remaining proof nodes. for (path, bytes) in account_nodes { @@ -161,11 +178,11 @@ impl SparseStateTrie { if let Some(root_node) = self.validate_root_node(&mut storage_nodes)? { // Reveal root node if it wasn't already. - let trie = self - .storages - .entry(account) - .or_default() - .reveal_root(root_node, self.retain_updates)?; + let trie = self.storages.entry(account).or_default().reveal_root_with_provider( + self.provider_factory.storage_node_provider(), + root_node, + self.retain_updates, + )?; // Reveal the remaining proof nodes. for (path, bytes) in storage_nodes { @@ -205,41 +222,6 @@ impl SparseStateTrie { Ok(Some(root_node)) } - /// Update or remove trie account based on new account info. This method will either recompute - /// 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<()> { - 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(SparseTrieError::Blind)? - } else if self.revealed.contains_key(&address) { - trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account"); - 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()) - }; - - if account.is_empty() && storage_root == EMPTY_ROOT_HASH { - trace!(target: "trie::sparse", ?address, "Removing account"); - self.remove_account_leaf(&nibbles) - } else { - trace!(target: "trie::sparse", ?address, "Updating account"); - 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, @@ -250,12 +232,6 @@ impl SparseStateTrie { Ok(()) } - /// Remove the account leaf node. - pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> { - self.state.remove_leaf(path)?; - Ok(()) - } - /// Update the leaf node of a storage trie at the provided address. pub fn update_storage_leaf( &mut self, @@ -263,18 +239,11 @@ impl SparseStateTrie { slot: Nibbles, value: Vec, ) -> SparseStateTrieResult<()> { - self.storages.entry(address).or_default().update_leaf(slot, value)?; - Ok(()) - } - - /// Update the leaf node of a storage trie at the provided address. - pub fn remove_storage_leaf( - &mut self, - address: B256, - slot: &Nibbles, - ) -> SparseStateTrieResult<()> { - self.storages.entry(address).or_default().remove_leaf(slot)?; - Ok(()) + if let Some(storage_trie) = self.storages.get_mut(&address) { + Ok(storage_trie.update_leaf(slot, value)?) + } else { + Err(SparseStateTrieError::Sparse(SparseTrieError::Blind)) + } } /// Wipe the storage trie at the provided address. @@ -329,6 +298,67 @@ impl SparseStateTrie { } } +impl SparseStateTrie +where + F: BlindedProviderFactory, + SparseTrieError: From<::Error> + + From<::Error>, +{ + /// Update or remove trie account based on new account info. This method will either recompute + /// 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<()> { + 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(SparseTrieError::Blind)? + } else if self.revealed.contains_key(&address) { + trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account"); + 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()) + }; + + if account.is_empty() && storage_root == EMPTY_ROOT_HASH { + trace!(target: "trie::sparse", ?address, "Removing account"); + self.remove_account_leaf(&nibbles) + } else { + trace!(target: "trie::sparse", ?address, "Updating account"); + 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()) + } + } + + /// Remove the account leaf node. + pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> { + self.state.remove_leaf(path)?; + Ok(()) + } + + /// Update the leaf node of a storage trie at the provided address. + pub fn remove_storage_leaf( + &mut self, + address: B256, + slot: &Nibbles, + ) -> SparseStateTrieResult<()> { + if let Some(storage_trie) = self.storages.get_mut(&address) { + Ok(storage_trie.remove_leaf(slot)?) + } else { + Err(SparseStateTrieError::Sparse(SparseTrieError::Blind)) + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 12a0f87e1..0dd6dc989 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -1,4 +1,7 @@ -use crate::{SparseTrieError, SparseTrieResult}; +use crate::{ + blinded::{BlindedProvider, DefaultBlindedProvider}, + SparseTrieError, SparseTrieResult, +}; use alloy_primitives::{ hex, keccak256, map::{HashMap, HashSet}, @@ -16,35 +19,31 @@ use std::{borrow::Cow, fmt}; /// Inner representation of the sparse trie. /// Sparse trie is blind by default until nodes are revealed. -#[derive(PartialEq, Eq, Default, Debug)] -pub enum SparseTrie { +#[derive(PartialEq, Eq, Debug)] +pub enum SparseTrie

{ /// None of the trie nodes are known. - #[default] Blind, /// The trie nodes have been revealed. - Revealed(Box), + Revealed(Box>), +} + +impl

Default for SparseTrie

{ + fn default() -> Self { + Self::Blind + } } impl SparseTrie { + /// Creates new blind trie. + pub const fn blind() -> Self { + Self::Blind + } + /// Creates new revealed empty trie. pub fn revealed_empty() -> Self { Self::Revealed(Box::default()) } - /// Returns `true` if the sparse trie has no revealed nodes. - pub const fn is_blind(&self) -> bool { - matches!(self, Self::Blind) - } - - /// Returns mutable reference to revealed sparse trie if the trie is not blind. - pub fn as_revealed_mut(&mut self) -> Option<&mut RevealedSparseTrie> { - if let Self::Revealed(revealed) = self { - Some(revealed) - } else { - None - } - } - /// Reveals the root node if the trie is blinded. /// /// # Returns @@ -55,8 +54,42 @@ impl SparseTrie { root: TrieNode, retain_updates: bool, ) -> SparseTrieResult<&mut RevealedSparseTrie> { + self.reveal_root_with_provider(Default::default(), root, retain_updates) + } +} + +impl

SparseTrie

{ + /// Returns `true` if the sparse trie has no revealed nodes. + pub const fn is_blind(&self) -> bool { + matches!(self, Self::Blind) + } + + /// Returns mutable reference to revealed sparse trie if the trie is not blind. + pub fn as_revealed_mut(&mut self) -> Option<&mut RevealedSparseTrie

> { + if let Self::Revealed(revealed) = self { + Some(revealed) + } else { + None + } + } + + /// Reveals the root node if the trie is blinded. + /// + /// # Returns + /// + /// Mutable reference to [`RevealedSparseTrie`]. + pub fn reveal_root_with_provider( + &mut self, + provider: P, + root: TrieNode, + retain_updates: bool, + ) -> SparseTrieResult<&mut RevealedSparseTrie

> { if self.is_blind() { - *self = Self::Revealed(Box::new(RevealedSparseTrie::from_root(root, retain_updates)?)) + *self = Self::Revealed(Box::new(RevealedSparseTrie::from_provider_and_root( + provider, + root, + retain_updates, + )?)) } Ok(self.as_revealed_mut().unwrap()) } @@ -68,13 +101,6 @@ impl SparseTrie { Ok(()) } - /// Remove the leaf node. - pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { - let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; - revealed.remove_leaf(path)?; - Ok(()) - } - /// Wipe the trie, removing all values and nodes, and replacing the root with an empty node. pub fn wipe(&mut self) -> SparseTrieResult<()> { let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; @@ -93,6 +119,19 @@ impl SparseTrie { } } +impl

SparseTrie

+where + P: BlindedProvider, + SparseTrieError: From, +{ + /// Remove the leaf node. + pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { + let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; + revealed.remove_leaf(path)?; + Ok(()) + } +} + /// The representation of revealed sparse trie. /// /// ## Invariants @@ -102,27 +141,29 @@ impl SparseTrie { /// The opposite is also true. /// - All keys in `values` collection are full leaf paths. #[derive(Clone, PartialEq, Eq)] -pub struct RevealedSparseTrie { +pub struct RevealedSparseTrie

{ + /// Blinded node provider. + provider: P, /// All trie nodes. nodes: HashMap, /// All leaf values. values: HashMap>, /// Prefix set. prefix_set: PrefixSetMut, - /// Reusable buffer for RLP encoding of nodes. - rlp_buf: Vec, /// Retained trie updates. updates: Option, + /// Reusable buffer for RLP encoding of nodes. + rlp_buf: Vec, } -impl fmt::Debug for RevealedSparseTrie { +impl

fmt::Debug for RevealedSparseTrie

{ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("RevealedSparseTrie") .field("nodes", &self.nodes) .field("values", &self.values) .field("prefix_set", &self.prefix_set) - .field("rlp_buf", &hex::encode(&self.rlp_buf)) .field("updates", &self.updates) + .field("rlp_buf", &hex::encode(&self.rlp_buf)) .finish() } } @@ -130,11 +171,12 @@ impl fmt::Debug for RevealedSparseTrie { impl Default for RevealedSparseTrie { fn default() -> Self { Self { + provider: Default::default(), nodes: HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]), values: HashMap::default(), prefix_set: PrefixSetMut::default(), - rlp_buf: Vec::new(), updates: None, + rlp_buf: Vec::new(), } } } @@ -143,6 +185,7 @@ impl RevealedSparseTrie { /// Create new revealed sparse trie from the given root node. pub fn from_root(node: TrieNode, retain_updates: bool) -> SparseTrieResult { let mut this = Self { + provider: Default::default(), nodes: HashMap::default(), values: HashMap::default(), prefix_set: PrefixSetMut::default(), @@ -153,6 +196,39 @@ impl RevealedSparseTrie { this.reveal_node(Nibbles::default(), node)?; Ok(this) } +} + +impl

RevealedSparseTrie

{ + /// Create new revealed sparse trie from the given root node. + pub fn from_provider_and_root( + provider: P, + node: TrieNode, + retain_updates: bool, + ) -> SparseTrieResult { + let mut this = Self { + provider, + nodes: HashMap::default(), + values: HashMap::default(), + prefix_set: PrefixSetMut::default(), + rlp_buf: Vec::new(), + updates: None, + } + .with_updates(retain_updates); + this.reveal_node(Nibbles::default(), node)?; + Ok(this) + } + + /// Set new blinded node provider on sparse trie. + pub fn with_provider(self, provider: BP) -> RevealedSparseTrie { + RevealedSparseTrie { + provider, + nodes: self.nodes, + values: self.values, + prefix_set: self.prefix_set, + updates: self.updates, + rlp_buf: self.rlp_buf, + } + } /// Set the retention of branch node updates and deletions. pub fn with_updates(mut self, retain_updates: bool) -> Self { @@ -357,176 +433,6 @@ impl RevealedSparseTrie { Ok(()) } - /// Remove leaf node from the trie. - pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { - self.prefix_set.insert(path.clone()); - self.values.remove(path); - - // If the path wasn't present in `values`, we still need to walk the trie and ensure that - // there is no node at the path. When a leaf node is a blinded `Hash`, it will have an entry - // in `nodes`, but not in the `values`. - - // If the path wasn't present in `values`, we still need to walk the trie and ensure that - // there is no node at the path. When a leaf node is a blinded `Hash`, it will have an entry - // in `nodes`, but not in the `values`. - - let mut removed_nodes = self.take_nodes_for_path(path)?; - trace!(target: "trie::sparse", ?path, ?removed_nodes, "Removed nodes for path"); - // Pop the first node from the stack which is the leaf node we want to remove. - let mut child = removed_nodes.pop().expect("leaf exists"); - #[cfg(debug_assertions)] - { - let mut child_path = child.path.clone(); - let SparseNode::Leaf { key, .. } = &child.node else { panic!("expected leaf node") }; - child_path.extend_from_slice_unchecked(key); - assert_eq!(&child_path, path); - } - - // If we don't have any other removed nodes, insert an empty node at the root. - if removed_nodes.is_empty() { - debug_assert!(self.nodes.is_empty()); - self.nodes.insert(Nibbles::default(), SparseNode::Empty); - - return Ok(()) - } - - // Walk the stack of removed nodes from the back and re-insert them back into the trie, - // adjusting the node type as needed. - while let Some(removed_node) = removed_nodes.pop() { - let removed_path = removed_node.path; - - let new_node = match &removed_node.node { - SparseNode::Empty => return Err(SparseTrieError::Blind), - SparseNode::Hash(hash) => { - return Err(SparseTrieError::BlindedNode { path: removed_path, hash: *hash }) - } - SparseNode::Leaf { .. } => { - unreachable!("we already popped the leaf node") - } - SparseNode::Extension { key, .. } => { - // If the node is an extension node, we need to look at its child to see if we - // need to merge them. - match &child.node { - SparseNode::Empty => return Err(SparseTrieError::Blind), - SparseNode::Hash(hash) => { - return Err(SparseTrieError::BlindedNode { - path: child.path, - hash: *hash, - }) - } - // For a leaf node, we collapse the extension node into a leaf node, - // extending the key. While it's impossible to encounter an extension node - // followed by a leaf node in a complete trie, it's possible here because we - // could have downgraded the extension node's child into a leaf node from - // another node type. - SparseNode::Leaf { key: leaf_key, .. } => { - self.nodes.remove(&child.path); - - let mut new_key = key.clone(); - new_key.extend_from_slice_unchecked(leaf_key); - SparseNode::new_leaf(new_key) - } - // For an extension node, we collapse them into one extension node, - // extending the key - SparseNode::Extension { key: extension_key, .. } => { - self.nodes.remove(&child.path); - - let mut new_key = key.clone(); - new_key.extend_from_slice_unchecked(extension_key); - SparseNode::new_ext(new_key) - } - // For a branch node, we just leave the extension node as-is. - SparseNode::Branch { .. } => removed_node.node, - } - } - SparseNode::Branch { mut state_mask, hash: _, store_in_db_trie: _ } => { - // If the node is a branch node, we need to check the number of children left - // after deleting the child at the given nibble. - - if let Some(removed_nibble) = removed_node.unset_branch_nibble { - state_mask.unset_bit(removed_nibble); - } - - // If only one child is left set in the branch node, we need to collapse it. - if state_mask.count_bits() == 1 { - let child_nibble = - state_mask.first_set_bit_index().expect("state mask is not empty"); - - // Get full path of the only child node left. - let mut child_path = removed_path.clone(); - child_path.push_unchecked(child_nibble); - - // Remove the only child node. - let child = self.nodes.get(&child_path).unwrap(); - - trace!(target: "trie::sparse", ?removed_path, ?child_path, ?child, "Branch node has only one child"); - - let mut delete_child = false; - let new_node = match child { - SparseNode::Empty => return Err(SparseTrieError::Blind), - SparseNode::Hash(hash) => { - return Err(SparseTrieError::BlindedNode { - path: child_path, - hash: *hash, - }) - } - // If the only child is a leaf node, we downgrade the branch node into a - // leaf node, prepending the nibble to the key, and delete the old - // child. - SparseNode::Leaf { key, .. } => { - delete_child = true; - - let mut new_key = Nibbles::from_nibbles_unchecked([child_nibble]); - new_key.extend_from_slice_unchecked(key); - SparseNode::new_leaf(new_key) - } - // If the only child node is an extension node, we downgrade the branch - // node into an even longer extension node, prepending the nibble to the - // key, and delete the old child. - SparseNode::Extension { key, .. } => { - delete_child = true; - - let mut new_key = Nibbles::from_nibbles_unchecked([child_nibble]); - new_key.extend_from_slice_unchecked(key); - SparseNode::new_ext(new_key) - } - // If the only child is a branch node, we downgrade the current branch - // node into a one-nibble extension node. - SparseNode::Branch { .. } => { - SparseNode::new_ext(Nibbles::from_nibbles_unchecked([child_nibble])) - } - }; - - if delete_child { - self.nodes.remove(&child_path); - } - - if let Some(updates) = self.updates.as_mut() { - updates.removed_nodes.insert(removed_path.clone()); - } - - new_node - } - // If more than one child is left set in the branch, we just re-insert it - // as-is. - else { - SparseNode::new_branch(state_mask) - } - } - }; - - child = RemovedSparseNode { - path: removed_path.clone(), - node: new_node.clone(), - unset_branch_nibble: None, - }; - trace!(target: "trie::sparse", ?removed_path, ?new_node, "Re-inserting the node"); - self.nodes.insert(removed_path, new_node); - } - - Ok(()) - } - /// Traverse trie nodes down to the leaf node and collect all nodes along the path. fn take_nodes_for_path(&mut self, path: &Nibbles) -> SparseTrieResult> { let mut current = Nibbles::default(); // Start traversal from the root @@ -621,10 +527,10 @@ impl RevealedSparseTrie { /// Wipe the trie, removing all values and nodes, and replacing the root with an empty node. pub fn wipe(&mut self) { - let updates_retained = self.updates.is_some(); - *self = Self::default(); + self.nodes = HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]); + self.values = HashMap::default(); self.prefix_set = PrefixSetMut::all(); - self.updates = updates_retained.then(SparseTrieUpdates::wiped); + self.updates = self.updates.is_some().then(SparseTrieUpdates::wiped); } /// Return the root of the sparse trie. @@ -901,6 +807,191 @@ impl RevealedSparseTrie { } } +impl

RevealedSparseTrie

+where + P: BlindedProvider, + SparseTrieError: From, +{ + /// Remove leaf node from the trie. + pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { + self.prefix_set.insert(path.clone()); + self.values.remove(path); + + // If the path wasn't present in `values`, we still need to walk the trie and ensure that + // there is no node at the path. When a leaf node is a blinded `Hash`, it will have an entry + // in `nodes`, but not in the `values`. + + // If the path wasn't present in `values`, we still need to walk the trie and ensure that + // there is no node at the path. When a leaf node is a blinded `Hash`, it will have an entry + // in `nodes`, but not in the `values`. + + let mut removed_nodes = self.take_nodes_for_path(path)?; + trace!(target: "trie::sparse", ?path, ?removed_nodes, "Removed nodes for path"); + // Pop the first node from the stack which is the leaf node we want to remove. + let mut child = removed_nodes.pop().expect("leaf exists"); + #[cfg(debug_assertions)] + { + let mut child_path = child.path.clone(); + let SparseNode::Leaf { key, .. } = &child.node else { panic!("expected leaf node") }; + child_path.extend_from_slice_unchecked(key); + assert_eq!(&child_path, path); + } + + // If we don't have any other removed nodes, insert an empty node at the root. + if removed_nodes.is_empty() { + debug_assert!(self.nodes.is_empty()); + self.nodes.insert(Nibbles::default(), SparseNode::Empty); + + return Ok(()) + } + + // Walk the stack of removed nodes from the back and re-insert them back into the trie, + // adjusting the node type as needed. + while let Some(removed_node) = removed_nodes.pop() { + let removed_path = removed_node.path; + + let new_node = match &removed_node.node { + SparseNode::Empty => return Err(SparseTrieError::Blind), + SparseNode::Hash(hash) => { + return Err(SparseTrieError::BlindedNode { path: removed_path, hash: *hash }) + } + SparseNode::Leaf { .. } => { + unreachable!("we already popped the leaf node") + } + SparseNode::Extension { key, .. } => { + // If the node is an extension node, we need to look at its child to see if we + // need to merge them. + match &child.node { + SparseNode::Empty => return Err(SparseTrieError::Blind), + SparseNode::Hash(hash) => { + return Err(SparseTrieError::BlindedNode { + path: child.path, + hash: *hash, + }) + } + // For a leaf node, we collapse the extension node into a leaf node, + // extending the key. While it's impossible to encounter an extension node + // followed by a leaf node in a complete trie, it's possible here because we + // could have downgraded the extension node's child into a leaf node from + // another node type. + SparseNode::Leaf { key: leaf_key, .. } => { + self.nodes.remove(&child.path); + + let mut new_key = key.clone(); + new_key.extend_from_slice_unchecked(leaf_key); + SparseNode::new_leaf(new_key) + } + // For an extension node, we collapse them into one extension node, + // extending the key + SparseNode::Extension { key: extension_key, .. } => { + self.nodes.remove(&child.path); + + let mut new_key = key.clone(); + new_key.extend_from_slice_unchecked(extension_key); + SparseNode::new_ext(new_key) + } + // For a branch node, we just leave the extension node as-is. + SparseNode::Branch { .. } => removed_node.node, + } + } + SparseNode::Branch { mut state_mask, hash: _, store_in_db_trie: _ } => { + // If the node is a branch node, we need to check the number of children left + // after deleting the child at the given nibble. + + if let Some(removed_nibble) = removed_node.unset_branch_nibble { + state_mask.unset_bit(removed_nibble); + } + + // If only one child is left set in the branch node, we need to collapse it. + if state_mask.count_bits() == 1 { + let child_nibble = + state_mask.first_set_bit_index().expect("state mask is not empty"); + + // Get full path of the only child node left. + let mut child_path = removed_path.clone(); + child_path.push_unchecked(child_nibble); + + trace!(target: "trie::sparse", ?removed_path, ?child_path, ?child, "Branch node has only one child"); + + if self.nodes.get(&child_path).unwrap().is_hash() { + trace!(target: "trie::sparse", ?child_path, "Retrieving remaining blinded branch child"); + if let Some(node) = self.provider.blinded_node(child_path.clone())? { + let decoded = TrieNode::decode(&mut &node[..])?; + trace!(target: "trie::sparse", ?child_path, ?decoded, "Revealing remaining blinded branch child"); + self.reveal_node(child_path.clone(), decoded)?; + } + } + + // Get the only child node. + let child = self.nodes.get(&child_path).unwrap(); + + let mut delete_child = false; + let new_node = match child { + SparseNode::Empty => return Err(SparseTrieError::Blind), + SparseNode::Hash(hash) => { + return Err(SparseTrieError::BlindedNode { + path: child_path, + hash: *hash, + }) + } + // If the only child is a leaf node, we downgrade the branch node into a + // leaf node, prepending the nibble to the key, and delete the old + // child. + SparseNode::Leaf { key, .. } => { + delete_child = true; + + let mut new_key = Nibbles::from_nibbles_unchecked([child_nibble]); + new_key.extend_from_slice_unchecked(key); + SparseNode::new_leaf(new_key) + } + // If the only child node is an extension node, we downgrade the branch + // node into an even longer extension node, prepending the nibble to the + // key, and delete the old child. + SparseNode::Extension { key, .. } => { + delete_child = true; + + let mut new_key = Nibbles::from_nibbles_unchecked([child_nibble]); + new_key.extend_from_slice_unchecked(key); + SparseNode::new_ext(new_key) + } + // If the only child is a branch node, we downgrade the current branch + // node into a one-nibble extension node. + SparseNode::Branch { .. } => { + SparseNode::new_ext(Nibbles::from_nibbles_unchecked([child_nibble])) + } + }; + + if delete_child { + self.nodes.remove(&child_path); + } + + if let Some(updates) = self.updates.as_mut() { + updates.removed_nodes.insert(removed_path.clone()); + } + + new_node + } + // If more than one child is left set in the branch, we just re-insert it + // as-is. + else { + SparseNode::new_branch(state_mask) + } + } + }; + + child = RemovedSparseNode { + path: removed_path.clone(), + node: new_node.clone(), + unset_branch_nibble: None, + }; + trace!(target: "trie::sparse", ?removed_path, ?new_node, "Re-inserting the node"); + self.nodes.insert(removed_path, new_node); + } + + Ok(()) + } +} + /// Enum representing sparse trie node type. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum SparseNodeType { @@ -1007,6 +1098,11 @@ impl SparseNode { pub const fn new_leaf(key: Nibbles) -> Self { Self::Leaf { key, hash: None } } + + /// Returns `true` if the node is a hash node. + pub const fn is_hash(&self) -> bool { + matches!(self, Self::Hash(_)) + } } #[derive(Debug)] @@ -1190,7 +1286,7 @@ mod tests { #[test] fn sparse_trie_is_blind() { - assert!(SparseTrie::default().is_blind()); + assert!(SparseTrie::blind().is_blind()); assert!(!SparseTrie::revealed_empty().is_blind()); }