chore(trie): address arena PR review feedback (#22996)

Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Brian Picciano
2026-03-12 17:04:02 +01:00
committed by GitHub
parent 6cd0f843a8
commit 9b53c4fa39
4 changed files with 56 additions and 60 deletions

View File

@@ -0,0 +1,5 @@
---
reth-trie-sparse: patch
---
Refactored arena trie internals by adding a `BranchChildIdx::sibling()` helper, deduplicating `Index`/`NodeArena` type aliases, and replacing `is_empty()` with a `drop_root()` method. Fixed a bug where `cursor.pop()` was called before checking if the leaf was the root node, which could cause incorrect dirty-state propagation.

View File

@@ -1,19 +1,13 @@
use super::{
branch_child_idx::{BranchChildIdx, BranchChildIter},
ArenaSparseNode, ArenaSparseNodeBranchChild, ArenaSparseNodeState,
ArenaSparseNode, ArenaSparseNodeBranchChild, ArenaSparseNodeState, Index, NodeArena,
};
use alloc::vec::Vec;
use reth_trie_common::Nibbles;
use slotmap::{DefaultKey, SlotMap};
use tracing::{instrument, trace};
const TRACE_TARGET: &str = "trie::arena::cursor";
/// Alias for the slotmap key type used as node references throughout the arena trie.
type Index = DefaultKey;
/// Alias for the slotmap used as the node arena throughout the arena trie.
type NodeArena = SlotMap<Index, ArenaSparseNode>;
/// An entry on the cursor's traversal stack, tracking an ancestor node during trie walks.
#[derive(Debug, Clone)]
pub(super) struct ArenaCursorStackEntry {
@@ -93,17 +87,14 @@ impl ArenaCursor {
self.stack.len() - 1
}
/// Returns `true` if the stack is empty.
pub(super) const fn is_empty(&self) -> bool {
self.stack.is_empty()
}
/// Clears the traversal stack and pushes the given root entry.
/// Replaces the root entry on the stack with a new one.
///
/// The stack must contain exactly the root (depth 0) or be empty (freshly constructed).
#[instrument(level = "trace", target = TRACE_TARGET, skip(self, arena))]
pub(super) fn reset(&mut self, arena: &NodeArena, idx: Index, path: Nibbles) {
debug_assert!(
self.stack.is_empty() && !self.needs_pop,
"cursor must be fully drained before reset; stack has {} entries, needs_pop={}",
self.stack.len() <= 1 && !self.needs_pop,
"cursor must be drained before reset; stack has {} entries, needs_pop={}",
self.stack.len(),
self.needs_pop,
);
@@ -150,18 +141,15 @@ impl ArenaCursor {
_ => false,
});
if child_is_dirty {
let parent_state = arena[parent.index].state_mut();
if !matches!(parent_state, ArenaSparseNodeState::Dirty) {
*parent_state = ArenaSparseNodeState::Dirty;
}
*arena[parent.index].state_mut() = ArenaSparseNodeState::Dirty;
}
}
entry
}
/// Drains the stack, propagating dirty state from each entry to its parent,
/// then removes the final (root) entry.
/// Drains the stack down to the root, propagating dirty state from each popped entry
/// to its parent. The root entry remains on the stack (there is no parent to propagate to).
#[instrument(level = "trace", target = TRACE_TARGET, skip_all)]
pub(super) fn drain(&mut self, arena: &mut NodeArena) {
trace!(target: TRACE_TARGET, "Draining stack");
@@ -169,7 +157,6 @@ impl ArenaCursor {
while self.stack.len() > 1 {
self.pop(arena);
}
self.stack.clear();
}
/// Returns the logical path of the branch at the top of the stack.

View File

@@ -1623,29 +1623,25 @@ impl ArenaParallelSparseTrie {
let child_nibble = head_path.last().expect("non-root leaf");
let parent_branch = arena[parent_idx].branch_ref();
if parent_branch.state_mask.count_bits() == 2 {
let child_idx = BranchChildIdx::new(parent_branch.state_mask, child_nibble)
.expect("leaf nibble not found in parent state_mask");
// With exactly 2 bits set the dense array has indices 0 and 1.
let sibling_dense_idx = 1 - child_idx.get();
if parent_branch.children[sibling_dense_idx].is_blinded() {
let sibling_nibble = parent_branch
.state_mask
.iter()
.find(|&n| n != child_nibble)
.expect("branch has two children");
let mut sibling_path = cursor.parent_logical_branch_path(arena);
sibling_path.push_unchecked(sibling_nibble);
trace!(target: TRACE_TARGET, ?full_path, ?sibling_path, "Removal would collapse branch onto blinded sibling, requesting proof");
return (
RemoveLeafResult::NeedsProof {
key,
proof_key: Self::nibbles_to_padded_b256(&sibling_path),
min_len: (sibling_path.len() as u8).min(64),
},
SubtrieCounterDeltas::default(),
);
}
if parent_branch.state_mask.count_bits() == 2 &&
parent_branch.sibling_child(child_nibble).is_blinded()
{
let sibling_nibble = parent_branch
.state_mask
.iter()
.find(|&n| n != child_nibble)
.expect("branch has two children");
let mut sibling_path = cursor.parent_logical_branch_path(arena);
sibling_path.push_unchecked(sibling_nibble);
trace!(target: TRACE_TARGET, ?full_path, ?sibling_path, "Removal would collapse branch onto blinded sibling, requesting proof");
return (
RemoveLeafResult::NeedsProof {
key,
proof_key: Self::nibbles_to_padded_b256(&sibling_path),
min_len: (sibling_path.len() as u8).min(64),
},
SubtrieCounterDeltas::default(),
);
}
}
@@ -1653,10 +1649,7 @@ impl ArenaParallelSparseTrie {
let removed_was_dirty =
matches!(arena[head_idx].state_ref(), Some(ArenaSparseNodeState::Dirty));
// Pop the leaf entry, propagating dirty state to the parent.
cursor.pop(arena);
if cursor.is_empty() {
if cursor.depth() == 0 {
// The leaf is the root — replace with EmptyRoot and reset the cursor
// so subsequent iterations can call seek normally.
arena.remove(head_idx);
@@ -1671,6 +1664,9 @@ impl ArenaParallelSparseTrie {
);
}
// Pop the leaf entry, propagating dirty state to the parent.
cursor.pop(arena);
// The parent must be a branch. Remove the leaf from it.
let parent_entry = cursor.head().expect("cursor is non-empty");
let parent_idx = parent_entry.index;
@@ -1738,10 +1734,7 @@ impl ArenaParallelSparseTrie {
return None;
}
let child_idx = BranchChildIdx::new(parent_branch.state_mask, child_nibble)
.expect("child nibble not in parent state_mask");
let sibling_dense_idx = 1 - child_idx.get();
if !parent_branch.children[sibling_dense_idx].is_blinded() {
if !parent_branch.sibling_child(child_nibble).is_blinded() {
return None;
}

View File

@@ -1,19 +1,13 @@
use super::{
branch_child_idx::{BranchChildIdx, BranchChildIter},
ArenaSparseSubtrie,
ArenaSparseSubtrie, Index, NodeArena,
};
use alloc::{boxed::Box, vec::Vec};
use alloy_primitives::{keccak256, B256};
use alloy_trie::{BranchNodeCompact, TrieMask};
use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNodeV2, RlpNode, TrieNodeV2};
use slotmap::{DefaultKey, SlotMap};
use smallvec::SmallVec;
/// Alias for the slotmap key type used as node references throughout the arena trie.
type Index = DefaultKey;
/// Alias for the slotmap used as the node arena throughout the arena trie.
type NodeArena = SlotMap<Index, ArenaSparseNode>;
/// Tracks whether a node's RLP encoding is cached or needs recomputation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum ArenaSparseNodeState {
@@ -106,6 +100,23 @@ impl ArenaSparseNodeBranch {
self.state = ArenaSparseNodeState::Dirty;
}
/// Returns a reference to the sibling child in a branch with exactly 2 children.
///
/// # Panics
///
/// Panics (debug) if the branch does not have exactly 2 children, or if `nibble` is not set.
pub(super) fn sibling_child(&self, nibble: u8) -> &ArenaSparseNodeBranchChild {
debug_assert_eq!(
self.state_mask.count_bits(),
2,
"sibling_child requires exactly 2 children"
);
let child_idx =
BranchChildIdx::new(self.state_mask, nibble).expect("nibble not found in state_mask");
// With exactly 2 children the dense array has indices 0 and 1.
&self.children[1 - child_idx.get()]
}
/// Iterates over `(nibble, &ArenaSparseNodeBranchChild)` pairs in nibble order.
pub(super) fn child_iter(
&self,