diff --git a/Cargo.toml b/Cargo.toml index 5319f09..d86f964 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,7 @@ exclude = [ [dependencies] bytes = "1.4.0" -enum-as-inner = "0.5.1" +enum-as-inner = "0.6.0" ethereum-types = "0.14.1" hex = "0.4.3" keccak-hash = "0.10.0" @@ -30,13 +30,17 @@ rlp = "0.5.2" serde = { version = "1.0.160", features = ["derive", "rc"] } [dev-dependencies] -eth_trie = "0.1.0" -pretty_env_logger = "0.4.0" +eth_trie = "0.4.0" +pretty_env_logger = "0.5.0" rand = "0.8.5" rlp-derive = "0.1.0" serde = { version = "1.0.160", features = ["derive"] } serde_json = "1.0.96" +[features] +default = ["trie_debug"] +trie_debug = [] + [lib] doc-scrape-examples = true diff --git a/src/debug_tools/common.rs b/src/debug_tools/common.rs new file mode 100644 index 0000000..b63bba9 --- /dev/null +++ b/src/debug_tools/common.rs @@ -0,0 +1,127 @@ +use std::fmt::{self, Display}; + +use crate::{ + nibbles::{Nibble, Nibbles}, + partial_trie::{Node, PartialTrie}, + utils::TrieNodeType, +}; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(super) enum PathSegment { + Empty, + Hash, + Branch(Nibble), + Extension(Nibbles), + Leaf(Nibbles), +} + +impl Display for PathSegment { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + PathSegment::Empty => write!(f, "Empty"), + PathSegment::Hash => write!(f, "Hash"), + PathSegment::Branch(nib) => write!(f, "Branch({})", nib), + PathSegment::Extension(nibs) => write!(f, "Extension({})", nibs), + PathSegment::Leaf(nibs) => write!(f, "Leaf({})", nibs), + } + } +} + +impl PathSegment { + pub(super) fn node_type(&self) -> TrieNodeType { + match self { + PathSegment::Empty => TrieNodeType::Empty, + PathSegment::Hash => TrieNodeType::Hash, + PathSegment::Branch(_) => TrieNodeType::Branch, + PathSegment::Extension(_) => TrieNodeType::Extension, + PathSegment::Leaf(_) => TrieNodeType::Leaf, + } + } + + pub(super) fn get_key_piece_from_seg_if_present(&self) -> Option { + match self { + PathSegment::Empty | PathSegment::Hash => None, + PathSegment::Branch(nib) => Some(Nibbles::from_nibble(*nib)), + PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => Some(*nibs), + } + } +} + +pub(super) fn get_segment_from_node_and_key_piece( + n: &Node, + k_piece: &Nibbles, +) -> PathSegment { + match TrieNodeType::from(n) { + TrieNodeType::Empty => PathSegment::Empty, + TrieNodeType::Hash => PathSegment::Hash, + TrieNodeType::Branch => PathSegment::Branch(k_piece.get_nibble(0)), + TrieNodeType::Extension => PathSegment::Extension(*k_piece), + TrieNodeType::Leaf => PathSegment::Leaf(*k_piece), + } +} + +/// Get the key piece from the given node if applicable. +/// +/// Note that there is no specific [`Nibble`] associated with a branch like +/// there are [`Nibbles`] with [Extension][`Node::Extension`] and +/// [Leaf][`Node::Leaf`] nodes, and the only way to get the `Nibble` +/// "associated" with a branch is to look at the next `Nibble` in the current +/// key as we traverse down it. +pub(super) fn get_key_piece_from_node_pulling_from_key_for_branches( + n: &Node, + curr_key: &Nibbles, +) -> Nibbles { + match n { + Node::Empty | Node::Hash(_) => Nibbles::default(), + Node::Branch { .. } => curr_key.get_next_nibbles(1), + Node::Extension { nibbles, child: _ } | Node::Leaf { nibbles, value: _ } => *nibbles, + } +} + +/// Get the key piece from the given node if applicable. Note that +/// [branch][`Node::Branch`]s have no [`Nibble`] directly associated with them. +pub(super) fn get_key_piece_from_node(n: &Node) -> Nibbles { + match n { + Node::Empty | Node::Hash(_) | Node::Branch { .. } => Nibbles::default(), + Node::Extension { nibbles, child: _ } | Node::Leaf { nibbles, value: _ } => *nibbles, + } +} + +#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)] +pub struct NodePath(pub(super) Vec); + +impl NodePath { + pub(super) fn dup_and_append(&self, seg: PathSegment) -> Self { + let mut duped_vec = self.0.clone(); + duped_vec.push(seg); + + Self(duped_vec) + } + + pub(super) fn append(&mut self, seg: PathSegment) { + self.0.push(seg); + } + + fn write_elem(f: &mut fmt::Formatter<'_>, seg: &PathSegment) -> fmt::Result { + write!(f, "{}", seg) + } +} + +impl Display for NodePath { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let num_elems = self.0.len(); + + // For everything but the last elem. + for seg in self.0.iter().take(num_elems.saturating_sub(1)) { + Self::write_elem(f, seg)?; + write!(f, " --> ")?; + } + + // Avoid the extra `-->` for the last elem. + if let Some(seg) = self.0.last() { + Self::write_elem(f, seg)?; + } + + Ok(()) + } +} diff --git a/src/debug_tools/diff.rs b/src/debug_tools/diff.rs new file mode 100644 index 0000000..b043d72 --- /dev/null +++ b/src/debug_tools/diff.rs @@ -0,0 +1,493 @@ +//! Diffing tools to compare two tries against each other. Useful when you want +//! to find where the tries diverge from one each other. +//! +//! There are a few considerations when implementing the logic to create a trie +//! diff: +//! - What should be reported when the trie node structures diverge (eg. two +//! different node types proceeding a given common node)? +//! - If there are multiple structural differences, how do we discover the +//! smallest difference to report back? +//! +//! If the node types between the tries (structure) are identical but some +//! values are different, then these types of diffs are easy to detect and +//! report the lowest difference. Structural differences are more challenging +//! and a bit hard to report well. There are two approaches (only one currently +//! is implemented) in how to detect structural differences: +//! - Top-down search +//! - Bottom-up search +//! +//! These two searches are somewhat self-explanatory: +//! - Top-down will find the highest point of a structural divergence and report +//! it. If there are multiple divergences, then only the one that is the +//! highest in the trie will be reported. +//! - Bottom-up (not implemented) is a lot more complex to implement, but will +//! attempt to find the smallest structural trie difference between the trie. +//! If there are multiple differences, then this will likely be what you want +//! to use. + +use std::fmt::{self, Debug}; +use std::{fmt::Display, ops::Deref}; + +use ethereum_types::H256; + +use super::common::{get_key_piece_from_node, get_segment_from_node_and_key_piece, NodePath}; +use crate::{ + nibbles::Nibbles, + partial_trie::{HashedPartialTrie, Node, PartialTrie}, + utils::TrieNodeType, +}; + +#[derive(Debug, Eq, PartialEq)] +pub struct TrieDiff { + pub latest_diff_res: Option, + // TODO: Later add a second pass for finding diffs from the bottom up (`earliest_diff_res`). +} + +impl Display for TrieDiff { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if let Some(diff) = &self.latest_diff_res { + write!(f, "{}", diff)?; + } + + Ok(()) + } +} + +#[derive(Copy, Clone, Debug)] +enum DiffDetectionState { + NodeTypesDiffer = 0, // Also implies that hashes differ. + HashDiffDetected, + NoDiffDetected, +} + +impl DiffDetectionState { + fn pick_most_significant_state(&self, other: &Self) -> Self { + match self.get_int_repr() > other.get_int_repr() { + false => *other, + true => *self, + } + } + + /// The integer representation also indicates the more "significant" state. + fn get_int_repr(self) -> usize { + self as usize + } +} + +/// A point (node) between the two tries where the children differ. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct DiffPoint { + pub depth: usize, + pub path: NodePath, + pub key: Nibbles, + pub a_info: NodeInfo, + pub b_info: NodeInfo, +} + +impl DiffPoint { + fn new( + child_a: &HashedPartialTrie, + child_b: &HashedPartialTrie, + parent_k: Nibbles, + path: NodePath, + ) -> Self { + let a_key = parent_k.merge_nibbles(&get_key_piece_from_node(child_a)); + let b_key = parent_k.merge_nibbles(&get_key_piece_from_node(child_b)); + + DiffPoint { + depth: 0, + path, + key: parent_k, + a_info: NodeInfo::new(child_a, a_key, get_value_from_node(child_a).cloned()), + b_info: NodeInfo::new(child_b, b_key, get_value_from_node(child_b).cloned()), + } + } +} + +// TODO: Redo display method so this is more readable... +impl Display for DiffPoint { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Point Diff {{depth: {}, ", self.depth)?; + write!(f, "Path: ({}), ", self.path)?; + write!(f, "Key: {:x} ", self.key)?; + write!(f, "A info: {} ", self.a_info)?; + write!(f, "B info: {}}}", self.b_info) + } +} + +/// Meta information for a node in a trie. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct NodeInfo { + key: Nibbles, + + /// The direct value associated with the node (only applicable to `Leaf` & + /// `Branch` nodes). + value: Option>, + node_type: TrieNodeType, + hash: H256, +} + +// TODO: Redo display method so this is more readable... +impl Display for NodeInfo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "(key: {:x} ", self.key)?; + + match &self.value { + Some(v) => write!(f, "Value: 0x{}, ", hex::encode(v))?, + None => write!(f, "Value: N/A, ")?, + } + + write!(f, "Node type: {} ", self.node_type)?; + write!(f, "Trie hash: {:x})", self.hash) + } +} + +impl NodeInfo { + fn new(n: &HashedPartialTrie, key: Nibbles, value: Option>) -> Self { + Self { + key, + value, + node_type: n.deref().into(), + hash: n.hash(), + } + } +} + +/// Create a diff between two tries. Will perform both types of diff searches +/// (top-down & bottom-up). +pub fn create_diff_between_tries(a: &HashedPartialTrie, b: &HashedPartialTrie) -> TrieDiff { + TrieDiff { + latest_diff_res: find_latest_diff_point_between_tries(a, b), + } +} + +// Only support `HashedPartialTrie` due to it being significantly faster to +// detect differences because of caching hashes. +fn find_latest_diff_point_between_tries( + a: &HashedPartialTrie, + b: &HashedPartialTrie, +) -> Option { + let state = DepthDiffPerCallState::new(a, b, Nibbles::default(), 0); + let mut longest_state = DepthNodeDiffState::default(); + + find_latest_diff_point_between_tries_rec(&state, &mut longest_state); + + // If there was a node diff, we always want to prioritize displaying this over a + // hash diff. The reasoning behind this is hash diffs can become sort of + // meaningless or misleading if the trie diverges at some point (eg. saying + // there is a hash diff deep in two separate trie structures doesn't make much + // sense). + longest_state + .longest_key_node_diff + .or(longest_state.longest_key_hash_diff) +} + +#[derive(Debug, Default)] +struct DepthNodeDiffState { + longest_key_node_diff: Option, + longest_key_hash_diff: Option, +} + +impl DepthNodeDiffState { + fn try_update_longest_divergence_key_hash(&mut self, state: &DepthDiffPerCallState) { + Self::replace_longest_field_if_our_key_is_larger( + &mut self.longest_key_hash_diff, + &state.curr_key, + state.a, + state.b, + state.curr_path.clone(), + ); + } + + fn try_update_longest_divergence_key_node(&mut self, state: &DepthDiffPerCallState) { + Self::replace_longest_field_if_our_key_is_larger( + &mut self.longest_key_node_diff, + &state.curr_key, + state.a, + state.b, + state.curr_path.clone(), + ); + } + + fn replace_longest_field_if_our_key_is_larger( + field: &mut Option, + parent_k: &Nibbles, + child_a: &HashedPartialTrie, + child_b: &HashedPartialTrie, + path: NodePath, + ) { + if field + .as_ref() + .map_or(true, |d_point| d_point.key.count < parent_k.count) + { + *field = Some(DiffPoint::new(child_a, child_b, *parent_k, path)); + } + } +} + +/// State that is copied per recursive call. +#[derive(Clone, Debug)] +struct DepthDiffPerCallState<'a> { + a: &'a HashedPartialTrie, + b: &'a HashedPartialTrie, + curr_key: Nibbles, + curr_depth: usize, + + // Horribly inefficient, but these are debug tools, so I think we get a pass. + curr_path: NodePath, +} + +impl<'a> DepthDiffPerCallState<'a> { + /// Exists solely to prevent construction of this type from going over + /// multiple lines. + fn new( + a: &'a HashedPartialTrie, + b: &'a HashedPartialTrie, + curr_key: Nibbles, + curr_depth: usize, + ) -> Self { + Self { + a, + b, + curr_key, + curr_depth, + curr_path: NodePath::default(), + } + } + + /// Note: The assumption here is that `a` and `b` are of the same node type + /// and have the key. + fn new_from_parent( + &self, + a: &'a HashedPartialTrie, + b: &'a HashedPartialTrie, + key_piece: &Nibbles, + ) -> Self { + let new_segment = get_segment_from_node_and_key_piece(self.a, key_piece); + let new_path = self.curr_path.dup_and_append(new_segment); + + Self { + a, + b, + curr_key: self.curr_key.merge_nibbles(key_piece), + curr_depth: self.curr_depth + 1, + curr_path: new_path, + } + } +} + +fn find_latest_diff_point_between_tries_rec( + state: &DepthDiffPerCallState, + depth_state: &mut DepthNodeDiffState, +) -> DiffDetectionState { + let a_hash = state.a.hash(); + let b_hash = state.b.hash(); + + // We're going to ignore node type differences if they have the same hash (only + // case I think where this can happen is if one is a hash node?). + if a_hash == b_hash { + return DiffDetectionState::NoDiffDetected; + } + + let a_type: TrieNodeType = state.a.deref().into(); + let b_type: TrieNodeType = state.b.deref().into(); + + let a_key_piece = get_key_piece_from_node(state.a); + let b_key_piece = get_key_piece_from_node(state.b); + + // Note that differences in a node's `value` will be picked up by a hash + // mismatch. + if (a_type, a_key_piece) == (b_type, b_key_piece) { + depth_state.try_update_longest_divergence_key_node(state); + DiffDetectionState::NodeTypesDiffer + } else { + match (&state.a.node, &state.b.node) { + (Node::Empty, Node::Empty) => DiffDetectionState::NoDiffDetected, + (Node::Hash(a_hash), Node::Hash(b_hash)) => { + create_diff_detection_state_based_from_hashes( + a_hash, + b_hash, + &state.new_from_parent(state.a, state.b, &Nibbles::default()), + depth_state, + ) + } + ( + Node::Branch { + children: a_children, + value: _a_value, + }, + Node::Branch { + children: b_children, + value: _b_value, + }, + ) => { + let mut most_significant_diff_found = DiffDetectionState::NoDiffDetected; + + for i in 0..16_usize { + let res = find_latest_diff_point_between_tries_rec( + &state.new_from_parent( + &a_children[i], + &b_children[i], + &Nibbles::from_nibble(i as u8), + ), + depth_state, + ); + most_significant_diff_found = + most_significant_diff_found.pick_most_significant_state(&res); + } + + if matches!( + most_significant_diff_found, + DiffDetectionState::NoDiffDetected + ) { + most_significant_diff_found + } else { + // Also run a hash check if we haven't picked anything up yet. + create_diff_detection_state_based_from_hash_and_gen_hashes(state, depth_state) + } + } + ( + Node::Extension { + nibbles: a_nibs, + child: a_child, + }, + Node::Extension { + nibbles: _b_nibs, + child: b_child, + }, + ) => find_latest_diff_point_between_tries_rec( + &state.new_from_parent(a_child, b_child, a_nibs), + depth_state, + ), + (Node::Leaf { .. }, Node::Leaf { .. }) => { + create_diff_detection_state_based_from_hash_and_gen_hashes(state, depth_state) + } + _ => unreachable!(), + } + } +} + +fn create_diff_detection_state_based_from_hash_and_gen_hashes( + state: &DepthDiffPerCallState, + depth_state: &mut DepthNodeDiffState, +) -> DiffDetectionState { + let a_hash = state.a.hash(); + let b_hash = state.b.hash(); + + create_diff_detection_state_based_from_hashes(&a_hash, &b_hash, state, depth_state) +} + +fn create_diff_detection_state_based_from_hashes( + a_hash: &H256, + b_hash: &H256, + state: &DepthDiffPerCallState, + depth_state: &mut DepthNodeDiffState, +) -> DiffDetectionState { + match a_hash == b_hash { + false => { + depth_state.try_update_longest_divergence_key_hash(state); + DiffDetectionState::HashDiffDetected + } + true => DiffDetectionState::NoDiffDetected, + } +} + +/// If the node type contains a value (without looking at the children), then +/// return it. +fn get_value_from_node(n: &Node) -> Option<&Vec> { + match n { + Node::Empty | Node::Hash(_) | Node::Extension { .. } => None, + Node::Branch { value, .. } | Node::Leaf { nibbles: _, value } => Some(value), + } +} + +#[cfg(test)] +mod tests { + use super::{create_diff_between_tries, DiffPoint, NodeInfo, NodePath}; + use crate::{ + nibbles::Nibbles, + partial_trie::{HashedPartialTrie, PartialTrie}, + utils::TrieNodeType, + }; + + #[test] + fn depth_single_node_hash_diffs_work() { + // TODO: Reduce duplication once we identify common structures across tests... + let mut a = HashedPartialTrie::default(); + a.insert(0x1234, vec![0]); + let a_hash = a.hash(); + + let mut b = a.clone(); + b.insert(0x1234, vec![1]); + let b_hash = b.hash(); + + let diff = create_diff_between_tries(&a, &b); + + let expected_a = NodeInfo { + key: 0x1234.into(), + value: Some(vec![0]), + node_type: TrieNodeType::Leaf, + hash: a_hash, + }; + + let expected_b = NodeInfo { + key: 0x1234.into(), + value: Some(vec![1]), + node_type: TrieNodeType::Leaf, + hash: b_hash, + }; + + let expected = DiffPoint { + depth: 0, + path: NodePath(vec![]), + key: Nibbles::default(), + a_info: expected_a, + b_info: expected_b, + }; + + assert_eq!(diff.latest_diff_res, Some(expected)); + } + + // TODO: Will finish these tests later (low-priority). + #[test] + #[ignore] + fn depth_single_node_node_diffs_work() { + todo!() + } + + #[test] + #[ignore] + fn depth_multi_node_single_node_hash_diffs_work() { + todo!() + } + + #[test] + #[ignore] + fn depth_multi_node_single_node_node_diffs_work() { + todo!() + } + + #[test] + #[ignore] + fn depth_massive_single_node_diff_tests() { + todo!() + } + + #[test] + #[ignore] + fn depth_multi_node_multi_node_hash_diffs_work() { + todo!() + } + + #[test] + #[ignore] + fn depth_multi_node_multi_node_node_diffs_work() { + todo!() + } + + #[test] + #[ignore] + fn depth_massive_multi_node_diff_tests() { + todo!() + } +} diff --git a/src/debug_tools/mod.rs b/src/debug_tools/mod.rs new file mode 100644 index 0000000..c05eb26 --- /dev/null +++ b/src/debug_tools/mod.rs @@ -0,0 +1,6 @@ +//! Additional methods that may be useful when diagnosing tries from this +//! library. + +pub mod common; +pub mod diff; +pub mod query; diff --git a/src/debug_tools/query.rs b/src/debug_tools/query.rs new file mode 100644 index 0000000..e75c2ed --- /dev/null +++ b/src/debug_tools/query.rs @@ -0,0 +1,313 @@ +//! Query tooling to report info on the path taken when searching down a trie +//! with a given key. + +use std::fmt::{self, Display}; + +use ethereum_types::H256; + +use super::common::{ + get_key_piece_from_node_pulling_from_key_for_branches, get_segment_from_node_and_key_piece, + NodePath, PathSegment, +}; +use crate::{ + nibbles::Nibbles, + partial_trie::{Node, PartialTrie, WrappedNode}, +}; + +/// Params controlling how much information is reported in the query output. +/// +/// By default, the node type along with its key piece is printed out per node +/// (eg. "Leaf(0x1234)"). Additional node specific information can be printed +/// out by enabling `include_node_specific_values`. +#[derive(Clone, Debug)] +pub struct DebugQueryParams { + /// Include (if applicable) the piece of the key that is contained by the + /// node (eg. ("0x1234")). + include_key_piece_per_node: bool, + + /// Include the type of node (eg "Branch"). + include_node_type: bool, + + /// Include additional data that is specific to the node type (eg. The mask + /// of a `Branch` or the hash of a `Hash` node). + include_node_specific_values: bool, +} + +impl Default for DebugQueryParams { + fn default() -> Self { + Self { + include_key_piece_per_node: true, + include_node_type: true, + include_node_specific_values: false, + } + } +} + +#[derive(Debug, Default)] +pub struct DebugQueryParamsBuilder { + params: DebugQueryParams, +} + +impl DebugQueryParamsBuilder { + /// Defaults to `true`. + pub fn print_key_pieces(mut self, enabled: bool) -> Self { + self.params.include_key_piece_per_node = enabled; + self + } + + /// Defaults to `true`. + pub fn print_node_type(mut self, enabled: bool) -> Self { + self.params.include_node_type = enabled; + self + } + + /// Defaults to `false`. + pub fn print_node_specific_values(mut self, enabled: bool) -> Self { + self.params.include_node_specific_values = enabled; + self + } + + pub fn build>(self, k: K) -> DebugQuery { + DebugQuery { + k: k.into(), + params: self.params, + } + } +} + +/// The payload to give to the query function. Construct this from the builder. +#[derive(Debug)] +pub struct DebugQuery { + k: Nibbles, + params: DebugQueryParams, +} + +impl From for DebugQuery { + fn from(k: Nibbles) -> Self { + Self { + k, + params: DebugQueryParams::default(), + } + } +} + +/// Extra data that is associated with a node. Only used if +/// `include_node_specific_values` is `true`. +#[derive(Clone, Debug)] +enum ExtraNodeSegmentInfo { + Hash(H256), + Branch { child_mask: u16 }, + Leaf { value: Vec }, +} + +impl Display for ExtraNodeSegmentInfo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ExtraNodeSegmentInfo::Hash(h) => write!(f, "Hash: {:x}", h), + ExtraNodeSegmentInfo::Branch { child_mask } => write!( + f, + "mask: {:#018b} (Num Children: {})", + child_mask, + count_non_empty_branch_children_from_mask(*child_mask) + ), + ExtraNodeSegmentInfo::Leaf { value } => write!(f, "Leaf Value: {}", hex::encode(value)), + } + } +} + +impl ExtraNodeSegmentInfo { + pub(super) fn from_node(n: &Node) -> Option { + match n { + Node::Empty | Node::Extension { .. } => None, + Node::Hash(h) => Some(ExtraNodeSegmentInfo::Hash(*h)), + Node::Branch { children, .. } => Some(ExtraNodeSegmentInfo::Branch { + child_mask: create_child_mask_from_children(children), + }), + Node::Leaf { value, .. } => Some(ExtraNodeSegmentInfo::Leaf { + value: value.clone(), + }), + } + } +} + +fn create_child_mask_from_children(children: &[WrappedNode; 16]) -> u16 { + let mut mask: u16 = 0; + + for (i, child) in children.iter().enumerate().take(16) { + if !matches!(child.as_ref(), Node::Empty) { + mask |= (1 << i) as u16; + } + } + + mask +} + +fn count_non_empty_branch_children_from_mask(mask: u16) -> usize { + let mut num_children = 0; + + for i in 0..16 { + num_children += ((mask & (1 << i)) > 0) as usize; + } + + num_children +} + +#[derive(Clone, Debug)] +pub struct DebugQueryOutput { + k: Nibbles, + node_path: NodePath, + extra_node_info: Vec>, + node_found: bool, + params: DebugQueryParams, +} + +impl Display for DebugQueryOutput { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.fmt_query_header(f)?; + + writeln!(f, "Query path:")?; + for (i, seg) in self + .node_path + .0 + .iter() + .take(self.node_path.0.len() - 1) + .enumerate() + { + Self::fmt_node_based_on_debug_params(f, seg, &self.extra_node_info[i], &self.params)?; + writeln!(f)?; + writeln!(f, "V")?; + } + + if let Some(last_seg) = self.node_path.0.last() { + Self::fmt_node_based_on_debug_params( + f, + last_seg, + &self.extra_node_info[self.node_path.0.len() - 1], + &self.params, + )?; + } + + Ok(()) + } +} + +impl DebugQueryOutput { + fn new(k: Nibbles, params: DebugQueryParams) -> Self { + Self { + k, + node_path: NodePath::default(), + extra_node_info: Vec::default(), + node_found: false, + params, + } + } + + // TODO: Make the output easier to read... + fn fmt_query_header(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "Query Result {{")?; + + writeln!(f, "Queried Key: {}", self.k)?; + writeln!(f, "Node found: {}", self.node_found)?; + + writeln!(f, "}}") + } + + // TODO: Make the output easier to read... + fn fmt_node_based_on_debug_params( + f: &mut fmt::Formatter<'_>, + seg: &PathSegment, + extra_seg_info: &Option, + params: &DebugQueryParams, + ) -> fmt::Result { + let node_type = seg.node_type(); + + if params.include_node_type { + write!(f, "{}", node_type)?; + } + + write!(f, "(")?; + + if params.include_key_piece_per_node { + if let Some(k_piece) = seg.get_key_piece_from_seg_if_present() { + write!(f, "key: {}", k_piece)?; + } + } + + if params.include_node_specific_values { + if let Some(extra_seg_info) = extra_seg_info { + if params.include_key_piece_per_node { + write!(f, ", ")?; + } + + write!(f, "Extra Seg Info: {}", extra_seg_info)?; + } + } + + write!(f, ")")?; + + Ok(()) + } +} + +/// Get debug information on the path taken when querying a key in a given trie. +pub fn get_path_from_query>( + trie: &Node, + q: Q, +) -> DebugQueryOutput { + let q = q.into(); + + let mut out = DebugQueryOutput::new(q.k, q.params); + get_path_from_query_rec(trie, &mut q.k.clone(), &mut out); + + out +} + +fn get_path_from_query_rec( + node: &Node, + curr_key: &mut Nibbles, + query_out: &mut DebugQueryOutput, +) { + let key_piece = get_key_piece_from_node_pulling_from_key_for_branches(node, curr_key); + let seg = get_segment_from_node_and_key_piece(node, &key_piece); + + query_out.node_path.append(seg); + query_out + .extra_node_info + .push(ExtraNodeSegmentInfo::from_node(node)); + + match node { + Node::Empty | Node::Hash(_) => (), + Node::Branch { children, value: _ } => { + let nib = curr_key.pop_next_nibble_front(); + + get_path_from_query_rec(&children[nib as usize], curr_key, query_out) + } + Node::Extension { nibbles, child } => { + get_next_nibbles_from_node_key_clamped(curr_key, nibbles.count); + get_path_from_query_rec(child, curr_key, query_out); + } + Node::Leaf { nibbles, value: _ } => { + let curr_key_next_nibs = + get_next_nibbles_from_node_key_clamped(curr_key, nibbles.count); + + if *nibbles == curr_key_next_nibs { + curr_key.pop_nibbles_front(curr_key_next_nibs.count); + } + } + } + + if curr_key.is_empty() { + query_out.node_found = true; + } +} + +/// Gets the next `n` [`Nibbles`] from the key and clamps it in the case of it +/// going out of range. +fn get_next_nibbles_from_node_key_clamped(key: &Nibbles, n_nibs: usize) -> Nibbles { + let num_nibs_to_get = n_nibs.min(key.count); + key.get_next_nibbles(num_nibs_to_get) +} + +// TODO: Create some simple tests... +#[cfg(test)] +mod tests {} diff --git a/src/lib.rs b/src/lib.rs index b7fd861..676e19b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,5 +19,8 @@ pub mod trie_ops; pub mod trie_subsets; mod utils; +#[cfg(feature = "trie_debug")] +pub mod debug_tools; + #[cfg(test)] mod testing_utils; diff --git a/src/utils.rs b/src/utils.rs index 4c0ea6b..9b87d86 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -5,7 +5,7 @@ use num_traits::PrimInt; use crate::partial_trie::{Node, PartialTrie}; -#[derive(Debug, Eq, PartialEq)] +#[derive(Clone, Debug, Eq, Hash, PartialEq)] /// Simplified trie node type to make logging cleaner. pub(crate) enum TrieNodeType { Empty, @@ -35,13 +35,15 @@ impl From<&Node> for TrieNodeType { impl Display for TrieNodeType { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - TrieNodeType::Empty => write!(f, "Empty"), - TrieNodeType::Hash => write!(f, "Hash"), - TrieNodeType::Branch => write!(f, "Branch"), - TrieNodeType::Extension => write!(f, "Extension"), - TrieNodeType::Leaf => write!(f, "Leaf"), - } + let s = match self { + TrieNodeType::Empty => "Empty", + TrieNodeType::Hash => "Hash", + TrieNodeType::Branch => "Branch", + TrieNodeType::Extension => "Extension", + TrieNodeType::Leaf => "Leaf", + }; + + write!(f, "{}", s) } }