Skip to content

Commit

Permalink
fix(trie): reveal extension child in sparse trie when updating a leaf (…
Browse files Browse the repository at this point in the history 
…#13183)
  • Loading branch information
@shekhirin
shekhirin authored Dec 9, 2024
1 parent 3bc7b00 commit da5ffc2
Show file tree
Hide file tree
Showing 2 changed files with 141 additions and 123 deletions.
48 changes: 24 additions & 24 deletions crates/trie/sparse/src/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -273,30 +273,6 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
Ok(Some(root_node))
}

/// Update the account leaf node.
pub fn update_account_leaf(
&mut self,
path: Nibbles,
value: Vec<u8>,
) -> 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<u8>,
) -> 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) {
Expand Down Expand Up @@ -354,6 +330,30 @@ where
SparseTrieError: From<<F::AccountNodeProvider as BlindedProvider>::Error>
+ From<<F::StorageNodeProvider as BlindedProvider>::Error>,
{
/// Update the account leaf node.
pub fn update_account_leaf(
&mut self,
path: Nibbles,
value: Vec<u8>,
) -> 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<u8>,
) -> 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.
///
Expand Down
216 changes: 117 additions & 99 deletions crates/trie/sparse/src/trie.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,6 @@ impl<P> SparseTrie<P> {
Ok(self.as_revealed_mut().unwrap())
}

/// Update the leaf node.
pub fn update_leaf(&mut self, path: Nibbles, value: Vec<u8>) -> 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)?;
Expand All @@ -134,6 +127,13 @@ where
P: BlindedProvider,
SparseTrieError: From<P::Error>,
{
/// Update the leaf node.
pub fn update_leaf(&mut self, path: Nibbles, value: Vec<u8>) -> 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)?;
Expand Down Expand Up @@ -372,98 +372,6 @@ impl<P> RevealedSparseTrie<P> {
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<u8>) -> 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(&current) {
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(&current) {
// 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<Vec<RemovedSparseNode>> {
let mut current = Nibbles::default(); // Start traversal from the root
Expand Down Expand Up @@ -866,6 +774,116 @@ where
P: BlindedProvider,
SparseTrieError: From<P::Error>,
{
/// Update the leaf node with provided value.
pub fn update_leaf(&mut self, path: Nibbles, value: Vec<u8>) -> 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(&current) {
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(&current) {
// 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(&current).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() {
Expand Down

0 comments on commit da5ffc2

Please sign in to comment.