RevealedSparseTrie
{ /// Return the root of the sparse trie. /// Updates all remaining dirty nodes before calculating the root. pub fn root(&mut self) -> B256 { - // take the current prefix set. + // Take the current prefix set let mut prefix_set = std::mem::take(&mut self.prefix_set).freeze(); - let rlp_node = self.rlp_node_allocate(Nibbles::default(), &mut prefix_set); + let rlp_node = self.rlp_node_allocate(&mut prefix_set); if let Some(root_hash) = rlp_node.as_hash() { root_hash } else { @@ -583,20 +583,38 @@ impl
RevealedSparseTrie
{
/// Update hashes of the nodes that are located at a level deeper than or equal to the provided
/// depth. Root node has a level of 0.
pub fn update_rlp_node_level(&mut self, depth: usize) {
- let mut prefix_set = self.prefix_set.clone().freeze();
+ // Take the current prefix set
+ let mut prefix_set = std::mem::take(&mut self.prefix_set).freeze();
let mut buffers = RlpNodeBuffers::default();
- let targets = self.get_changed_nodes_at_depth(&mut prefix_set, depth);
- for target in targets {
- buffers.path_stack.push((target, Some(true)));
+ // Get the nodes that have changed at the given depth.
+ let (targets, new_prefix_set) = self.get_changed_nodes_at_depth(&mut prefix_set, depth);
+ // Update the prefix set to the prefix set of the nodes that stil need to be updated.
+ self.prefix_set = new_prefix_set;
+
+ trace!(target: "trie::sparse", ?depth, ?targets, "Updating nodes at depth");
+ for (level, path) in targets {
+ buffers.path_stack.push(RlpNodePathStackItem {
+ level,
+ path,
+ is_in_prefix_set: Some(true),
+ });
self.rlp_node(&mut prefix_set, &mut buffers);
}
}
- /// Returns a list of paths to the nodes that were changed according to the prefix set and are
- /// located at the provided depth when counting from the root node. If there's a leaf at a
- /// depth less than the provided depth, it will be included in the result.
- fn get_changed_nodes_at_depth(&self, prefix_set: &mut PrefixSet, depth: usize) -> Vec RevealedSparseTrie {
continue
}
- targets.push(path);
+ targets.push((level, path));
}
SparseNode::Extension { key, hash, store_in_db_trie: _ } => {
if hash.is_some() && !prefix_set.contains(&path) {
@@ -616,8 +634,10 @@ impl RevealedSparseTrie {
}
if level >= depth {
- targets.push(path);
+ targets.push((level, path));
} else {
+ unchanged_prefix_set.insert(path.clone());
+
path.extend_from_slice_unchecked(key);
paths.push((path, level + 1));
}
@@ -628,8 +648,10 @@ impl RevealedSparseTrie {
}
if level >= depth {
- targets.push(path);
+ targets.push((level, path));
} else {
+ unchanged_prefix_set.insert(path.clone());
+
for bit in CHILD_INDEX_RANGE.rev() {
if state_mask.is_bit_set(bit) {
let mut child_path = path.clone();
@@ -642,16 +664,16 @@ impl RevealedSparseTrie {
}
}
- targets
+ (targets, unchanged_prefix_set)
}
- /// Look up or calculate the RLP of the node at the given path.
+ /// Look up or calculate the RLP of the node at the root path.
///
/// # Panics
///
/// If the node at provided path does not exist.
- pub fn rlp_node_allocate(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode {
- let mut buffers = RlpNodeBuffers::new_with_path(path);
+ pub fn rlp_node_allocate(&mut self, prefix_set: &mut PrefixSet) -> RlpNode {
+ let mut buffers = RlpNodeBuffers::new_with_root_path();
self.rlp_node(prefix_set, &mut buffers)
}
@@ -665,14 +687,29 @@ impl RevealedSparseTrie {
prefix_set: &mut PrefixSet,
buffers: &mut RlpNodeBuffers,
) -> RlpNode {
- 'main: while let Some((path, mut is_in_prefix_set)) = buffers.path_stack.pop() {
+ let starting_path = buffers.path_stack.last().map(|item| item.path.clone());
+
+ 'main: while let Some(RlpNodePathStackItem { level, path, mut is_in_prefix_set }) =
+ buffers.path_stack.pop()
+ {
+ let node = self.nodes.get_mut(&path).unwrap();
+ trace!(
+ target: "trie::sparse",
+ ?starting_path,
+ ?level,
+ ?path,
+ ?is_in_prefix_set,
+ ?node,
+ "Popped node from path stack"
+ );
+
// Check if the path is in the prefix set.
// First, check the cached value. If it's `None`, then check the prefix set, and update
// the cached value.
let mut prefix_set_contains =
|path: &Nibbles| *is_in_prefix_set.get_or_insert_with(|| prefix_set.contains(path));
- let (rlp_node, node_type) = match self.nodes.get_mut(&path).unwrap() {
+ let (rlp_node, node_type) = match node {
SparseNode::Empty => (RlpNode::word_rlp(&EMPTY_ROOT_HASH), SparseNodeType::Empty),
SparseNode::Hash(hash) => (RlpNode::word_rlp(hash), SparseNodeType::Hash),
SparseNode::Leaf { key, hash } => {
@@ -698,8 +735,12 @@ impl RevealedSparseTrie {
RlpNode::word_rlp(&hash),
SparseNodeType::Extension { store_in_db_trie: Some(store_in_db_trie) },
)
- } else if buffers.rlp_node_stack.last().is_some_and(|e| e.0 == child_path) {
- let (_, child, child_node_type) = buffers.rlp_node_stack.pop().unwrap();
+ } else if buffers.rlp_node_stack.last().is_some_and(|e| e.path == child_path) {
+ let RlpNodeStackItem {
+ path: _,
+ rlp_node: child,
+ node_type: child_node_type,
+ } = buffers.rlp_node_stack.pop().unwrap();
self.rlp_buf.clear();
let rlp_node = ExtensionNodeRef::new(key, &child).rlp(&mut self.rlp_buf);
*hash = rlp_node.as_hash();
@@ -726,7 +767,14 @@ impl RevealedSparseTrie {
)
} else {
// need to get rlp node for child first
- buffers.path_stack.extend([(path, is_in_prefix_set), (child_path, None)]);
+ buffers.path_stack.extend([
+ RlpNodePathStackItem { level, path, is_in_prefix_set },
+ RlpNodePathStackItem {
+ level: level + 1,
+ path: child_path,
+ is_in_prefix_set: None,
+ },
+ ]);
continue
}
}
@@ -734,11 +782,13 @@ impl RevealedSparseTrie {
if let Some((hash, store_in_db_trie)) =
hash.zip(*store_in_db_trie).filter(|_| !prefix_set_contains(&path))
{
- buffers.rlp_node_stack.push((
+ buffers.rlp_node_stack.push(RlpNodeStackItem {
path,
- RlpNode::word_rlp(&hash),
- SparseNodeType::Branch { store_in_db_trie: Some(store_in_db_trie) },
- ));
+ rlp_node: RlpNode::word_rlp(&hash),
+ node_type: SparseNodeType::Branch {
+ store_in_db_trie: Some(store_in_db_trie),
+ },
+ });
continue
}
let retain_updates = self.updates.is_some() && prefix_set_contains(&path);
@@ -763,8 +813,12 @@ impl RevealedSparseTrie {
let mut hash_mask = TrieMask::default();
let mut hashes = Vec::new();
for (i, child_path) in buffers.branch_child_buf.iter().enumerate() {
- if buffers.rlp_node_stack.last().is_some_and(|e| &e.0 == child_path) {
- let (_, child, child_node_type) = buffers.rlp_node_stack.pop().unwrap();
+ if buffers.rlp_node_stack.last().is_some_and(|e| &e.path == child_path) {
+ let RlpNodeStackItem {
+ path: _,
+ rlp_node: child,
+ node_type: child_node_type,
+ } = buffers.rlp_node_stack.pop().unwrap();
// Update the masks only if we need to retain trie updates
if retain_updates {
@@ -815,10 +869,18 @@ impl RevealedSparseTrie {
added_children = true;
} else {
debug_assert!(!added_children);
- buffers.path_stack.push((path, is_in_prefix_set));
- buffers
- .path_stack
- .extend(buffers.branch_child_buf.drain(..).map(|p| (p, None)));
+ buffers.path_stack.push(RlpNodePathStackItem {
+ level,
+ path,
+ is_in_prefix_set,
+ });
+ buffers.path_stack.extend(buffers.branch_child_buf.drain(..).map(
+ |path| RlpNodePathStackItem {
+ level: level + 1,
+ path,
+ is_in_prefix_set: None,
+ },
+ ));
continue 'main
}
}
@@ -893,11 +955,23 @@ impl RevealedSparseTrie {
)
}
};
- buffers.rlp_node_stack.push((path, rlp_node, node_type));
+
+ trace!(
+ target: "trie::sparse",
+ ?starting_path,
+ ?level,
+ ?path,
+ ?node,
+ ?node_type,
+ ?is_in_prefix_set,
+ "Added node to rlp node stack"
+ );
+
+ buffers.rlp_node_stack.push(RlpNodeStackItem { path, rlp_node, node_type });
}
debug_assert_eq!(buffers.rlp_node_stack.len(), 1);
- buffers.rlp_node_stack.pop().unwrap().1
+ buffers.rlp_node_stack.pop().unwrap().rlp_node
}
}
@@ -1361,10 +1435,10 @@ struct RemovedSparseNode {
/// Collection of reusable buffers for [`RevealedSparseTrie::rlp_node`].
#[derive(Debug, Default)]
pub struct RlpNodeBuffers {
- /// Stack of paths we need rlp nodes for and whether the path is in the prefix set.
- path_stack: Vec<(Nibbles, Option