jleony/nanoreth
1
0
Fork 0
mirror of https://github.com/hl-archive-node/nanoreth.git synced 2025-12-06 10:59:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(trie): update prefix set on the call to RevealedSparseTrie::update_rlp_node_level (#14108)

Browse Source
This commit is contained in:
Alexey Shekhirin
2025-01-30 19:32:09 +00:00
committed by GitHub
parent 2c56061b90
commit 590b58f978
1 changed files with 185 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

247
crates/trie/sparse/src/trie.rs
View File

@ -570,9 +570,9 @@ impl<P> RevealedSparseTrie<P> {
/// 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<P> RevealedSparseTrie<P> {
/// 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<Nibbles> {
/// Returns a list of levels and 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.
///
/// Additionally, returns a new prefix set containing the paths that will not be updated, thus
/// need re-calculation.
fn get_changed_nodes_at_depth(
&self,
prefix_set: &mut PrefixSet,
depth: usize,
) -> (Vec<(usize, Nibbles)>, PrefixSetMut) {
let mut unchanged_prefix_set = PrefixSetMut::default();
let mut paths = Vec::from([(Nibbles::default(), 0)]);
let mut targets = Vec::new();
@ -608,7 +626,7 @@ impl<P> RevealedSparseTrie<P> {
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<P> RevealedSparseTrie<P> {
}
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<P> RevealedSparseTrie<P> {
}
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<P> RevealedSparseTrie<P> {
}
}
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<P> RevealedSparseTrie<P> {
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<P> RevealedSparseTrie<P> {
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<P> RevealedSparseTrie<P> {
)
} 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<P> RevealedSparseTrie<P> {
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<P> RevealedSparseTrie<P> {
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<P> RevealedSparseTrie<P> {
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<P> RevealedSparseTrie<P> {
)
}
};
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<bool>)>,
/// Stack of rlp nodes
rlp_node_stack: Vec<(Nibbles, RlpNode, SparseNodeType)>,
/// Stack of RLP node paths
path_stack: Vec<RlpNodePathStackItem>,
/// Stack of RLP nodes
rlp_node_stack: Vec<RlpNodeStackItem>,
/// Reusable branch child path
branch_child_buf: SmallVec<[Nibbles; 16]>,
/// Reusable branch value stack
@ -1372,10 +1446,14 @@ pub struct RlpNodeBuffers {
}
impl RlpNodeBuffers {
/// Creates a new instance of buffers with the given path on the stack.
fn new_with_path(path: Nibbles) -> Self {
/// Creates a new instance of buffers with the root path on the stack.
fn new_with_root_path() -> Self {
Self {
path_stack: vec![(path, None)],
path_stack: vec![RlpNodePathStackItem {
level: 0,
path: Nibbles::default(),
is_in_prefix_set: None,
}],
rlp_node_stack: Vec::new(),
branch_child_buf: SmallVec::<[Nibbles; 16]>::new_const(),
branch_value_stack_buf: SmallVec::<[RlpNode; 16]>::new_const(),
@ -1383,6 +1461,28 @@ impl RlpNodeBuffers {
}
}
/// RLP node path stack item.
#[derive(Debug)]
struct RlpNodePathStackItem {
/// Level at which the node is located. Higher numbers correspond to lower levels in the trie.
level: usize,
/// Path to the node.
path: Nibbles,
/// Whether the path is in the prefix set. If [`None`], then unknown yet.
is_in_prefix_set: Option<bool>,
}
/// RLP node stack item.
#[derive(Debug)]
struct RlpNodeStackItem {
/// Path to the node.
path: Nibbles,
/// RLP node.
rlp_node: RlpNode,
/// Type of the node.
node_type: SparseNodeType,
}
/// The aggregation of sparse trie updates.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SparseTrieUpdates {
@ -2498,39 +2598,62 @@ mod tests {
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 0),
vec![Nibbles::default()]
(vec![(0, Nibbles::default())], PrefixSetMut::default())
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 1),
vec![Nibbles::from_nibbles_unchecked([0x5])]
(vec![(1, Nibbles::from_nibbles_unchecked([0x5]))], [Nibbles::default()].into())
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 2),
vec![
Nibbles::from_nibbles_unchecked([0x5, 0x0]),
Nibbles::from_nibbles_unchecked([0x5, 0x2]),
Nibbles::from_nibbles_unchecked([0x5, 0x3])
]
(
vec![
(2, Nibbles::from_nibbles_unchecked([0x5, 0x0])),
(2, Nibbles::from_nibbles_unchecked([0x5, 0x2])),
(2, Nibbles::from_nibbles_unchecked([0x5, 0x3]))
],
[Nibbles::default(), Nibbles::from_nibbles_unchecked([0x5])].into()
)
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 3),
vec![
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3]),
Nibbles::from_nibbles_unchecked([0x5, 0x2]),
Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1]),
Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3])
]
(
vec![
(3, Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3])),
(2, Nibbles::from_nibbles_unchecked([0x5, 0x2])),
(3, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1])),
(3, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3]))
],
[
Nibbles::default(),
Nibbles::from_nibbles_unchecked([0x5]),
Nibbles::from_nibbles_unchecked([0x5, 0x0]),
Nibbles::from_nibbles_unchecked([0x5, 0x3])
]
.into()
)
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 4),
vec![
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x1]),
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x3]),
Nibbles::from_nibbles_unchecked([0x5, 0x2]),
Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1]),
Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x0]),
Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x2])
]
(
vec![
(4, Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x1])),
(4, Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x3])),
(2, Nibbles::from_nibbles_unchecked([0x5, 0x2])),
(3, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1])),
(4, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x0])),
(4, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x2]))
],
[
Nibbles::default(),
Nibbles::from_nibbles_unchecked([0x5]),
Nibbles::from_nibbles_unchecked([0x5, 0x0]),
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3]),
Nibbles::from_nibbles_unchecked([0x5, 0x3]),
Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3])
]
.into()
)
);
}