Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(trie): reveal extension child in sparse trie when updating a leaf #13183

Merged
merged 2 commits into from
Dec 9, 2024
Merged

fix(trie): reveal extension child in sparse trie when updating a leaf #13183

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c1487fe
fix(trie): reveal extension child in sparse trie when updating a leaf
shekhirin Dec 6, 2024
522273f
query only if updates are enabled
shekhirin Dec 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 24 additions & 24 deletions crates/trie/sparse/src/state.rs
View file
Open in desktop
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 @@ -355,6 +331,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
210 changes: 111 additions & 99 deletions crates/trie/sparse/src/trie.rs
View file
Open in desktop
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,110 @@ 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);

// Check if the extension node child is a hash that needs to be revealed
if self.nodes.get(&current).unwrap().is_hash() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

only if updates are enabled

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

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
Loading