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): update prefix set on the call to RevealedSparseTrie::update_rlp_node_level #14108

Merged
merged 13 commits into from
Jan 30, 2025
Merged

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

Changes from 11 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
a090a37
feat(trie): log each popped node path in sparse trie
shekhirin Jan 30, 2025
0b34a59
add node type
shekhirin Jan 30, 2025
606b154
better logs
shekhirin Jan 30, 2025
bb6b206
log node
shekhirin Jan 30, 2025
1259d77
log nodes at depth update
shekhirin Jan 30, 2025
929f88a
take prefix set when updating at level
shekhirin Jan 30, 2025
ef3342d
generate new prefix set
shekhirin Jan 30, 2025
6fad18f
fix test and logic
shekhirin Jan 30, 2025
37055bb
rename var and add comment
shekhirin Jan 30, 2025
b15f4dc
fix test
shekhirin Jan 30, 2025
2278214
better comments
shekhirin Jan 30, 2025
c7a78d8
use structs instead of tuples
shekhirin Jan 30, 2025
95dc36b
set correct levels
shekhirin Jan 30, 2025
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
144 changes: 107 additions & 37 deletions crates/trie/sparse/src/trie.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -570,7 +570,7 @@ 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);
if let Some(root_hash) = rlp_node.as_hash() {
Expand All @@ -583,20 +583,34 @@ 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);
// 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 target in targets {
buffers.path_stack.push((target, Some(true)));
buffers.path_stack.push((0, target, 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> {
///
/// Additionally, returns a new prefix set containing the paths that will not be updated, thus
/// need re-calculation.
shekhirin marked this conversation as resolved.
Show resolved Hide resolved
fn get_changed_nodes_at_depth(
&self,
prefix_set: &mut PrefixSet,
depth: usize,
) -> (Vec<Nibbles>, PrefixSetMut) {
let mut unchanged_prefix_set = PrefixSetMut::default();
let mut paths = Vec::from([(Nibbles::default(), 0)]);
let mut targets = Vec::new();

Expand All @@ -618,6 +632,8 @@ impl<P> RevealedSparseTrie<P> {
if level >= depth {
targets.push(path);
} else {
unchanged_prefix_set.insert(path.clone());

path.extend_from_slice_unchecked(key);
paths.push((path, level + 1));
}
Expand All @@ -630,6 +646,8 @@ impl<P> RevealedSparseTrie<P> {
if level >= depth {
targets.push(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();
Expand All @@ -642,7 +660,7 @@ impl<P> RevealedSparseTrie<P> {
}
}

targets
(targets, unchanged_prefix_set)
}

/// Look up or calculate the RLP of the node at the given path.
Expand All @@ -665,14 +683,27 @@ 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(|(_, path, _)| path).cloned();

'main: while let Some((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 } => {
Expand Down Expand Up @@ -726,7 +757,10 @@ 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([
(level, path, is_in_prefix_set),
(level + 1, child_path, None),
]);
continue
}
}
Expand Down Expand Up @@ -815,10 +849,10 @@ 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((level, path, is_in_prefix_set));
buffers.path_stack.extend(
buffers.branch_child_buf.drain(..).map(|p| (level + 1, p, None)),
);
continue 'main
}
}
Expand Down Expand Up @@ -893,6 +927,18 @@ impl<P> RevealedSparseTrie<P> {
)
}
};

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((path, rlp_node, node_type));
}

Expand Down Expand Up @@ -1361,8 +1407,9 @@ 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 levels of depth starting from the node where the calculation began, paths we need
/// rlp nodes for, whether the path is in the prefix set.
path_stack: Vec<(usize, Nibbles, Option<bool>)>,
shekhirin marked this conversation as resolved.
Show resolved Hide resolved
/// Stack of rlp nodes
rlp_node_stack: Vec<(Nibbles, RlpNode, SparseNodeType)>,
/// Reusable branch child path
Expand All @@ -1375,7 +1422,7 @@ impl RlpNodeBuffers {
/// Creates a new instance of buffers with the given path on the stack.
fn new_with_path(path: Nibbles) -> Self {
Self {
path_stack: vec![(path, None)],
path_stack: vec![(0, path, None)],
rlp_node_stack: Vec::new(),
branch_child_buf: SmallVec::<[Nibbles; 16]>::new_const(),
branch_value_stack_buf: SmallVec::<[RlpNode; 16]>::new_const(),
Expand Down Expand Up @@ -2498,39 +2545,62 @@ mod tests {

assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 0),
vec![Nibbles::default()]
(vec![Nibbles::default()], PrefixSetMut::default())
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 1),
vec![Nibbles::from_nibbles_unchecked([0x5])]
(vec![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![
Nibbles::from_nibbles_unchecked([0x5, 0x0]),
Nibbles::from_nibbles_unchecked([0x5, 0x2]),
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![
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])
],
[
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![
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])
],
[
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()
)
);
}

Expand Down
Loading