From 9b53c4fa397a0a79ca38a8490c13ce0d8fd5cf8d Mon Sep 17 00:00:00 2001 From: Brian Picciano Date: Thu, 12 Mar 2026 17:04:02 +0100 Subject: [PATCH] chore(trie): address arena PR review feedback (#22996) Co-authored-by: Amp --- .changelog/dull-clams-pack.md | 5 +++ crates/trie/sparse/src/arena/cursor.rs | 31 +++++---------- crates/trie/sparse/src/arena/mod.rs | 55 +++++++++++--------------- crates/trie/sparse/src/arena/nodes.rs | 25 ++++++++---- 4 files changed, 56 insertions(+), 60 deletions(-) create mode 100644 .changelog/dull-clams-pack.md diff --git a/.changelog/dull-clams-pack.md b/.changelog/dull-clams-pack.md new file mode 100644 index 0000000000..fd521a9539 --- /dev/null +++ b/.changelog/dull-clams-pack.md @@ -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. diff --git a/crates/trie/sparse/src/arena/cursor.rs b/crates/trie/sparse/src/arena/cursor.rs index e11911f6b1..fbbd79e022 100644 --- a/crates/trie/sparse/src/arena/cursor.rs +++ b/crates/trie/sparse/src/arena/cursor.rs @@ -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; - /// 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. diff --git a/crates/trie/sparse/src/arena/mod.rs b/crates/trie/sparse/src/arena/mod.rs index ae49e162b4..10f8040d1e 100644 --- a/crates/trie/sparse/src/arena/mod.rs +++ b/crates/trie/sparse/src/arena/mod.rs @@ -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; } diff --git a/crates/trie/sparse/src/arena/nodes.rs b/crates/trie/sparse/src/arena/nodes.rs index cbc43231bb..8127ced6c0 100644 --- a/crates/trie/sparse/src/arena/nodes.rs +++ b/crates/trie/sparse/src/arena/nodes.rs @@ -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; - /// 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,