diff --git a/crates/trie/sparse/src/arena/mod.rs b/crates/trie/sparse/src/arena/mod.rs index 8a80dbfcb2..ae49e162b4 100644 --- a/crates/trie/sparse/src/arena/mod.rs +++ b/crates/trie/sparse/src/arena/mod.rs @@ -26,6 +26,9 @@ use slotmap::{DefaultKey, Key as _, SlotMap}; use smallvec::SmallVec; use tracing::{instrument, trace}; +#[cfg(feature = "trie-debug")] +use crate::debug_recorder::{LeafUpdateRecord, ProofTrieNodeRecord, RecordedOp, TrieDebugRecorder}; + /// 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. @@ -512,6 +515,9 @@ pub struct ArenaParallelSparseTrie { cleared_subtries: Vec>, /// Thresholds controlling when parallelism is enabled for different operations. parallelism_thresholds: ArenaParallelismThresholds, + /// Debug recorder for tracking mutating operations. + #[cfg(feature = "trie-debug")] + debug_recorder: TrieDebugRecorder, } impl ArenaParallelSparseTrie { @@ -524,6 +530,162 @@ impl ArenaParallelSparseTrie { self } + /// Returns the arena indexes of all [`ArenaSparseNode::Subtrie`] nodes in the upper arena. + fn all_subtries(&self) -> SmallVec<[Index; 16]> { + self.upper_arena + .iter() + .filter_map(|(idx, node)| matches!(node, ArenaSparseNode::Subtrie(_)).then_some(idx)) + .collect() + } + + /// Resets the debug recorder and records the current trie state as `SetRoot` + `RevealNodes` + /// ops, representing the initial state at the beginning of a block (after pruning). + /// + /// Walks the upper arena and all subtries depth-first using the cursor, converting each + /// node into a [`crate::debug_recorder::ProofTrieNodeRecord`]. + #[cfg(feature = "trie-debug")] + fn record_initial_state(&mut self) { + use crate::debug_recorder::{NodeStateRecord, TrieNodeRecord}; + use alloy_primitives::hex; + use alloy_trie::nodes::{BranchNode, TrieNode}; + + fn state_to_record(state: &ArenaSparseNodeState) -> NodeStateRecord { + match state { + ArenaSparseNodeState::Revealed => NodeStateRecord::Revealed, + ArenaSparseNodeState::Cached { rlp_node } => { + NodeStateRecord::Cached { rlp_node: hex::encode(rlp_node.as_ref()) } + } + ArenaSparseNodeState::Dirty => NodeStateRecord::Dirty, + } + } + + /// Converts an [`ArenaSparseNode`] into a [`ProofTrieNodeRecord`] at the given path. + /// For branch children, resolves revealed children's cached RLP from `arena`. + /// Returns `None` for subtrie/taken-subtrie nodes (handled separately). + fn node_to_record( + arena: &NodeArena, + idx: Index, + path: Nibbles, + ) -> Option { + match &arena[idx] { + ArenaSparseNode::EmptyRoot => Some(ProofTrieNodeRecord { + path, + node: TrieNodeRecord(TrieNode::EmptyRoot), + masks: None, + short_key: None, + state: None, + }), + ArenaSparseNode::Branch(b) => { + let stack = b + .children + .iter() + .map(|child| match child { + ArenaSparseNodeBranchChild::Blinded(rlp) => rlp.clone(), + ArenaSparseNodeBranchChild::Revealed(child_idx) => { + // After pruning / root(), all nodes have cached RLP. + arena[*child_idx] + .state_ref() + .and_then(|s| s.cached_rlp_node()) + .cloned() + .unwrap_or_default() + } + }) + .collect(); + Some(ProofTrieNodeRecord { + path, + node: TrieNodeRecord(TrieNode::Branch(BranchNode::new( + stack, + b.state_mask, + ))), + masks: Some(( + b.branch_masks.hash_mask.get(), + b.branch_masks.tree_mask.get(), + )), + short_key: (!b.short_key.is_empty()).then_some(b.short_key), + state: Some(state_to_record(&b.state)), + }) + } + ArenaSparseNode::Leaf { key, value, state, .. } => Some(ProofTrieNodeRecord { + path, + node: TrieNodeRecord(TrieNode::Leaf(alloy_trie::nodes::LeafNode::new( + *key, + value.clone(), + ))), + masks: None, + short_key: None, + state: Some(state_to_record(state)), + }), + ArenaSparseNode::Subtrie(_) | ArenaSparseNode::TakenSubtrie => None, + } + } + + /// Walks an arena depth-first using `cursor` and collects all nodes as records. + fn collect_records( + arena: &mut NodeArena, + root: Index, + root_path: Nibbles, + cursor: &mut ArenaCursor, + result: &mut Vec, + ) { + cursor.reset(arena, root, root_path); + + // The cursor starts with root on the stack but `next` only yields children. + if let Some(record) = node_to_record(arena, root, root_path) { + result.push(record); + } + + loop { + match cursor.next(arena, |_, node| { + matches!(node, ArenaSparseNode::Branch(_) | ArenaSparseNode::Leaf { .. }) + }) { + NextResult::Done => break, + NextResult::Branch | NextResult::NonBranch => { + let head = cursor.head().expect("cursor is non-empty"); + if let Some(record) = node_to_record(arena, head.index, head.path) { + result.push(record); + } + } + } + } + } + + let mut nodes = Vec::new(); + + // Collect from the upper arena. + collect_records( + &mut self.upper_arena, + self.root, + Nibbles::default(), + &mut self.buffers.cursor, + &mut nodes, + ); + + // Collect from all subtries. + for (_, node) in &mut self.upper_arena { + if let ArenaSparseNode::Subtrie(subtrie) = node { + collect_records( + &mut subtrie.arena, + subtrie.root, + subtrie.path, + &mut self.buffers.cursor, + &mut nodes, + ); + } + } + + // Reset the recorder and record that we pruned, then the initial state. + self.debug_recorder.reset(); + self.debug_recorder.record(RecordedOp::Prune); + + // First node is the root → SetRoot, remaining → RevealNodes. + if let Some(root_record) = nodes.first() { + self.debug_recorder.record(RecordedOp::SetRoot { node: root_record.clone() }); + } + if nodes.len() > 1 { + self.debug_recorder.record(RecordedOp::RevealNodes { nodes: nodes[1..].to_vec() }); + } + } + /// Takes a cleared [`ArenaSparseSubtrie`] from the pool (or creates a new one) and /// pre-allocates a root slot with a placeholder. The caller must overwrite /// `subtrie.arena[subtrie.root]` before use. @@ -1953,6 +2115,8 @@ impl Default for ArenaParallelSparseTrie { buffers: ArenaTrieBuffers::default(), cleared_subtries: Vec::new(), parallelism_thresholds: ArenaParallelismThresholds::default(), + #[cfg(feature = "trie-debug")] + debug_recorder: Default::default(), } } } @@ -1980,6 +2144,15 @@ impl SparseTrie for ArenaParallelSparseTrie { masks: Option, retain_updates: bool, ) -> SparseTrieResult<()> { + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::SetRoot { + node: ProofTrieNodeRecord::from_proof_trie_node_v2(&ProofTrieNodeV2 { + path: Nibbles::default(), + node: root.clone(), + masks, + }), + }); + debug_assert!( matches!(self.upper_arena[self.root], ArenaSparseNode::EmptyRoot), "set_root called on a trie that already has revealed nodes" @@ -2042,6 +2215,11 @@ impl SparseTrie for ArenaParallelSparseTrie { return Ok(()); } + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::RevealNodes { + nodes: nodes.iter().map(ProofTrieNodeRecord::from_proof_trie_node_v2).collect(), + }); + if matches!(self.upper_arena[self.root], ArenaSparseNode::EmptyRoot) { trace!(target: TRACE_TARGET, "Skipping reveal_nodes on empty root"); return Ok(()); @@ -2190,6 +2368,9 @@ impl SparseTrie for ArenaParallelSparseTrie { #[instrument(level = "trace", target = TRACE_TARGET, skip_all, ret)] fn root(&mut self) -> B256 { + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::Root); + self.update_subtrie_hashes(); // Merge buffered subtrie updates into self.updates before hashing the upper trie, @@ -2215,6 +2396,9 @@ impl SparseTrie for ArenaParallelSparseTrie { #[instrument(level = "trace", target = TRACE_TARGET, skip_all)] fn update_subtrie_hashes(&mut self) { + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::UpdateSubtrieHashes); + trace!(target: TRACE_TARGET, "Updating subtrie hashes"); // Only descend if the root is a branch; otherwise there are no subtries. @@ -2342,12 +2526,18 @@ impl SparseTrie for ArenaParallelSparseTrie { #[instrument(level = "trace", target = TRACE_TARGET, skip_all)] fn clear(&mut self) { - for (_, node) in self.upper_arena.drain() { - if let ArenaSparseNode::Subtrie(mut subtrie) = node { + #[cfg(feature = "trie-debug")] + self.debug_recorder.reset(); + + for idx in self.all_subtries() { + if let ArenaSparseNode::Subtrie(mut subtrie) = + self.upper_arena.remove(idx).expect("subtrie exists in arena") + { subtrie.clear(); self.cleared_subtries.push(subtrie); } } + self.upper_arena.clear(); self.root = self.upper_arena.insert(ArenaSparseNode::EmptyRoot); if let Some(updates) = self.updates.as_mut() { updates.clear() @@ -2541,6 +2731,9 @@ impl SparseTrie for ArenaParallelSparseTrie { } } + #[cfg(feature = "trie-debug")] + self.record_initial_state(); + pruned } @@ -2559,6 +2752,12 @@ impl SparseTrie for ArenaParallelSparseTrie { return Ok(()); } + #[cfg(feature = "trie-debug")] + let recorded_updates: Vec<_> = + updates.iter().map(|(k, v)| (*k, LeafUpdateRecord::from(v))).collect(); + #[cfg(feature = "trie-debug")] + let mut recorded_proof_targets: Vec<(B256, u8)> = Vec::new(); + // Drain and sort updates lexicographically by nibbles path. let mut sorted: Vec<_> = updates.drain().map(|(key, update)| (key, Nibbles::unpack(key), update)).collect(); @@ -2585,6 +2784,8 @@ impl SparseTrie for ArenaParallelSparseTrie { let min_len = (logical_len as u8 + 1).min(64); trace!(target: TRACE_TARGET, ?key, min_len, "Update hit blinded node, requesting proof"); proof_required_fn(key, min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((key, min_len)); updates.insert(key, update.clone()); } // Subtrie — forward all consecutive updates under this subtrie's prefix. @@ -2612,6 +2813,8 @@ impl SparseTrie for ArenaParallelSparseTrie { ) { trace!(target: TRACE_TARGET, proof_key = ?proof.key, proof_min_len = proof.min_len, "Subtrie collapse would need blinded sibling, requesting proof"); proof_required_fn(proof.key, proof.min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((proof.key, proof.min_len)); for &(key, _, ref update) in subtrie_updates { updates.insert(key, update.clone()); } @@ -2643,6 +2846,8 @@ impl SparseTrie for ArenaParallelSparseTrie { for (target_idx, proof) in subtrie.required_proofs.drain(..) { proof_required_fn(proof.key, proof.min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((proof.key, proof.min_len)); let (key, _, ref update) = subtrie_updates[target_idx]; updates.insert(key, update.clone()); } @@ -2703,6 +2908,8 @@ impl SparseTrie for ArenaParallelSparseTrie { match result { RemoveLeafResult::NeedsProof { key, proof_key, min_len } => { proof_required_fn(proof_key, min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((proof_key, min_len)); let update = mem::replace(&mut sorted[update_idx].2, LeafUpdate::Touched); updates.insert(key, update); @@ -2738,6 +2945,13 @@ impl SparseTrie for ArenaParallelSparseTrie { if taken.is_empty() { #[cfg(debug_assertions)] self.debug_assert_subtrie_structure(); + + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::UpdateLeaves { + updates: recorded_updates, + proof_targets: recorded_proof_targets, + }); + return Ok(()); } @@ -2760,6 +2974,8 @@ impl SparseTrie for ArenaParallelSparseTrie { let subtrie_updates = &sorted[range]; for (target_idx, proof) in subtrie.required_proofs.drain(..) { proof_required_fn(proof.key, proof.min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((proof.key, proof.min_len)); let (key, _, ref update) = subtrie_updates[target_idx]; updates.insert(key, update.clone()); } @@ -2806,9 +3022,20 @@ impl SparseTrie for ArenaParallelSparseTrie { #[cfg(debug_assertions)] self.debug_assert_subtrie_structure(); + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::UpdateLeaves { + updates: recorded_updates, + proof_targets: recorded_proof_targets, + }); + Ok(()) } + #[cfg(feature = "trie-debug")] + fn take_debug_recorder(&mut self) -> TrieDebugRecorder { + core::mem::take(&mut self.debug_recorder) + } + fn commit_updates( &mut self, _updated: &HashMap, diff --git a/crates/trie/sparse/src/debug_recorder.rs b/crates/trie/sparse/src/debug_recorder.rs index dad5a9d590..412d8b8cf1 100644 --- a/crates/trie/sparse/src/debug_recorder.rs +++ b/crates/trie/sparse/src/debug_recorder.rs @@ -1,7 +1,7 @@ //! Debug recorder for tracking mutating operations on sparse tries. //! //! This module is only available with the `trie-debug` feature and provides -//! infrastructure for recording all mutations to a [`crate::ParallelSparseTrie`] +//! infrastructure for recording all mutations to a [`crate::ArenaParallelSparseTrie`] //! for post-hoc debugging of state root mismatches. use alloc::{string::String, vec::Vec}; @@ -60,9 +60,6 @@ pub enum RecordedOp { UpdateLeaves { /// The leaf updates that were applied. updates: Vec<(B256, LeafUpdateRecord)>, - /// Keys remaining in the updates map after the call (i.e. those that could not be applied - /// due to blinded nodes). - remaining_keys: Vec, /// Proof targets returned via the callback as `(key, min_len)` pairs. proof_targets: Vec<(B256, u8)>, }, @@ -75,6 +72,9 @@ pub enum RecordedOp { /// The root trie node that was set. node: ProofTrieNodeRecord, }, + /// Records a `prune` call. Emitted before the post-prune `SetRoot`/`RevealNodes` + /// so consumers know the following initial state is the result of pruning. + Prune, } /// A serializable record of a proof trie node. @@ -86,6 +86,28 @@ pub struct ProofTrieNodeRecord { pub node: TrieNodeRecord, /// The branch node masks `(hash_mask, tree_mask)` stored as raw `u16` values, if present. pub masks: Option<(u16, u16)>, + /// The short key (extension key) of a branch node. When present, the branch's logical path + /// is `path` extended by this key. The arena trie merges extension nodes into their child + /// branch, so this replaces separate `Extension` node records. + #[serde(skip_serializing_if = "Option::is_none")] + pub short_key: Option, + /// The node's state (`Revealed`, `Cached`, or `Dirty`), if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub state: Option, +} + +/// A serializable record of a node's state, mirroring the arena trie's `ArenaSparseNodeState`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub enum NodeStateRecord { + /// The node has been revealed but its RLP encoding is not cached. + Revealed, + /// The node has a cached RLP encoding that is still valid. + Cached { + /// The cached RLP-encoded node, hex-encoded. + rlp_node: String, + }, + /// The node has been modified and its RLP encoding needs recomputation. + Dirty, } impl ProofTrieNodeRecord { @@ -95,25 +117,32 @@ impl ProofTrieNodeRecord { path: node.path, node: TrieNodeRecord(node.node.clone()), masks: node.masks.map(|masks| (masks.hash_mask.get(), masks.tree_mask.get())), + short_key: None, + state: None, } } /// Creates a record from a [`reth_trie_common::ProofTrieNodeV2`]. pub fn from_proof_trie_node_v2(node: &reth_trie_common::ProofTrieNodeV2) -> Self { use reth_trie_common::TrieNodeV2; - let trie_node = match &node.node { - TrieNodeV2::EmptyRoot => TrieNode::EmptyRoot, - TrieNodeV2::Leaf(leaf) => TrieNode::Leaf(leaf.clone()), - TrieNodeV2::Extension(ext) => TrieNode::Extension(ext.clone()), - TrieNodeV2::Branch(branch) => TrieNode::Branch(alloy_trie::nodes::BranchNode::new( - branch.stack.clone(), - branch.state_mask, - )), + let (trie_node, short_key) = match &node.node { + TrieNodeV2::EmptyRoot => (TrieNode::EmptyRoot, None), + TrieNodeV2::Leaf(leaf) => (TrieNode::Leaf(leaf.clone()), None), + TrieNodeV2::Extension(ext) => (TrieNode::Extension(ext.clone()), None), + TrieNodeV2::Branch(branch) => ( + TrieNode::Branch(alloy_trie::nodes::BranchNode::new( + branch.stack.clone(), + branch.state_mask, + )), + (!branch.key.is_empty()).then_some(branch.key), + ), }; Self { path: node.path, node: TrieNodeRecord(trie_node), masks: node.masks.map(|masks| (masks.hash_mask.get(), masks.tree_mask.get())), + short_key, + state: None, } } } diff --git a/crates/trie/sparse/src/parallel.rs b/crates/trie/sparse/src/parallel.rs index 36cb3675ec..f4a5441b5f 100644 --- a/crates/trie/sparse/src/parallel.rs +++ b/crates/trie/sparse/src/parallel.rs @@ -1338,7 +1338,6 @@ impl SparseTrie for ParallelSparseTrie { #[cfg(feature = "trie-debug")] self.debug_recorder.record(RecordedOp::UpdateLeaves { updates: recorded_updates, - remaining_keys: updates.keys().copied().collect(), proof_targets: recorded_proof_targets, });