From da5ffc24c229f8b0b459213b536ad8e7396c7348 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Mon, 9 Dec 2024 15:12:36 +0000 Subject: [PATCH] fix(trie): reveal extension child in sparse trie when updating a leaf (#13183) --- crates/trie/sparse/src/state.rs | 48 +++---- crates/trie/sparse/src/trie.rs | 216 +++++++++++++++++--------------- 2 files changed, 141 insertions(+), 123 deletions(-) diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index ec51df8982c3..6638632f0adc 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -273,30 +273,6 @@ impl SparseStateTrie { Ok(Some(root_node)) } - /// Update the account leaf node. - pub fn update_account_leaf( - &mut self, - path: Nibbles, - value: Vec, - ) -> SparseStateTrieResult<()> { - self.state.update_leaf(path, value)?; - Ok(()) - } - - /// Update the leaf node of a storage trie at the provided address. - pub fn update_storage_leaf( - &mut self, - address: B256, - slot: Nibbles, - value: Vec, - ) -> SparseStateTrieResult<()> { - 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. pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> { if let Some(trie) = self.storages.get_mut(&address) { @@ -354,6 +330,30 @@ where SparseTrieError: From<::Error> + From<::Error>, { + /// Update the account leaf node. + pub fn update_account_leaf( + &mut self, + path: Nibbles, + value: Vec, + ) -> SparseStateTrieResult<()> { + self.state.update_leaf(path, value)?; + Ok(()) + } + + /// Update the leaf node of a storage trie at the provided address. + pub fn update_storage_leaf( + &mut self, + address: B256, + slot: Nibbles, + value: Vec, + ) -> SparseStateTrieResult<()> { + if let Some(storage_trie) = self.storages.get_mut(&address) { + Ok(storage_trie.update_leaf(slot, value)?) + } else { + Err(SparseStateTrieError::Sparse(SparseTrieError::Blind)) + } + } + /// 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. /// diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 8fff0819bcba..e4d4ff701f3e 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -104,13 +104,6 @@ impl

SparseTrie

{ Ok(self.as_revealed_mut().unwrap()) } - /// Update the leaf node. - pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseTrieResult<()> { - let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; - revealed.update_leaf(path, value)?; - 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)?; @@ -134,6 +127,13 @@ where P: BlindedProvider, SparseTrieError: From, { + /// Update the leaf node. + pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseTrieResult<()> { + let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; + revealed.update_leaf(path, value)?; + Ok(()) + } + /// Remove the leaf node. pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; @@ -372,98 +372,6 @@ impl

RevealedSparseTrie

{ self.reveal_node(path, TrieNode::decode(&mut &child[..])?, None) } - /// Update the leaf node with provided value. - pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseTrieResult<()> { - self.prefix_set.insert(path.clone()); - let existing = self.values.insert(path.clone(), value); - if existing.is_some() { - // trie structure unchanged, return immediately - return Ok(()) - } - - let mut current = Nibbles::default(); - while let Some(node) = self.nodes.get_mut(¤t) { - match node { - SparseNode::Empty => { - *node = SparseNode::new_leaf(path); - break - } - SparseNode::Hash(hash) => { - return Err(SparseTrieError::BlindedNode { path: current, hash: *hash }) - } - SparseNode::Leaf { key: current_key, .. } => { - current.extend_from_slice_unchecked(current_key); - - // this leaf is being updated - if current == path { - unreachable!("we already checked leaf presence in the beginning"); - } - - // find the common prefix - let common = current.common_prefix_length(&path); - - // update existing node - let new_ext_key = current.slice(current.len() - current_key.len()..common); - *node = SparseNode::new_ext(new_ext_key); - - // create a branch node and corresponding leaves - self.nodes.insert( - current.slice(..common), - SparseNode::new_split_branch(current[common], path[common]), - ); - self.nodes.insert( - path.slice(..=common), - SparseNode::new_leaf(path.slice(common + 1..)), - ); - self.nodes.insert( - current.slice(..=common), - SparseNode::new_leaf(current.slice(common + 1..)), - ); - - break; - } - SparseNode::Extension { key, .. } => { - current.extend_from_slice(key); - if !path.starts_with(¤t) { - // find the common prefix - let common = current.common_prefix_length(&path); - - *key = current.slice(current.len() - key.len()..common); - - // create state mask for new branch node - // NOTE: this might overwrite the current extension node - let branch = SparseNode::new_split_branch(current[common], path[common]); - self.nodes.insert(current.slice(..common), branch); - - // create new leaf - let new_leaf = SparseNode::new_leaf(path.slice(common + 1..)); - self.nodes.insert(path.slice(..=common), new_leaf); - - // recreate extension to previous child if needed - let key = current.slice(common + 1..); - if !key.is_empty() { - self.nodes.insert(current.slice(..=common), SparseNode::new_ext(key)); - } - - break; - } - } - SparseNode::Branch { state_mask, .. } => { - let nibble = path[current.len()]; - current.push_unchecked(nibble); - if !state_mask.is_bit_set(nibble) { - state_mask.set_bit(nibble); - let new_leaf = SparseNode::new_leaf(path.slice(current.len()..)); - self.nodes.insert(current, new_leaf); - break; - } - } - }; - } - - 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 @@ -866,6 +774,116 @@ where P: BlindedProvider, SparseTrieError: From, { + /// Update the leaf node with provided value. + pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseTrieResult<()> { + self.prefix_set.insert(path.clone()); + let existing = self.values.insert(path.clone(), value); + if existing.is_some() { + // trie structure unchanged, return immediately + return Ok(()) + } + + let mut current = Nibbles::default(); + while let Some(node) = self.nodes.get_mut(¤t) { + match node { + SparseNode::Empty => { + *node = SparseNode::new_leaf(path); + break + } + SparseNode::Hash(hash) => { + return Err(SparseTrieError::BlindedNode { path: current, hash: *hash }) + } + SparseNode::Leaf { key: current_key, .. } => { + current.extend_from_slice_unchecked(current_key); + + // this leaf is being updated + if current == path { + unreachable!("we already checked leaf presence in the beginning"); + } + + // find the common prefix + let common = current.common_prefix_length(&path); + + // update existing node + let new_ext_key = current.slice(current.len() - current_key.len()..common); + *node = SparseNode::new_ext(new_ext_key); + + // create a branch node and corresponding leaves + self.nodes.insert( + current.slice(..common), + SparseNode::new_split_branch(current[common], path[common]), + ); + self.nodes.insert( + path.slice(..=common), + SparseNode::new_leaf(path.slice(common + 1..)), + ); + self.nodes.insert( + current.slice(..=common), + SparseNode::new_leaf(current.slice(common + 1..)), + ); + + break; + } + SparseNode::Extension { key, .. } => { + current.extend_from_slice(key); + + if !path.starts_with(¤t) { + // find the common prefix + let common = current.common_prefix_length(&path); + *key = current.slice(current.len() - key.len()..common); + + // If branch node updates retention is enabled, we need to query the + // extension node child to later set the hash mask for a parent branch node + // correctly. + if self.updates.is_some() { + // Check if the extension node child is a hash that needs to be revealed + if self.nodes.get(¤t).unwrap().is_hash() { + if let Some(node) = self.provider.blinded_node(current.clone())? { + let decoded = TrieNode::decode(&mut &node[..])?; + trace!(target: "trie::sparse", ?current, ?decoded, "Revealing extension node child"); + // We'll never have to update the revealed child node, only + // remove or do nothing, so + // we can safely ignore the hash mask here and + // pass `None`. + self.reveal_node(current.clone(), decoded, None)?; + } + } + } + + // create state mask for new branch node + // NOTE: this might overwrite the current extension node + let branch = SparseNode::new_split_branch(current[common], path[common]); + self.nodes.insert(current.slice(..common), branch); + + // create new leaf + let new_leaf = SparseNode::new_leaf(path.slice(common + 1..)); + self.nodes.insert(path.slice(..=common), new_leaf); + + // recreate extension to previous child if needed + let key = current.slice(common + 1..); + if !key.is_empty() { + self.nodes.insert(current.slice(..=common), SparseNode::new_ext(key)); + } + + break; + } + } + SparseNode::Branch { state_mask, .. } => { + let nibble = path[current.len()]; + current.push_unchecked(nibble); + if !state_mask.is_bit_set(nibble) { + state_mask.set_bit(nibble); + let new_leaf = SparseNode::new_leaf(path.slice(current.len()..)); + self.nodes.insert(current, new_leaf); + break; + } + } + }; + } + + Ok(()) + } + /// Remove leaf node from the trie. pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { if self.values.remove(path).is_none() {