Skip to content

Commit

Permalink
perf(trie): reduce allocations in sparse trie rlp node calculation (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
shekhirin authored Oct 28, 2024
1 parent 380e237 commit af5ae5a
Showing 1 changed file with 57 additions and 29 deletions.
86 changes: 57 additions & 29 deletions crates/trie/sparse/src/trie.rs
Original file line number Diff line number Diff line change
Expand Up @@ -558,7 +558,7 @@ impl RevealedSparseTrie {
pub fn root(&mut self) -> B256 {
// take the current prefix set.
let mut prefix_set = std::mem::take(&mut self.prefix_set).freeze();
let root_rlp = self.rlp_node(Nibbles::default(), &mut prefix_set);
let root_rlp = self.rlp_node_allocate(Nibbles::default(), &mut prefix_set);
if let Some(root_hash) = root_rlp.as_hash() {
root_hash
} else {
Expand All @@ -570,10 +570,12 @@ impl RevealedSparseTrie {
/// 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();
let mut buffers = RlpNodeBuffers::default();

let targets = self.get_changed_nodes_at_depth(&mut prefix_set, depth);
for target in targets {
self.rlp_node(target, &mut prefix_set);
buffers.path_stack.push((target, Some(true)));
self.rlp_node(&mut prefix_set, &mut buffers);
}
}

Expand Down Expand Up @@ -629,17 +631,13 @@ impl RevealedSparseTrie {
targets
}

fn rlp_node(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode {
// stack of paths we need rlp nodes for
let mut path_stack = Vec::from([(path, None)]);
// stack of rlp nodes
let mut rlp_node_stack = Vec::<(Nibbles, RlpNode)>::new();
// reusable branch child path
let mut branch_child_buf = SmallVec::<[Nibbles; 16]>::new_const();
// reusable branch value stack
let mut branch_value_stack_buf = SmallVec::<[RlpNode; 16]>::new_const();

'main: while let Some((path, mut is_in_prefix_set)) = path_stack.pop() {
fn rlp_node_allocate(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode {
let mut buffers = RlpNodeBuffers::new_with_path(path);
self.rlp_node(prefix_set, &mut buffers)
}

fn rlp_node(&mut self, prefix_set: &mut PrefixSet, buffers: &mut RlpNodeBuffers) -> RlpNode {
'main: while let Some((path, mut is_in_prefix_set)) = buffers.path_stack.pop() {
// 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.
Expand Down Expand Up @@ -667,63 +665,68 @@ impl RevealedSparseTrie {
child_path.extend_from_slice_unchecked(key);
if let Some(hash) = hash.filter(|_| !prefix_set_contains(&path)) {
RlpNode::word_rlp(&hash)
} else if rlp_node_stack.last().map_or(false, |e| e.0 == child_path) {
let (_, child) = rlp_node_stack.pop().unwrap();
} else if buffers.rlp_node_stack.last().map_or(false, |e| e.0 == child_path) {
let (_, child) = 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();
rlp_node
} else {
// need to get rlp node for child first
path_stack.extend([(path, is_in_prefix_set), (child_path, None)]);
buffers.path_stack.extend([(path, is_in_prefix_set), (child_path, None)]);
continue
}
}
SparseNode::Branch { state_mask, hash } => {
if let Some(hash) = hash.filter(|_| !prefix_set_contains(&path)) {
rlp_node_stack.push((path, RlpNode::word_rlp(&hash)));
buffers.rlp_node_stack.push((path, RlpNode::word_rlp(&hash)));
continue
}

branch_child_buf.clear();
buffers.branch_child_buf.clear();
// Walk children in a reverse order from `f` to `0`, so we pop the `0` first
// from the stack.
for bit in CHILD_INDEX_RANGE.rev() {
if state_mask.is_bit_set(bit) {
let mut child = path.clone();
child.push_unchecked(bit);
branch_child_buf.push(child);
buffers.branch_child_buf.push(child);
}
}

branch_value_stack_buf.resize(branch_child_buf.len(), Default::default());
buffers
.branch_value_stack_buf
.resize(buffers.branch_child_buf.len(), Default::default());
let mut added_children = false;
for (i, child_path) in branch_child_buf.iter().enumerate() {
if rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) {
let (_, child) = rlp_node_stack.pop().unwrap();
for (i, child_path) in buffers.branch_child_buf.iter().enumerate() {
if buffers.rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) {
let (_, child) = buffers.rlp_node_stack.pop().unwrap();
// Insert children in the resulting buffer in a normal order, because
// initially we iterated in reverse.
branch_value_stack_buf[branch_child_buf.len() - i - 1] = child;
buffers.branch_value_stack_buf
[buffers.branch_child_buf.len() - i - 1] = child;
added_children = true;
} else {
debug_assert!(!added_children);
path_stack.push((path, is_in_prefix_set));
path_stack.extend(branch_child_buf.drain(..).map(|p| (p, None)));
buffers.path_stack.push((path, is_in_prefix_set));
buffers
.path_stack
.extend(buffers.branch_child_buf.drain(..).map(|p| (p, None)));
continue 'main
}
}

self.rlp_buf.clear();
let rlp_node = BranchNodeRef::new(&branch_value_stack_buf, *state_mask)
let rlp_node = BranchNodeRef::new(&buffers.branch_value_stack_buf, *state_mask)
.rlp(&mut self.rlp_buf);
*hash = rlp_node.as_hash();
rlp_node
}
};
rlp_node_stack.push((path, rlp_node));
buffers.rlp_node_stack.push((path, rlp_node));
}

rlp_node_stack.pop().unwrap().1
buffers.rlp_node_stack.pop().unwrap().1
}
}

Expand Down Expand Up @@ -803,6 +806,31 @@ struct RemovedSparseNode {
unset_branch_nibble: Option<u8>,
}

/// Collection of reusable buffers for [`RevealedSparseTrie::rlp_node`].
#[derive(Debug, Default)]
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)>,
/// Reusable branch child path
branch_child_buf: SmallVec<[Nibbles; 16]>,
/// Reusable branch value stack
branch_value_stack_buf: SmallVec<[RlpNode; 16]>,
}

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)],
rlp_node_stack: Vec::new(),
branch_child_buf: SmallVec::<[Nibbles; 16]>::new_const(),
branch_value_stack_buf: SmallVec::<[RlpNode; 16]>::new_const(),
}
}
}

#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
Expand Down

0 comments on commit af5ae5a

Please sign in to comment.