From 0cd57485341b5e70604761b01144619574f6a419 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Mon, 29 Jan 2024 12:42:57 -0700
Subject: [PATCH 01/16] Added some more tests and traces

- One test is a regression test. Next commit will fix bug.
---
 mpt_trie/src/trie_subsets.rs | 104 ++++++++++++++++++++++++++++++-----
 1 file changed, 89 insertions(+), 15 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index 5a540d899..04658d046 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -8,6 +8,7 @@
 use std::sync::Arc;
 
 use ethereum_types::H256;
+use log::trace;
 use thiserror::Error;
 
 use crate::{
@@ -206,6 +207,7 @@ impl<N: PartialTrie> TrackedNodeInfo<N> {
     }
 }
 
+// TODO: Make this interface also work with &[ ... ]...
 /// Create a [`PartialTrie`] subset from a base trie given an iterator of keys
 /// of nodes that may or may not exist in the trie. All nodes traversed by the
 /// keys will not be hashed out in the trie subset. If the key does not exist in
@@ -221,6 +223,7 @@ where
     create_trie_subset_intern(&mut tracked_trie, keys_involved.into_iter())
 }
 
+// TODO: Make this interface also work with &[ ... ]...
 /// Create [`PartialTrie`] subsets from a given base `PartialTrie` given a
 /// iterator of keys per subset needed. See [`create_trie_subset`] for more
 /// info.
@@ -265,6 +268,12 @@ fn mark_nodes_that_are_needed<N: PartialTrie>(
 ) -> SubsetTrieResult<()> {
     trie.info.touched = true;
 
+    trace!(
+        "Sub-trie marking at {:x}, (type: {})",
+        curr_nibbles,
+        TrieNodeType::from(trie.info.underlying_node.deref())
+    );
+
     match &mut trie.node {
         TrackedNodeIntern::Empty => Ok(()),
         TrackedNodeIntern::Hash => match curr_nibbles.is_empty() {
@@ -384,38 +393,69 @@ mod tests {
                 nibbles: nibbles.into(),
             }
         }
+
+        fn with_key(&self, k: &Nibbles) -> bool {
+            self.nibbles.reverse() == *k
+        }
+    }
+
+    fn get_node_full_nibbles_from_list<'a>(
+        nodes: &'a [NodeFullNibbles],
+        k: &'a Nibbles,
+    ) -> Option<&'a NodeFullNibbles> {
+        nodes.iter().find(|x| x.with_key(k))
+    }
+
+    fn get_all_nodes_in_trie(trie: &TrieType) -> Vec<NodeFullNibbles> {
+        get_nodes_in_trie_intern(trie, false)
     }
 
     fn get_all_non_empty_and_hash_nodes_in_trie(trie: &TrieType) -> Vec<NodeFullNibbles> {
+        get_nodes_in_trie_intern(trie, true)
+    }
+
+    fn get_nodes_in_trie_intern(
+        trie: &TrieType,
+        return_on_empty_or_hash: bool,
+    ) -> Vec<NodeFullNibbles> {
         let mut nodes = Vec::new();
-        get_all_non_empty_and_non_hash_nodes_in_trie_intern(trie, Nibbles::default(), &mut nodes);
+        get_nodes_in_trie_intern_rec(
+            trie,
+            Nibbles::default(),
+            &mut nodes,
+            return_on_empty_or_hash,
+        );
 
         nodes
     }
 
-    fn get_all_non_empty_and_non_hash_nodes_in_trie_intern(
+    fn get_nodes_in_trie_intern_rec(
         trie: &TrieType,
         mut curr_nibbles: Nibbles,
         nodes: &mut Vec<NodeFullNibbles>,
+        return_on_empty_or_hash: bool,
     ) {
         match &trie.node {
-            Node::Empty | Node::Hash(_) => return,
+            Node::Empty | Node::Hash(_) => match return_on_empty_or_hash {
+                false => (),
+                true => return,
+            },
             Node::Branch { children, .. } => {
                 for (i, c) in children.iter().enumerate() {
-                    get_all_non_empty_and_non_hash_nodes_in_trie_intern(
+                    get_nodes_in_trie_intern_rec(
                         c,
                         curr_nibbles.merge_nibble(i as u8),
                         nodes,
+                        return_on_empty_or_hash,
                     )
                 }
             }
-            Node::Extension { nibbles, child } => {
-                get_all_non_empty_and_non_hash_nodes_in_trie_intern(
-                    child,
-                    curr_nibbles.merge_nibbles(nibbles),
-                    nodes,
-                )
-            }
+            Node::Extension { nibbles, child } => get_nodes_in_trie_intern_rec(
+                child,
+                curr_nibbles.merge_nibbles(nibbles),
+                nodes,
+                return_on_empty_or_hash,
+            ),
             Node::Leaf { nibbles, .. } => curr_nibbles = curr_nibbles.merge_nibbles(nibbles),
         };
 
@@ -448,7 +488,7 @@ mod tests {
 
     #[test]
     fn encountering_a_hash_node_returns_err() {
-        let trie = HashedPartialTrie::new(Node::Hash(H256::zero()));
+        let trie = TrieType::new(Node::Hash(H256::zero()));
         let res = create_trie_subset(&trie, once(0x1234));
 
         assert!(res.is_err())
@@ -573,6 +613,25 @@ mod tests {
         )));
     }
 
+    fn assert_nodes_are_hashed<K: Clone + Into<Nibbles>>(trie: &TrieType, keys: &[K]) {
+        let nodes = get_all_nodes_in_trie(trie);
+
+        println!("{:#?}", nodes);
+
+        for k in keys.iter().map(|k| k.clone().into()) {
+            let node = get_node_full_nibbles_from_list(&nodes, &k);
+
+            match node {
+                Some(v) => {
+                    if v.n_type != TrieNodeType::Hash {
+                        panic!("Expected trie node at {:x} to be a hash node but it wasn't! (found a {} node instead)", k, v.n_type)
+                    }
+                }
+                None => panic!("Expected a hash node at {:x} but no node was found!", k),
+            }
+        }
+    }
+
     #[test]
     fn all_leafs_of_keys_to_create_subset_are_included_in_subset_for_giant_trie() {
         let (_, trie_subsets, keys_of_subsets) = create_massive_trie_and_subsets(9009);
@@ -585,15 +644,30 @@ mod tests {
 
     #[test]
     fn hash_of_single_leaf_trie_partial_trie_matches_original_trie() {
-        let mut trie = TrieType::default();
-        trie.insert(0x1234, vec![0]);
+        let trie = create_trie_with_large_entry_nodes(&[0x0]);
 
         let base_hash = trie.hash();
-        let partial_trie = create_trie_subset(&trie, vec![0x1234]).unwrap();
+        let partial_trie = create_trie_subset(&trie, [0x1234]).unwrap();
 
         assert_eq!(base_hash, partial_trie.hash());
     }
 
+    #[test]
+    fn sub_trie_that_includes_branch_but_not_children_hashes_out_children() {
+        let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x12345, 0x12346, 0x1234f]);
+        let partial_trie = create_trie_subset(&trie, [0x1234f]).unwrap();
+
+        assert_nodes_are_hashed(&partial_trie, &[0x12345, 0x12346]);
+    }
+
+    #[test]
+    fn sub_trie_for_non_existent_key_that_hits_branch_leaf_hashes_out_leaf() {
+        let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x1234589, 0x12346]);
+        let partial_trie = create_trie_subset(&trie, [0x1234567]).unwrap();
+
+        assert_nodes_are_hashed(&partial_trie, &[0x1234589, 0x12346]);
+    }
+
     #[test]
     fn hash_of_branch_partial_tries_matches_original_trie() {
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x56, 0x12345]);

From b70af7217364f28da4494f407c0f54c21f23718b Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Mon, 29 Jan 2024 13:40:13 -0700
Subject: [PATCH 02/16] Fixed sub-set bug

- Should squash this though, as the regression test still fails.
---
 mpt_trie/src/trie_subsets.rs | 48 +++++++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 15 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index 04658d046..3f4e12a75 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -262,12 +262,17 @@ where
     Ok(create_partial_trie_subset_from_tracked_trie(tracked_trie))
 }
 
+/// For a given key, mark every node that we encounter that is part of the key.
+/// Note that this means non-existent keys passed into this function will mark
+/// nodes to not be hashed that are part of the given key. For example:
+/// - Relevant nodes in trie: [B(0x), B(0x1), L(0x123)]
+/// - For the key `0x1`, the marked nodes would be [B(0x), B(0x1)].
+/// - For the key `0x12`, the marked nodes still would be [B(0x), B(0x1)].
+/// - For the key `0x123`, the marked nodes would be [B(0x), B(0x1), L(0x123)].
 fn mark_nodes_that_are_needed<N: PartialTrie>(
     trie: &mut TrackedNode<N>,
     curr_nibbles: &mut Nibbles,
 ) -> SubsetTrieResult<()> {
-    trie.info.touched = true;
-
     trace!(
         "Sub-trie marking at {:x}, (type: {})",
         curr_nibbles,
@@ -275,35 +280,50 @@ fn mark_nodes_that_are_needed<N: PartialTrie>(
     );
 
     match &mut trie.node {
-        TrackedNodeIntern::Empty => Ok(()),
+        TrackedNodeIntern::Empty => {
+            trie.info.touched = true;
+        }
         TrackedNodeIntern::Hash => match curr_nibbles.is_empty() {
-            false => Err(SubsetTrieError::UnexpectedKey(
-                *curr_nibbles,
-                format!("{:?}", trie),
-            )),
-            true => Ok(()),
+            false => {
+                return Err(SubsetTrieError::UnexpectedKey(
+                    *curr_nibbles,
+                    format!("{:?}", trie),
+                ))
+            }
+            true => {
+                trie.info.touched = true;
+            }
         },
         // Note: If we end up supporting non-fixed sized keys, then we need to also check value.
         TrackedNodeIntern::Branch(children) => {
+            trie.info.touched = true;
+
             // Check against branch value.
             if curr_nibbles.is_empty() {
                 return Ok(());
             }
 
             let nib = curr_nibbles.pop_next_nibble_front();
-            mark_nodes_that_are_needed(&mut children[nib as usize], curr_nibbles)
+            return mark_nodes_that_are_needed(&mut children[nib as usize], curr_nibbles);
         }
         TrackedNodeIntern::Extension(child) => {
             let nibbles = trie.info.get_nibbles_expected();
             let r = curr_nibbles.pop_nibbles_front(nibbles.count);
 
-            match r.nibbles_are_identical_up_to_smallest_count(nibbles) {
-                false => Ok(()),
-                true => mark_nodes_that_are_needed(child, curr_nibbles),
+            if r.nibbles_are_identical_up_to_smallest_count(nibbles) {
+                trie.info.touched = true;
+                return mark_nodes_that_are_needed(child, curr_nibbles);
+            }
+        }
+        TrackedNodeIntern::Leaf => {
+            let (k, _) = trie.info.get_leaf_nibbles_and_value_expected();
+            if k == curr_nibbles {
+                trie.info.touched = true;
             }
         }
-        TrackedNodeIntern::Leaf => Ok(()),
     }
+
+    Ok(())
 }
 
 fn create_partial_trie_subset_from_tracked_trie<N: PartialTrie>(
@@ -616,8 +636,6 @@ mod tests {
     fn assert_nodes_are_hashed<K: Clone + Into<Nibbles>>(trie: &TrieType, keys: &[K]) {
         let nodes = get_all_nodes_in_trie(trie);
 
-        println!("{:#?}", nodes);
-
         for k in keys.iter().map(|k| k.clone().into()) {
             let node = get_node_full_nibbles_from_list(&nodes, &k);
 

From efcff576cdd8da752644ae85025cc87280953372 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Mon, 29 Jan 2024 16:28:56 -0700
Subject: [PATCH 03/16] Added another large battery test

---
 mpt_trie/src/trie_subsets.rs | 72 ++++++++++++++++++++++++++++++------
 1 file changed, 61 insertions(+), 11 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index 3f4e12a75..f3660696e 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -374,7 +374,10 @@ fn reset_tracked_trie_state<N: PartialTrie>(tracked_node: &mut TrackedNode<N>) {
 
 #[cfg(test)]
 mod tests {
-    use std::{collections::HashSet, iter::once};
+    use std::{
+        collections::{HashMap, HashSet},
+        iter::once,
+    };
 
     use ethereum_types::H256;
 
@@ -633,19 +636,47 @@ mod tests {
         )));
     }
 
-    fn assert_nodes_are_hashed<K: Clone + Into<Nibbles>>(trie: &TrieType, keys: &[K]) {
+    fn assert_nodes_are_leaf_nodes<K: Clone + Into<Nibbles>, I: IntoIterator<Item = K>>(
+        trie: &TrieType,
+        keys: I,
+    ) {
+        assert_keys_point_to_nodes_of_type(
+            trie,
+            keys.into_iter().map(|k| (k.into(), TrieNodeType::Leaf)),
+        )
+    }
+
+    fn assert_nodes_are_hash_nodes<K: Clone + Into<Nibbles>, I: IntoIterator<Item = K>>(
+        trie: &TrieType,
+        keys: I,
+    ) {
+        assert_keys_point_to_nodes_of_type(
+            trie,
+            keys.into_iter().map(|k| (k.into(), TrieNodeType::Hash)),
+        )
+    }
+
+    fn assert_keys_point_to_nodes_of_type(
+        trie: &TrieType,
+        keys: impl Iterator<Item = (Nibbles, TrieNodeType)>,
+    ) {
         let nodes = get_all_nodes_in_trie(trie);
+        let keys_to_node_types: HashMap<_, _> =
+            HashMap::from_iter(nodes.into_iter().map(|n| (n.nibbles, n.n_type)));
 
-        for k in keys.iter().map(|k| k.clone().into()) {
-            let node = get_node_full_nibbles_from_list(&nodes, &k);
+        for (k, expected_n_type) in keys.map(|(k, t)| (k, t)) {
+            let actual_n_type_opt = keys_to_node_types.get(&k);
 
-            match node {
-                Some(v) => {
-                    if v.n_type != TrieNodeType::Hash {
-                        panic!("Expected trie node at {:x} to be a hash node but it wasn't! (found a {} node instead)", k, v.n_type)
+            match actual_n_type_opt {
+                Some(actual_n_type) => {
+                    if *actual_n_type != expected_n_type {
+                        panic!("Expected trie node at {:x} to be a {} node but it wasn't! (found a {} node instead)", k, expected_n_type, actual_n_type)
                     }
                 }
-                None => panic!("Expected a hash node at {:x} but no node was found!", k),
+                None => panic!(
+                    "Expected a {} node at {:x} but no node was found!",
+                    expected_n_type, k
+                ),
             }
         }
     }
@@ -675,7 +706,7 @@ mod tests {
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x12345, 0x12346, 0x1234f]);
         let partial_trie = create_trie_subset(&trie, [0x1234f]).unwrap();
 
-        assert_nodes_are_hashed(&partial_trie, &[0x12345, 0x12346]);
+        assert_nodes_are_hash_nodes(&partial_trie, [0x12345, 0x12346]);
     }
 
     #[test]
@@ -683,7 +714,7 @@ mod tests {
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x1234589, 0x12346]);
         let partial_trie = create_trie_subset(&trie, [0x1234567]).unwrap();
 
-        assert_nodes_are_hashed(&partial_trie, &[0x1234589, 0x12346]);
+        assert_nodes_are_hash_nodes(&partial_trie, [0x1234589, 0x12346]);
     }
 
     #[test]
@@ -713,6 +744,25 @@ mod tests {
             .all(|p_tree| p_tree.hash() == base_hash))
     }
 
+    #[test]
+    fn giant_random_partial_tries_hashes_leaves_correctly() {
+        let (base_trie, trie_subsets, leaf_keys_per_trie) = create_massive_trie_and_subsets(9011);
+        let all_keys: Vec<Nibbles> = base_trie.keys().collect();
+
+        for (partial_trie, leaf_trie_keys) in
+            trie_subsets.into_iter().zip(leaf_keys_per_trie.into_iter())
+        {
+            let leaf_keys_lookup: HashSet<Nibbles> = leaf_trie_keys.iter().cloned().collect();
+            let keys_of_hash_nodes = all_keys
+                .iter()
+                .filter(|k| !leaf_keys_lookup.contains(k))
+                .cloned();
+
+            assert_nodes_are_leaf_nodes(&partial_trie, leaf_trie_keys);
+            assert_nodes_are_hash_nodes(&partial_trie, keys_of_hash_nodes);
+        }
+    }
+
     fn create_massive_trie_and_subsets(seed: u64) -> (TrieType, Vec<TrieType>, Vec<Vec<Nibbles>>) {
         let trie_size = MASSIVE_TEST_NUM_SUB_TRIES * MASSIVE_TEST_NUM_SUB_TRIE_SIZE;
 

From d7d6a811111f38f1592e43621ffd599bba6ed45b Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Wed, 31 Jan 2024 09:43:23 -0700
Subject: [PATCH 04/16] Few small test fixes

---
 mpt_trie/src/trie_subsets.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index f3660696e..f7e7860ee 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -662,9 +662,9 @@ mod tests {
     ) {
         let nodes = get_all_nodes_in_trie(trie);
         let keys_to_node_types: HashMap<_, _> =
-            HashMap::from_iter(nodes.into_iter().map(|n| (n.nibbles, n.n_type)));
+            HashMap::from_iter(nodes.into_iter().map(|n| (n.nibbles.reverse(), n.n_type)));
 
-        for (k, expected_n_type) in keys.map(|(k, t)| (k, t)) {
+        for (k, expected_n_type) in keys {
             let actual_n_type_opt = keys_to_node_types.get(&k);
 
             match actual_n_type_opt {
@@ -714,7 +714,8 @@ mod tests {
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x1234589, 0x12346]);
         let partial_trie = create_trie_subset(&trie, [0x1234567]).unwrap();
 
-        assert_nodes_are_hash_nodes(&partial_trie, [0x1234589, 0x12346]);
+        // Note that `0x1234589` gets hashed at the branch slot at `0x12345`.
+        assert_nodes_are_hash_nodes(&partial_trie, [0x12345, 0x12346]);
     }
 
     #[test]

From 9643885ea6126aa3ebaa7049de66f33bd504904f Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Thu, 8 Feb 2024 13:09:38 -0700
Subject: [PATCH 05/16] Moved some types around in preparation for new logic

---
 mpt_trie/src/debug_tools/common.rs |  99 +--------------------
 mpt_trie/src/debug_tools/diff.rs   |   3 +-
 mpt_trie/src/debug_tools/query.rs  |  10 +--
 mpt_trie/src/lib.rs                |   2 +-
 mpt_trie/src/trie_subsets.rs       |   2 +-
 mpt_trie/src/utils.rs              | 138 ++++++++++++++++++++++++++++-
 6 files changed, 145 insertions(+), 109 deletions(-)

diff --git a/mpt_trie/src/debug_tools/common.rs b/mpt_trie/src/debug_tools/common.rs
index 5eba83624..cb38d5c47 100644
--- a/mpt_trie/src/debug_tools/common.rs
+++ b/mpt_trie/src/debug_tools/common.rs
@@ -1,66 +1,9 @@
 //! Common utilities for the debugging tools.
-use std::fmt::{self, Display};
-
 use crate::{
-    nibbles::{Nibble, Nibbles},
+    nibbles::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<Nibbles> {
-        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<T: PartialTrie>(
-    n: &Node<T>,
-    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
@@ -87,43 +30,3 @@ pub(super) fn get_key_piece_from_node<T: PartialTrie>(n: &Node<T>) -> Nibbles {
         Node::Extension { nibbles, child: _ } | Node::Leaf { nibbles, value: _ } => *nibbles,
     }
 }
-
-#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
-/// A vector of path segments representing a path in the trie.
-pub struct NodePath(pub(super) Vec<PathSegment>);
-
-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/mpt_trie/src/debug_tools/diff.rs b/mpt_trie/src/debug_tools/diff.rs
index 3ed857492..c10179c69 100644
--- a/mpt_trie/src/debug_tools/diff.rs
+++ b/mpt_trie/src/debug_tools/diff.rs
@@ -30,7 +30,8 @@ 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 super::common::get_key_piece_from_node;
+use crate::utils::{get_segment_from_node_and_key_piece, NodePath};
 use crate::{
     nibbles::Nibbles,
     partial_trie::{HashedPartialTrie, Node, PartialTrie},
diff --git a/mpt_trie/src/debug_tools/query.rs b/mpt_trie/src/debug_tools/query.rs
index 8ce9ddfad..a09976343 100644
--- a/mpt_trie/src/debug_tools/query.rs
+++ b/mpt_trie/src/debug_tools/query.rs
@@ -5,13 +5,11 @@ 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 super::common::get_key_piece_from_node_pulling_from_key_for_branches;
 use crate::{
     nibbles::Nibbles,
     partial_trie::{Node, PartialTrie, WrappedNode},
+    utils::{get_segment_from_node_and_key_piece, NodePath, PathSegment},
 };
 
 /// Params controlling how much information is reported in the query output.
@@ -159,7 +157,9 @@ fn count_non_empty_branch_children_from_mask(mask: u16) -> usize {
 /// of the path used for searching for a key in the trie.
 pub struct DebugQueryOutput {
     k: Nibbles,
-    node_path: NodePath,
+
+    /// The nodes hit during the query.
+    pub node_path: NodePath,
     extra_node_info: Vec<Option<ExtraNodeSegmentInfo>>,
     node_found: bool,
     params: DebugQueryParams,
diff --git a/mpt_trie/src/lib.rs b/mpt_trie/src/lib.rs
index c359e0430..6964409ac 100644
--- a/mpt_trie/src/lib.rs
+++ b/mpt_trie/src/lib.rs
@@ -20,7 +20,7 @@ pub mod partial_trie;
 mod trie_hashing;
 pub mod trie_ops;
 pub mod trie_subsets;
-mod utils;
+pub mod utils;
 
 #[cfg(feature = "trie_debug")]
 pub mod debug_tools;
diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index f7e7860ee..5ae8c4b75 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -384,7 +384,7 @@ mod tests {
     use super::{create_trie_subset, create_trie_subsets};
     use crate::{
         nibbles::Nibbles,
-        partial_trie::{HashedPartialTrie, Node, PartialTrie},
+        partial_trie::{Node, PartialTrie},
         testing_utils::{
             create_trie_with_large_entry_nodes, generate_n_random_fixed_trie_value_entries,
             handmade_trie_1, TrieType,
diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs
index 9b87d8606..f784f6117 100644
--- a/mpt_trie/src/utils.rs
+++ b/mpt_trie/src/utils.rs
@@ -1,17 +1,35 @@
-use std::{fmt::Display, ops::BitAnd, sync::Arc};
+//! Various types and logic that don't fit well into any other module.
+
+use std::{
+    fmt::{self, Display},
+    ops::BitAnd,
+    sync::Arc,
+};
 
 use ethereum_types::{H256, U512};
 use num_traits::PrimInt;
 
-use crate::partial_trie::{Node, PartialTrie};
+use crate::{
+    nibbles::{Nibble, Nibbles},
+    partial_trie::{Node, PartialTrie},
+};
 
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 /// Simplified trie node type to make logging cleaner.
-pub(crate) enum TrieNodeType {
+pub enum TrieNodeType {
+    /// Empty node.
     Empty,
+
+    /// Hash node.
     Hash,
+
+    /// Branch node.
     Branch,
+
+    /// Extension node.
     Extension,
+
+    /// Leaf node.
     Leaf,
 }
 
@@ -58,3 +76,117 @@ pub(crate) fn create_mask_of_1s(amt: usize) -> U512 {
 pub(crate) fn bytes_to_h256(b: &[u8; 32]) -> H256 {
     keccak_hash::H256::from_slice(b)
 }
+
+/// Minimal key information of "segments" (nodes) used to construct trie
+/// "traces" of a trie query. Unlike [`TrieNodeType`], this type also contains
+/// the key piece of the node if applicable (eg. [`Node::Empty`] &
+/// [`Node::Hash`] do not have associated key pieces).
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub enum PathSegment {
+    /// Empty node.
+    Empty,
+
+    /// Hash node.
+    Hash,
+
+    /// Branch node along with the nibble of the child taken.
+    Branch(Nibble),
+
+    /// Extension node along with the key piece of the node.
+    Extension(Nibbles),
+
+    /// Leaf node along wth the key piece of the node.
+    Leaf(Nibbles),
+}
+
+/// A vector of path segments representing a path in the trie.
+#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
+pub struct NodePath(pub Vec<PathSegment>);
+
+impl NodePath {
+    pub(crate) fn dup_and_append(&self, seg: PathSegment) -> Self {
+        let mut duped_vec = self.0.clone();
+        duped_vec.push(seg);
+
+        Self(duped_vec)
+    }
+
+    pub(crate) 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(())
+    }
+}
+
+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 {
+    /// Get the node type of the [`PathSegment`].
+    pub 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,
+        }
+    }
+
+    /// Extracts the key piece used by the segment (if applicable).
+    pub fn get_key_piece_from_seg_if_present(&self) -> Option<Nibbles> {
+        match self {
+            PathSegment::Empty | PathSegment::Hash => None,
+            PathSegment::Branch(nib) => Some(Nibbles::from_nibble(*nib)),
+            PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => Some(*nibs),
+        }
+    }
+}
+
+/// Creates a [`PathSegment`] given a node and a key we are querying.
+///
+/// This function is intended to be used during a trie query as we are
+/// traversing down a trie. Depending on the current node, we pop off nibbles
+/// and use these to create `PathSegment`s.
+pub fn get_segment_from_node_and_key_piece<T: PartialTrie>(
+    n: &Node<T>,
+    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),
+    }
+}

From c59cdd9e60afba0ca4598bc3550f211a420010c5 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Fri, 9 Feb 2024 19:14:30 -0800
Subject: [PATCH 06/16] Logic for trie queries

- Note that this is different from `debug_query`. This logic is a lot
  more light-weight and meant for actual logic (instead of for debugging
  purposes).
---
 mpt_trie/src/lib.rs           |   1 +
 mpt_trie/src/special_query.rs | 153 ++++++++++++++++++++++++++++++++++
 2 files changed, 154 insertions(+)
 create mode 100644 mpt_trie/src/special_query.rs

diff --git a/mpt_trie/src/lib.rs b/mpt_trie/src/lib.rs
index 6964409ac..48c75ae3e 100644
--- a/mpt_trie/src/lib.rs
+++ b/mpt_trie/src/lib.rs
@@ -17,6 +17,7 @@
 
 pub mod nibbles;
 pub mod partial_trie;
+mod special_query;
 mod trie_hashing;
 pub mod trie_ops;
 pub mod trie_subsets;
diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
new file mode 100644
index 000000000..f6b71acd9
--- /dev/null
+++ b/mpt_trie/src/special_query.rs
@@ -0,0 +1,153 @@
+//! Specialized queries that users of the library may need that require
+//! knowledge of the private internal trie state.
+
+use crate::{
+    nibbles::Nibbles,
+    partial_trie::{Node, PartialTrie, WrappedNode},
+    utils::PathSegment,
+};
+
+#[derive(Debug)]
+struct PathSegmentIterator<N: PartialTrie> {
+    curr_node: WrappedNode<N>,
+    curr_key: Nibbles,
+
+    // Although wrapping `curr_node` in an option might be more "Rust like", the logic is a lot
+    // cleaner with a bool.
+    terminated: bool,
+}
+
+impl<T: PartialTrie> Iterator for PathSegmentIterator<T> {
+    type Item = PathSegment;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.terminated {
+            return None;
+        }
+
+        match self.curr_node.as_ref() {
+            Node::Empty => {
+                self.terminated = true;
+                Some(PathSegment::Empty)
+            }
+            Node::Hash(_) => {
+                self.terminated = true;
+                Some(PathSegment::Hash)
+            }
+            Node::Branch { children, .. } => {
+                if self.curr_key.is_empty() {
+                    self.terminated = true;
+                    return None;
+                }
+
+                let nib = self.curr_key.pop_next_nibble_front();
+                self.curr_node = children[nib as usize].clone();
+
+                Some(PathSegment::Branch(nib))
+            }
+            Node::Extension { nibbles, child } => {
+                match self
+                    .curr_key
+                    .nibbles_are_identical_up_to_smallest_count(nibbles)
+                {
+                    false => {
+                        self.terminated = true;
+                        None
+                    }
+                    true => {
+                        pop_nibbles_clamped(&mut self.curr_key, nibbles.count);
+                        let res = Some(PathSegment::Extension(*nibbles));
+                        self.curr_node = child.clone();
+
+                        res
+                    }
+                }
+            }
+            Node::Leaf { nibbles, .. } => {
+                self.terminated = true;
+
+                match self.curr_key == *nibbles {
+                    false => None,
+                    true => Some(PathSegment::Leaf(*nibbles)),
+                }
+            }
+        }
+    }
+}
+
+fn pop_nibbles_clamped(nibbles: &mut Nibbles, n: usize) -> Nibbles {
+    let n_nibs_to_pop = nibbles.count.min(n);
+    nibbles.pop_nibbles_front(n_nibs_to_pop)
+}
+
+// TODO: Move to a blanket impl...
+// TODO: Could make this return an `Iterator` with some work...
+/// Returns all nodes in the trie that are traversed given a query (key).
+///
+/// Note that if the key does not match the entire key of a node (eg. the
+/// remaining key is `0x34` but the next key is a leaf with the key `0x3456`),
+/// then the leaf will not appear in the query output.
+pub fn get_path_for_query<K, T: PartialTrie>(
+    trie: &Node<T>,
+    k: K,
+) -> impl Iterator<Item = PathSegment>
+where
+    K: Into<Nibbles>,
+{
+    PathSegmentIterator {
+        curr_node: trie.clone().into(),
+        curr_key: k.into(),
+        terminated: false,
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::str::FromStr;
+
+    use super::get_path_for_query;
+    use crate::{nibbles::Nibbles, testing_utils::handmade_trie_1, utils::PathSegment};
+
+    #[test]
+    fn query_iter_works() {
+        let (trie, ks) = handmade_trie_1();
+
+        // ks --> vec![0x1234, 0x1324, 0x132400005_u64, 0x2001, 0x2002];
+        let res = vec![
+            vec![
+                PathSegment::Branch(1),
+                PathSegment::Branch(2),
+                PathSegment::Leaf(0x34.into()),
+            ],
+            vec![
+                PathSegment::Branch(1),
+                PathSegment::Branch(3),
+                PathSegment::Extension(0x24.into()),
+            ],
+            vec![
+                PathSegment::Branch(1),
+                PathSegment::Branch(3),
+                PathSegment::Extension(0x24.into()),
+                PathSegment::Branch(0),
+                PathSegment::Leaf(Nibbles::from_str("0x0005").unwrap()),
+            ],
+            vec![
+                PathSegment::Branch(2),
+                PathSegment::Extension(Nibbles::from_str("0x00").unwrap()),
+                PathSegment::Branch(0x1),
+                PathSegment::Leaf(Nibbles::default()),
+            ],
+            vec![
+                PathSegment::Branch(2),
+                PathSegment::Extension(Nibbles::from_str("0x00").unwrap()),
+                PathSegment::Branch(0x2),
+                PathSegment::Leaf(Nibbles::default()),
+            ],
+        ];
+
+        for (q, expected) in ks.into_iter().zip(res.into_iter()) {
+            let res: Vec<_> = get_path_for_query(&trie.node, q).collect();
+            assert_eq!(res, expected)
+        }
+    }
+}

From 3323697c46b293424343c0cc3856e61f5f0f5050 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Fri, 9 Feb 2024 19:26:07 -0800
Subject: [PATCH 07/16] Added docs to query logic

---
 mpt_trie/src/special_query.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
index f6b71acd9..b73abdbc0 100644
--- a/mpt_trie/src/special_query.rs
+++ b/mpt_trie/src/special_query.rs
@@ -9,7 +9,10 @@ use crate::{
 
 #[derive(Debug)]
 struct PathSegmentIterator<N: PartialTrie> {
+    /// The next node in the trie to query with the remaining key.
     curr_node: WrappedNode<N>,
+
+    /// The remaining part of the key as we traverse down the trie.
     curr_key: Nibbles,
 
     // Although wrapping `curr_node` in an option might be more "Rust like", the logic is a lot
@@ -35,6 +38,7 @@ impl<T: PartialTrie> Iterator for PathSegmentIterator<T> {
                 Some(PathSegment::Hash)
             }
             Node::Branch { children, .. } => {
+                // Our query key has ended. Stop here.
                 if self.curr_key.is_empty() {
                     self.terminated = true;
                     return None;
@@ -51,6 +55,7 @@ impl<T: PartialTrie> Iterator for PathSegmentIterator<T> {
                     .nibbles_are_identical_up_to_smallest_count(nibbles)
                 {
                     false => {
+                        // Only a partial match. Stop.
                         self.terminated = true;
                         None
                     }
@@ -75,6 +80,8 @@ impl<T: PartialTrie> Iterator for PathSegmentIterator<T> {
     }
 }
 
+/// Attempt to pop `n` nibbles from the given [`Nibbles`] and "clamp" the
+/// nibbles popped by not popping more nibbles than exist.
 fn pop_nibbles_clamped(nibbles: &mut Nibbles, n: usize) -> Nibbles {
     let n_nibs_to_pop = nibbles.count.min(n);
     nibbles.pop_nibbles_front(n_nibs_to_pop)

From cca640a79310369d2deb46aede726831825c1ea0 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Fri, 9 Feb 2024 20:05:29 -0800
Subject: [PATCH 08/16] Some refactoring related to the query logic

---
 mpt_trie/src/debug_tools/diff.rs  | 16 +++++-----
 mpt_trie/src/debug_tools/query.rs |  6 ++--
 mpt_trie/src/special_query.rs     | 15 ++++------
 mpt_trie/src/utils.rs             | 50 +++++++++++++++++++++----------
 4 files changed, 51 insertions(+), 36 deletions(-)

diff --git a/mpt_trie/src/debug_tools/diff.rs b/mpt_trie/src/debug_tools/diff.rs
index c10179c69..d6271cfb1 100644
--- a/mpt_trie/src/debug_tools/diff.rs
+++ b/mpt_trie/src/debug_tools/diff.rs
@@ -31,7 +31,7 @@ use std::{fmt::Display, ops::Deref};
 use ethereum_types::H256;
 
 use super::common::get_key_piece_from_node;
-use crate::utils::{get_segment_from_node_and_key_piece, NodePath};
+use crate::utils::{get_segment_from_node_and_key_piece, TriePath};
 use crate::{
     nibbles::Nibbles,
     partial_trie::{HashedPartialTrie, Node, PartialTrie},
@@ -84,7 +84,7 @@ pub struct DiffPoint {
     /// The depth of the point in both tries.
     pub depth: usize,
     /// The path of the point in both tries.
-    pub path: NodePath,
+    pub path: TriePath,
     /// The node key in both tries.
     pub key: Nibbles,
     /// The node info in the first trie.
@@ -98,7 +98,7 @@ impl DiffPoint {
         child_a: &HashedPartialTrie,
         child_b: &HashedPartialTrie,
         parent_k: Nibbles,
-        path: NodePath,
+        path: TriePath,
     ) -> 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));
@@ -223,7 +223,7 @@ impl DepthNodeDiffState {
         parent_k: &Nibbles,
         child_a: &HashedPartialTrie,
         child_b: &HashedPartialTrie,
-        path: NodePath,
+        path: TriePath,
     ) {
         if field
             .as_ref()
@@ -243,7 +243,7 @@ struct DepthDiffPerCallState<'a> {
     curr_depth: usize,
 
     // Horribly inefficient, but these are debug tools, so I think we get a pass.
-    curr_path: NodePath,
+    curr_path: TriePath,
 }
 
 impl<'a> DepthDiffPerCallState<'a> {
@@ -260,7 +260,7 @@ impl<'a> DepthDiffPerCallState<'a> {
             b,
             curr_key,
             curr_depth,
-            curr_path: NodePath::default(),
+            curr_path: TriePath::default(),
         }
     }
 
@@ -412,7 +412,7 @@ fn get_value_from_node<T: PartialTrie>(n: &Node<T>) -> Option<&Vec<u8>> {
 
 #[cfg(test)]
 mod tests {
-    use super::{create_diff_between_tries, DiffPoint, NodeInfo, NodePath};
+    use super::{create_diff_between_tries, DiffPoint, NodeInfo, TriePath};
     use crate::{
         nibbles::Nibbles,
         partial_trie::{HashedPartialTrie, PartialTrie},
@@ -448,7 +448,7 @@ mod tests {
 
         let expected = DiffPoint {
             depth: 0,
-            path: NodePath(vec![]),
+            path: TriePath(vec![]),
             key: Nibbles::default(),
             a_info: expected_a,
             b_info: expected_b,
diff --git a/mpt_trie/src/debug_tools/query.rs b/mpt_trie/src/debug_tools/query.rs
index a09976343..d5ad82a17 100644
--- a/mpt_trie/src/debug_tools/query.rs
+++ b/mpt_trie/src/debug_tools/query.rs
@@ -9,7 +9,7 @@ use super::common::get_key_piece_from_node_pulling_from_key_for_branches;
 use crate::{
     nibbles::Nibbles,
     partial_trie::{Node, PartialTrie, WrappedNode},
-    utils::{get_segment_from_node_and_key_piece, NodePath, PathSegment},
+    utils::{get_segment_from_node_and_key_piece, PathSegment, TriePath},
 };
 
 /// Params controlling how much information is reported in the query output.
@@ -159,7 +159,7 @@ pub struct DebugQueryOutput {
     k: Nibbles,
 
     /// The nodes hit during the query.
-    pub node_path: NodePath,
+    pub node_path: TriePath,
     extra_node_info: Vec<Option<ExtraNodeSegmentInfo>>,
     node_found: bool,
     params: DebugQueryParams,
@@ -199,7 +199,7 @@ impl DebugQueryOutput {
     fn new(k: Nibbles, params: DebugQueryParams) -> Self {
         Self {
             k,
-            node_path: NodePath::default(),
+            node_path: TriePath::default(),
             extra_node_info: Vec::default(),
             node_found: false,
             params,
diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
index b73abdbc0..69558207f 100644
--- a/mpt_trie/src/special_query.rs
+++ b/mpt_trie/src/special_query.rs
@@ -8,7 +8,7 @@ use crate::{
 };
 
 #[derive(Debug)]
-struct PathSegmentIterator<N: PartialTrie> {
+pub struct TriePathIter<N: PartialTrie> {
     /// The next node in the trie to query with the remaining key.
     curr_node: WrappedNode<N>,
 
@@ -20,7 +20,7 @@ struct PathSegmentIterator<N: PartialTrie> {
     terminated: bool,
 }
 
-impl<T: PartialTrie> Iterator for PathSegmentIterator<T> {
+impl<T: PartialTrie> Iterator for TriePathIter<T> {
     type Item = PathSegment;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -94,14 +94,11 @@ fn pop_nibbles_clamped(nibbles: &mut Nibbles, n: usize) -> Nibbles {
 /// Note that if the key does not match the entire key of a node (eg. the
 /// remaining key is `0x34` but the next key is a leaf with the key `0x3456`),
 /// then the leaf will not appear in the query output.
-pub fn get_path_for_query<K, T: PartialTrie>(
-    trie: &Node<T>,
-    k: K,
-) -> impl Iterator<Item = PathSegment>
+pub fn path_for_query<K, T: PartialTrie>(trie: &Node<T>, k: K) -> impl Iterator<Item = PathSegment>
 where
     K: Into<Nibbles>,
 {
-    PathSegmentIterator {
+    TriePathIter {
         curr_node: trie.clone().into(),
         curr_key: k.into(),
         terminated: false,
@@ -112,7 +109,7 @@ where
 mod test {
     use std::str::FromStr;
 
-    use super::get_path_for_query;
+    use super::path_for_query;
     use crate::{nibbles::Nibbles, testing_utils::handmade_trie_1, utils::PathSegment};
 
     #[test]
@@ -153,7 +150,7 @@ mod test {
         ];
 
         for (q, expected) in ks.into_iter().zip(res.into_iter()) {
-            let res: Vec<_> = get_path_for_query(&trie.node, q).collect();
+            let res: Vec<_> = path_for_query(&trie.node, q).collect();
             assert_eq!(res, expected)
         }
     }
diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs
index f784f6117..b6c0047af 100644
--- a/mpt_trie/src/utils.rs
+++ b/mpt_trie/src/utils.rs
@@ -12,6 +12,7 @@ use num_traits::PrimInt;
 use crate::{
     nibbles::{Nibble, Nibbles},
     partial_trie::{Node, PartialTrie},
+    special_query::TriePathIter,
 };
 
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
@@ -101,9 +102,28 @@ pub enum PathSegment {
 
 /// A vector of path segments representing a path in the trie.
 #[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
-pub struct NodePath(pub Vec<PathSegment>);
+pub struct TriePath(pub Vec<PathSegment>);
 
-impl NodePath {
+impl Display for TriePath {
+    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(())
+    }
+}
+
+impl TriePath {
     pub(crate) fn dup_and_append(&self, seg: PathSegment) -> Self {
         let mut duped_vec = self.0.clone();
         duped_vec.push(seg);
@@ -120,22 +140,20 @@ impl NodePath {
     }
 }
 
-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, " --> ")?;
-        }
+/// A trie path that is constructed lazily.
+#[derive(Debug)]
+pub struct TriePathLazy<T: PartialTrie>(TriePathIter<T>);
 
-        // Avoid the extra `-->` for the last elem.
-        if let Some(seg) = self.0.last() {
-            Self::write_elem(f, seg)?;
-        }
+impl<T: PartialTrie> TriePathLazy<T> {
+    /// Extract the key from the trie path.
+    pub fn into_key(self) -> Nibbles {
+        todo!()
+    }
+}
 
-        Ok(())
+impl<T: PartialTrie> From<TriePathIter<T>> for TriePathLazy<T> {
+    fn from(v: TriePathIter<T>) -> Self {
+        Self(v)
     }
 }
 

From ee14801a3e93123edc8256c0a3d490725a7fb832 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Fri, 9 Feb 2024 20:13:43 -0800
Subject: [PATCH 09/16] Added logic to reconstruct a key from a trie path

---
 mpt_trie/src/utils.rs | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs
index b6c0047af..ce0c61af1 100644
--- a/mpt_trie/src/utils.rs
+++ b/mpt_trie/src/utils.rs
@@ -142,12 +142,24 @@ impl TriePath {
 
 /// A trie path that is constructed lazily.
 #[derive(Debug)]
-pub struct TriePathLazy<T: PartialTrie>(TriePathIter<T>);
+pub struct TriePathLazy<T: PartialTrie>(pub TriePathIter<T>);
 
 impl<T: PartialTrie> TriePathLazy<T> {
     /// Extract the key from the trie path.
     pub fn into_key(self) -> Nibbles {
-        todo!()
+        let mut key = Nibbles::default();
+
+        for seg in self.0 {
+            match seg {
+                PathSegment::Empty | PathSegment::Hash => (),
+                PathSegment::Branch(nib) => key.push_nibble_front(nib),
+                PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => {
+                    key.push_nibbles_front(&nibs)
+                }
+            }
+        }
+
+        key
     }
 }
 

From 7383858bf8b76460c34f72164dee917cc7d5e75a Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Wed, 14 Feb 2024 17:01:50 -0700
Subject: [PATCH 10/16] Added missing function `push_nibbles_back`

- Also added missing tests for all `push_*` functions.
---
 mpt_trie/src/nibbles.rs | 91 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 88 insertions(+), 3 deletions(-)

diff --git a/mpt_trie/src/nibbles.rs b/mpt_trie/src/nibbles.rs
index 94eb8cafc..f5d63a043 100644
--- a/mpt_trie/src/nibbles.rs
+++ b/mpt_trie/src/nibbles.rs
@@ -363,11 +363,11 @@ impl Nibbles {
     /// Appends `Nibbles` to the front.
     ///
     /// # Panics
-    /// Panics if appending the `Nibble` causes an overflow (total nibbles >
+    /// Panics if appending the `Nibbles` causes an overflow (total nibbles >
     /// 64).
     pub fn push_nibbles_front(&mut self, n: &Self) {
         let new_count = self.count + n.count;
-        assert!(new_count <= 64);
+        self.nibbles_append_safety_asserts(new_count);
 
         let shift_amt = 4 * self.count;
 
@@ -375,6 +375,21 @@ impl Nibbles {
         self.packed |= n.packed << shift_amt;
     }
 
+    /// Appends `Nibbles` to the back.
+    ///
+    /// # Panics
+    /// Panics if appending the `Nibbles` causes an overflow (total nibbles >
+    /// 64).
+    pub fn push_nibbles_back(&mut self, n: &Self) {
+        let new_count = self.count + n.count;
+        self.nibbles_append_safety_asserts(new_count);
+
+        let shift_amt = 4 * n.count;
+
+        self.count = new_count;
+        self.packed = (self.packed << shift_amt) | n.packed;
+    }
+
     /// Gets the nibbles at the range specified, where `0` is the next nibble.
     ///
     /// # Panics
@@ -765,6 +780,10 @@ impl Nibbles {
         assert!(n < 16);
     }
 
+    fn nibbles_append_safety_asserts(&self, new_count: usize) {
+        assert!(new_count <= 64);
+    }
+
     // TODO: REMOVE BEFORE NEXT CRATE VERSION! THIS IS A TEMP HACK!
     /// Converts to u256 returning an error if not possible.
     pub fn try_into_u256(&self) -> Result<U256, String> {
@@ -787,6 +806,9 @@ mod tests {
     use super::{Nibble, Nibbles, ToNibbles};
     use crate::nibbles::FromHexPrefixError;
 
+    const LONG_ZERO_NIBS_STR_LEN_63: &str =
+        "0x000000000000000000000000000000000000000000000000000000000000000";
+
     #[test]
     fn get_nibble_works() {
         let n = Nibbles::from(0x1234);
@@ -898,6 +920,69 @@ mod tests {
         assert_eq!(res, expected_resulting_nibbles);
     }
 
+    #[test]
+    fn push_nibble_front_works() {
+        test_and_assert_nib_push_func(Nibbles::default(), 0x1, |n| n.push_nibble_front(0x1));
+        test_and_assert_nib_push_func(0x1, 0x21, |n| n.push_nibble_front(0x2));
+        test_and_assert_nib_push_func(
+            Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
+            Nibbles::from_str("0x1000000000000000000000000000000000000000000000000000000000000000")
+                .unwrap(),
+            |n| n.push_nibble_front(0x1),
+        );
+    }
+
+    #[test]
+    fn push_nibble_back_works() {
+        test_and_assert_nib_push_func(Nibbles::default(), 0x1, |n| n.push_nibble_back(0x1));
+        test_and_assert_nib_push_func(0x1, 0x12, |n| n.push_nibble_back(0x2));
+        test_and_assert_nib_push_func(
+            Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
+            Nibbles::from_str("0x0000000000000000000000000000000000000000000000000000000000000001")
+                .unwrap(),
+            |n| n.push_nibble_back(0x1),
+        );
+    }
+
+    #[test]
+    fn push_nibbles_front_works() {
+        test_and_assert_nib_push_func(Nibbles::default(), 0x1234, |n| {
+            n.push_nibbles_front(&0x1234.into())
+        });
+        test_and_assert_nib_push_func(0x1234, 0x5671234, |n| n.push_nibbles_front(&0x567.into()));
+        test_and_assert_nib_push_func(
+            Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
+            Nibbles::from_str("0x1000000000000000000000000000000000000000000000000000000000000000")
+                .unwrap(),
+            |n| n.push_nibbles_front(&0x1.into()),
+        );
+    }
+
+    #[test]
+    fn push_nibbles_back_works() {
+        test_and_assert_nib_push_func(Nibbles::default(), 0x1234, |n| {
+            n.push_nibbles_back(&0x1234.into())
+        });
+        test_and_assert_nib_push_func(0x1234, 0x1234567, |n| n.push_nibbles_back(&0x567.into()));
+        test_and_assert_nib_push_func(
+            Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
+            Nibbles::from_str("0x0000000000000000000000000000000000000000000000000000000000000001")
+                .unwrap(),
+            |n| n.push_nibbles_back(&0x1.into()),
+        );
+    }
+
+    fn test_and_assert_nib_push_func<F: Fn(&mut Nibbles), S: Into<Nibbles>, E: Into<Nibbles>>(
+        starting_nibs: S,
+        expected: E,
+        f: F,
+    ) {
+        let mut nibs = starting_nibs.into();
+        (f)(&mut nibs);
+
+        assert_eq!(nibs, expected.into());
+    }
+
     #[test]
     fn get_next_nibbles_works() {
         let n: Nibbles = 0x1234.into();
@@ -1179,7 +1264,7 @@ mod tests {
     fn nibbles_from_h256_works() {
         assert_eq!(
             format!("{:x}", Nibbles::from_h256_be(H256::from_low_u64_be(0))),
-            "0x0000000000000000000000000000000000000000000000000000000000000000"
+            "0x0000000000000000000000000000000000000000000000000000000000000000",
         );
         assert_eq!(
             format!("{:x}", Nibbles::from_h256_be(H256::from_low_u64_be(2048))),

From cf1f147773bcc2f6e736938d7ff3d65735d9ae97 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Wed, 14 Feb 2024 17:21:24 -0700
Subject: [PATCH 11/16] Finished segment to key + cleanup

---
 mpt_trie/src/lib.rs           |   2 +-
 mpt_trie/src/special_query.rs |   3 +-
 mpt_trie/src/utils.rs         | 104 ++++++++++++++++++++++++----------
 3 files changed, 77 insertions(+), 32 deletions(-)

diff --git a/mpt_trie/src/lib.rs b/mpt_trie/src/lib.rs
index 48c75ae3e..85988e4c4 100644
--- a/mpt_trie/src/lib.rs
+++ b/mpt_trie/src/lib.rs
@@ -17,7 +17,7 @@
 
 pub mod nibbles;
 pub mod partial_trie;
-mod special_query;
+pub mod special_query;
 mod trie_hashing;
 pub mod trie_ops;
 pub mod trie_subsets;
diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
index 69558207f..d338cfe7e 100644
--- a/mpt_trie/src/special_query.rs
+++ b/mpt_trie/src/special_query.rs
@@ -7,6 +7,7 @@ use crate::{
     utils::PathSegment,
 };
 
+/// An iterator for a trie query. Note that this iterator is lazy.
 #[derive(Debug)]
 pub struct TriePathIter<N: PartialTrie> {
     /// The next node in the trie to query with the remaining key.
@@ -94,7 +95,7 @@ fn pop_nibbles_clamped(nibbles: &mut Nibbles, n: usize) -> Nibbles {
 /// Note that if the key does not match the entire key of a node (eg. the
 /// remaining key is `0x34` but the next key is a leaf with the key `0x3456`),
 /// then the leaf will not appear in the query output.
-pub fn path_for_query<K, T: PartialTrie>(trie: &Node<T>, k: K) -> impl Iterator<Item = PathSegment>
+pub fn path_for_query<K, T: PartialTrie>(trie: &Node<T>, k: K) -> TriePathIter<T>
 where
     K: Into<Nibbles>,
 {
diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs
index ce0c61af1..2b871493c 100644
--- a/mpt_trie/src/utils.rs
+++ b/mpt_trie/src/utils.rs
@@ -1,6 +1,7 @@
 //! Various types and logic that don't fit well into any other module.
 
 use std::{
+    borrow::Borrow,
     fmt::{self, Display},
     ops::BitAnd,
     sync::Arc,
@@ -12,7 +13,6 @@ use num_traits::PrimInt;
 use crate::{
     nibbles::{Nibble, Nibbles},
     partial_trie::{Node, PartialTrie},
-    special_query::TriePathIter,
 };
 
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
@@ -100,6 +100,30 @@ pub enum PathSegment {
     Leaf(Nibbles),
 }
 
+/// Trait for a type that can be converted into a trie key ([`Nibbles`]).
+pub trait IntoTrieKey {
+    /// Reconstruct the key of the type.
+    fn into_key(self) -> Nibbles;
+}
+
+impl<P: Borrow<PathSegment>, T: Iterator<Item = P>> IntoTrieKey for T {
+    fn into_key(self) -> Nibbles {
+        let mut key = Nibbles::default();
+
+        for seg in self {
+            match seg.borrow() {
+                PathSegment::Empty | PathSegment::Hash => (),
+                PathSegment::Branch(nib) => key.push_nibble_back(*nib),
+                PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => {
+                    key.push_nibbles_back(nibs)
+                }
+            }
+        }
+
+        key
+    }
+}
+
 /// A vector of path segments representing a path in the trie.
 #[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
 pub struct TriePath(pub Vec<PathSegment>);
@@ -123,7 +147,34 @@ impl Display for TriePath {
     }
 }
 
+impl IntoIterator for TriePath {
+    type Item = PathSegment;
+
+    type IntoIter = <Vec<Self::Item> as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl From<Vec<PathSegment>> for TriePath {
+    fn from(v: Vec<PathSegment>) -> Self {
+        Self(v)
+    }
+}
+
+impl FromIterator<PathSegment> for TriePath {
+    fn from_iter<T: IntoIterator<Item = PathSegment>>(iter: T) -> Self {
+        Self(Vec::from_iter(iter))
+    }
+}
+
 impl TriePath {
+    /// Get an iterator of the individual path segments in the [`TriePath`].
+    pub fn iter(&self) -> impl Iterator<Item = &'_ PathSegment> {
+        self.0.iter()
+    }
+
     pub(crate) fn dup_and_append(&self, seg: PathSegment) -> Self {
         let mut duped_vec = self.0.clone();
         duped_vec.push(seg);
@@ -140,35 +191,6 @@ impl TriePath {
     }
 }
 
-/// A trie path that is constructed lazily.
-#[derive(Debug)]
-pub struct TriePathLazy<T: PartialTrie>(pub TriePathIter<T>);
-
-impl<T: PartialTrie> TriePathLazy<T> {
-    /// Extract the key from the trie path.
-    pub fn into_key(self) -> Nibbles {
-        let mut key = Nibbles::default();
-
-        for seg in self.0 {
-            match seg {
-                PathSegment::Empty | PathSegment::Hash => (),
-                PathSegment::Branch(nib) => key.push_nibble_front(nib),
-                PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => {
-                    key.push_nibbles_front(&nibs)
-                }
-            }
-        }
-
-        key
-    }
-}
-
-impl<T: PartialTrie> From<TriePathIter<T>> for TriePathLazy<T> {
-    fn from(v: TriePathIter<T>) -> Self {
-        Self(v)
-    }
-}
-
 impl Display for PathSegment {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
@@ -220,3 +242,25 @@ pub fn get_segment_from_node_and_key_piece<T: PartialTrie>(
         TrieNodeType::Leaf => PathSegment::Leaf(*k_piece),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use super::{IntoTrieKey, PathSegment, TriePath};
+    use crate::nibbles::Nibbles;
+
+    #[test]
+    fn path_from_query_works() {
+        let query_path: TriePath = vec![
+            PathSegment::Branch(1),
+            PathSegment::Branch(2),
+            PathSegment::Extension(0x34.into()),
+            PathSegment::Leaf(0x567.into()),
+        ]
+        .into();
+
+        let reconstructed_key = query_path.iter().into_key();
+        assert_eq!(reconstructed_key, Nibbles::from_str("0x1234567").unwrap());
+    }
+}

From 3030cd4e7b80677b6d5df955db365225bc6a69fe Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Wed, 14 Feb 2024 22:49:44 -0700
Subject: [PATCH 12/16] Fixed a bug in `create_trie_subsets`

- The issue was `reset_tracked_trie_state` never traversing past
  the first node.
---
 mpt_trie/src/trie_subsets.rs | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index 5ae8c4b75..3bdd1d936 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -362,10 +362,14 @@ fn create_partial_trie_subset_from_tracked_trie<N: PartialTrie>(
 
 fn reset_tracked_trie_state<N: PartialTrie>(tracked_node: &mut TrackedNode<N>) {
     match tracked_node.node {
-        TrackedNodeIntern::Branch(ref mut children) => {
-            children.iter_mut().for_each(|c| c.info.reset())
+        TrackedNodeIntern::Branch(ref mut children) => children.iter_mut().for_each(|c| {
+            c.info.reset();
+            reset_tracked_trie_state(c);
+        }),
+        TrackedNodeIntern::Extension(ref mut child) => {
+            child.info.reset();
+            reset_tracked_trie_state(child);
         }
-        TrackedNodeIntern::Extension(ref mut child) => child.info.reset(),
         TrackedNodeIntern::Empty | TrackedNodeIntern::Hash | TrackedNodeIntern::Leaf => {
             tracked_node.info.reset()
         }
@@ -760,7 +764,17 @@ mod tests {
                 .cloned();
 
             assert_nodes_are_leaf_nodes(&partial_trie, leaf_trie_keys);
-            assert_nodes_are_hash_nodes(&partial_trie, keys_of_hash_nodes);
+
+            // We have no idea were the paths to the hashed out nodes will start in the
+            // trie, so the best we can do is to check that they don't exist (if we traverse
+            // over a `Hash` node, we return `None`.)
+            assert_all_keys_do_not_exist(&partial_trie, keys_of_hash_nodes);
+        }
+    }
+
+    fn assert_all_keys_do_not_exist(trie: &TrieType, ks: impl Iterator<Item = Nibbles>) {
+        for k in ks {
+            assert!(trie.get(k).is_none());
         }
     }
 

From 695c9a390b12c90dbd58b5d7191823af1612c02e Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Wed, 14 Feb 2024 22:54:01 -0700
Subject: [PATCH 13/16] Subset tests now start the logger

---
 mpt_trie/src/trie_subsets.rs | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index 3bdd1d936..8573d105c 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -390,8 +390,8 @@ mod tests {
         nibbles::Nibbles,
         partial_trie::{Node, PartialTrie},
         testing_utils::{
-            create_trie_with_large_entry_nodes, generate_n_random_fixed_trie_value_entries,
-            handmade_trie_1, TrieType,
+            common_setup, create_trie_with_large_entry_nodes,
+            generate_n_random_fixed_trie_value_entries, handmade_trie_1, TrieType,
         },
         trie_ops::ValOrHash,
         utils::TrieNodeType,
@@ -497,6 +497,8 @@ mod tests {
 
     #[test]
     fn empty_trie_does_not_return_err_on_query() {
+        common_setup();
+
         let trie = TrieType::default();
         let nibbles: Nibbles = 0x1234.into();
         let res = create_trie_subset(&trie, once(nibbles));
@@ -506,6 +508,8 @@ mod tests {
 
     #[test]
     fn non_existent_key_does_not_return_err() {
+        common_setup();
+
         let mut trie = TrieType::default();
         trie.insert(0x1234, vec![0, 1, 2]);
         let res = create_trie_subset(&trie, once(0x5678));
@@ -515,6 +519,8 @@ mod tests {
 
     #[test]
     fn encountering_a_hash_node_returns_err() {
+        common_setup();
+
         let trie = TrieType::new(Node::Hash(H256::zero()));
         let res = create_trie_subset(&trie, once(0x1234));
 
@@ -523,6 +529,8 @@ mod tests {
 
     #[test]
     fn single_node_trie_is_queryable() {
+        common_setup();
+
         let mut trie = TrieType::default();
         trie.insert(0x1234, vec![0, 1, 2]);
         let trie_subset = create_trie_subset(&trie, once(0x1234)).unwrap();
@@ -532,6 +540,8 @@ mod tests {
 
     #[test]
     fn multi_node_trie_returns_proper_subset() {
+        common_setup();
+
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x56, 0x12345_u64]);
 
         let trie_subset = create_trie_subset(&trie, vec![0x1234, 0x56]).unwrap();
@@ -544,6 +554,8 @@ mod tests {
 
     #[test]
     fn intermediate_nodes_are_included_in_subset() {
+        common_setup();
+
         let (trie, ks_nibbles) = handmade_trie_1();
         let trie_subset_all = create_trie_subset(&trie, ks_nibbles.iter().cloned()).unwrap();
 
@@ -687,6 +699,8 @@ mod tests {
 
     #[test]
     fn all_leafs_of_keys_to_create_subset_are_included_in_subset_for_giant_trie() {
+        common_setup();
+
         let (_, trie_subsets, keys_of_subsets) = create_massive_trie_and_subsets(9009);
 
         for (sub_trie, ks_used) in trie_subsets.into_iter().zip(keys_of_subsets.into_iter()) {
@@ -707,6 +721,8 @@ mod tests {
 
     #[test]
     fn sub_trie_that_includes_branch_but_not_children_hashes_out_children() {
+        common_setup();
+
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x12345, 0x12346, 0x1234f]);
         let partial_trie = create_trie_subset(&trie, [0x1234f]).unwrap();
 
@@ -715,6 +731,8 @@ mod tests {
 
     #[test]
     fn sub_trie_for_non_existent_key_that_hits_branch_leaf_hashes_out_leaf() {
+        common_setup();
+
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x1234589, 0x12346]);
         let partial_trie = create_trie_subset(&trie, [0x1234567]).unwrap();
 
@@ -724,6 +742,7 @@ mod tests {
 
     #[test]
     fn hash_of_branch_partial_tries_matches_original_trie() {
+        common_setup();
         let trie = create_trie_with_large_entry_nodes(&[0x1234, 0x56, 0x12345]);
 
         let base_hash: H256 = trie.hash();
@@ -741,6 +760,8 @@ mod tests {
 
     #[test]
     fn hash_of_giant_random_partial_tries_matches_original_trie() {
+        common_setup();
+
         let (base_trie, trie_subsets, _) = create_massive_trie_and_subsets(9010);
         let base_hash = base_trie.hash();
 
@@ -751,6 +772,8 @@ mod tests {
 
     #[test]
     fn giant_random_partial_tries_hashes_leaves_correctly() {
+        common_setup();
+
         let (base_trie, trie_subsets, leaf_keys_per_trie) = create_massive_trie_and_subsets(9011);
         let all_keys: Vec<Nibbles> = base_trie.keys().collect();
 

From 1ed8454a1560119a4e6ebfbbe6109ed9876e06eb Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Wed, 14 Feb 2024 23:03:51 -0700
Subject: [PATCH 14/16] Renamed `PathSegment` --> `TrieSegment`

---
 mpt_trie/src/debug_tools/query.rs |  4 +-
 mpt_trie/src/special_query.rs     | 56 ++++++++++----------
 mpt_trie/src/utils.rs             | 88 +++++++++++++++----------------
 3 files changed, 73 insertions(+), 75 deletions(-)

diff --git a/mpt_trie/src/debug_tools/query.rs b/mpt_trie/src/debug_tools/query.rs
index d5ad82a17..dd7d3b654 100644
--- a/mpt_trie/src/debug_tools/query.rs
+++ b/mpt_trie/src/debug_tools/query.rs
@@ -9,7 +9,7 @@ use super::common::get_key_piece_from_node_pulling_from_key_for_branches;
 use crate::{
     nibbles::Nibbles,
     partial_trie::{Node, PartialTrie, WrappedNode},
-    utils::{get_segment_from_node_and_key_piece, PathSegment, TriePath},
+    utils::{get_segment_from_node_and_key_piece, TriePath, TrieSegment},
 };
 
 /// Params controlling how much information is reported in the query output.
@@ -219,7 +219,7 @@ impl DebugQueryOutput {
     // TODO: Make the output easier to read...
     fn fmt_node_based_on_debug_params(
         f: &mut fmt::Formatter<'_>,
-        seg: &PathSegment,
+        seg: &TrieSegment,
         extra_seg_info: &Option<ExtraNodeSegmentInfo>,
         params: &DebugQueryParams,
     ) -> fmt::Result {
diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
index d338cfe7e..329726772 100644
--- a/mpt_trie/src/special_query.rs
+++ b/mpt_trie/src/special_query.rs
@@ -4,7 +4,7 @@
 use crate::{
     nibbles::Nibbles,
     partial_trie::{Node, PartialTrie, WrappedNode},
-    utils::PathSegment,
+    utils::TrieSegment,
 };
 
 /// An iterator for a trie query. Note that this iterator is lazy.
@@ -22,7 +22,7 @@ pub struct TriePathIter<N: PartialTrie> {
 }
 
 impl<T: PartialTrie> Iterator for TriePathIter<T> {
-    type Item = PathSegment;
+    type Item = TrieSegment;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.terminated {
@@ -32,11 +32,11 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
         match self.curr_node.as_ref() {
             Node::Empty => {
                 self.terminated = true;
-                Some(PathSegment::Empty)
+                Some(TrieSegment::Empty)
             }
             Node::Hash(_) => {
                 self.terminated = true;
-                Some(PathSegment::Hash)
+                Some(TrieSegment::Hash)
             }
             Node::Branch { children, .. } => {
                 // Our query key has ended. Stop here.
@@ -48,7 +48,7 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
                 let nib = self.curr_key.pop_next_nibble_front();
                 self.curr_node = children[nib as usize].clone();
 
-                Some(PathSegment::Branch(nib))
+                Some(TrieSegment::Branch(nib))
             }
             Node::Extension { nibbles, child } => {
                 match self
@@ -62,7 +62,7 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
                     }
                     true => {
                         pop_nibbles_clamped(&mut self.curr_key, nibbles.count);
-                        let res = Some(PathSegment::Extension(*nibbles));
+                        let res = Some(TrieSegment::Extension(*nibbles));
                         self.curr_node = child.clone();
 
                         res
@@ -74,7 +74,7 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
 
                 match self.curr_key == *nibbles {
                     false => None,
-                    true => Some(PathSegment::Leaf(*nibbles)),
+                    true => Some(TrieSegment::Leaf(*nibbles)),
                 }
             }
         }
@@ -88,8 +88,6 @@ fn pop_nibbles_clamped(nibbles: &mut Nibbles, n: usize) -> Nibbles {
     nibbles.pop_nibbles_front(n_nibs_to_pop)
 }
 
-// TODO: Move to a blanket impl...
-// TODO: Could make this return an `Iterator` with some work...
 /// Returns all nodes in the trie that are traversed given a query (key).
 ///
 /// Note that if the key does not match the entire key of a node (eg. the
@@ -111,7 +109,7 @@ mod test {
     use std::str::FromStr;
 
     use super::path_for_query;
-    use crate::{nibbles::Nibbles, testing_utils::handmade_trie_1, utils::PathSegment};
+    use crate::{nibbles::Nibbles, testing_utils::handmade_trie_1, utils::TrieSegment};
 
     #[test]
     fn query_iter_works() {
@@ -120,33 +118,33 @@ mod test {
         // ks --> vec![0x1234, 0x1324, 0x132400005_u64, 0x2001, 0x2002];
         let res = vec![
             vec![
-                PathSegment::Branch(1),
-                PathSegment::Branch(2),
-                PathSegment::Leaf(0x34.into()),
+                TrieSegment::Branch(1),
+                TrieSegment::Branch(2),
+                TrieSegment::Leaf(0x34.into()),
             ],
             vec![
-                PathSegment::Branch(1),
-                PathSegment::Branch(3),
-                PathSegment::Extension(0x24.into()),
+                TrieSegment::Branch(1),
+                TrieSegment::Branch(3),
+                TrieSegment::Extension(0x24.into()),
             ],
             vec![
-                PathSegment::Branch(1),
-                PathSegment::Branch(3),
-                PathSegment::Extension(0x24.into()),
-                PathSegment::Branch(0),
-                PathSegment::Leaf(Nibbles::from_str("0x0005").unwrap()),
+                TrieSegment::Branch(1),
+                TrieSegment::Branch(3),
+                TrieSegment::Extension(0x24.into()),
+                TrieSegment::Branch(0),
+                TrieSegment::Leaf(Nibbles::from_str("0x0005").unwrap()),
             ],
             vec![
-                PathSegment::Branch(2),
-                PathSegment::Extension(Nibbles::from_str("0x00").unwrap()),
-                PathSegment::Branch(0x1),
-                PathSegment::Leaf(Nibbles::default()),
+                TrieSegment::Branch(2),
+                TrieSegment::Extension(Nibbles::from_str("0x00").unwrap()),
+                TrieSegment::Branch(0x1),
+                TrieSegment::Leaf(Nibbles::default()),
             ],
             vec![
-                PathSegment::Branch(2),
-                PathSegment::Extension(Nibbles::from_str("0x00").unwrap()),
-                PathSegment::Branch(0x2),
-                PathSegment::Leaf(Nibbles::default()),
+                TrieSegment::Branch(2),
+                TrieSegment::Extension(Nibbles::from_str("0x00").unwrap()),
+                TrieSegment::Branch(0x2),
+                TrieSegment::Leaf(Nibbles::default()),
             ],
         ];
 
diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs
index 2b871493c..59400ee82 100644
--- a/mpt_trie/src/utils.rs
+++ b/mpt_trie/src/utils.rs
@@ -83,7 +83,7 @@ pub(crate) fn bytes_to_h256(b: &[u8; 32]) -> H256 {
 /// the key piece of the node if applicable (eg. [`Node::Empty`] &
 /// [`Node::Hash`] do not have associated key pieces).
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
-pub enum PathSegment {
+pub enum TrieSegment {
     /// Empty node.
     Empty,
 
@@ -106,15 +106,15 @@ pub trait IntoTrieKey {
     fn into_key(self) -> Nibbles;
 }
 
-impl<P: Borrow<PathSegment>, T: Iterator<Item = P>> IntoTrieKey for T {
+impl<P: Borrow<TrieSegment>, T: Iterator<Item = P>> IntoTrieKey for T {
     fn into_key(self) -> Nibbles {
         let mut key = Nibbles::default();
 
         for seg in self {
             match seg.borrow() {
-                PathSegment::Empty | PathSegment::Hash => (),
-                PathSegment::Branch(nib) => key.push_nibble_back(*nib),
-                PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => {
+                TrieSegment::Empty | TrieSegment::Hash => (),
+                TrieSegment::Branch(nib) => key.push_nibble_back(*nib),
+                TrieSegment::Extension(nibs) | TrieSegment::Leaf(nibs) => {
                     key.push_nibbles_back(nibs)
                 }
             }
@@ -126,7 +126,7 @@ impl<P: Borrow<PathSegment>, T: Iterator<Item = P>> IntoTrieKey for T {
 
 /// A vector of path segments representing a path in the trie.
 #[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
-pub struct TriePath(pub Vec<PathSegment>);
+pub struct TriePath(pub Vec<TrieSegment>);
 
 impl Display for TriePath {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -148,7 +148,7 @@ impl Display for TriePath {
 }
 
 impl IntoIterator for TriePath {
-    type Item = PathSegment;
+    type Item = TrieSegment;
 
     type IntoIter = <Vec<Self::Item> as IntoIterator>::IntoIter;
 
@@ -157,89 +157,89 @@ impl IntoIterator for TriePath {
     }
 }
 
-impl From<Vec<PathSegment>> for TriePath {
-    fn from(v: Vec<PathSegment>) -> Self {
+impl From<Vec<TrieSegment>> for TriePath {
+    fn from(v: Vec<TrieSegment>) -> Self {
         Self(v)
     }
 }
 
-impl FromIterator<PathSegment> for TriePath {
-    fn from_iter<T: IntoIterator<Item = PathSegment>>(iter: T) -> Self {
+impl FromIterator<TrieSegment> for TriePath {
+    fn from_iter<T: IntoIterator<Item = TrieSegment>>(iter: T) -> Self {
         Self(Vec::from_iter(iter))
     }
 }
 
 impl TriePath {
     /// Get an iterator of the individual path segments in the [`TriePath`].
-    pub fn iter(&self) -> impl Iterator<Item = &'_ PathSegment> {
+    pub fn iter(&self) -> impl Iterator<Item = &'_ TrieSegment> {
         self.0.iter()
     }
 
-    pub(crate) fn dup_and_append(&self, seg: PathSegment) -> Self {
+    pub(crate) fn dup_and_append(&self, seg: TrieSegment) -> Self {
         let mut duped_vec = self.0.clone();
         duped_vec.push(seg);
 
         Self(duped_vec)
     }
 
-    pub(crate) fn append(&mut self, seg: PathSegment) {
+    pub(crate) fn append(&mut self, seg: TrieSegment) {
         self.0.push(seg);
     }
 
-    fn write_elem(f: &mut fmt::Formatter<'_>, seg: &PathSegment) -> fmt::Result {
+    fn write_elem(f: &mut fmt::Formatter<'_>, seg: &TrieSegment) -> fmt::Result {
         write!(f, "{}", seg)
     }
 }
 
-impl Display for PathSegment {
+impl Display for TrieSegment {
     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),
+            TrieSegment::Empty => write!(f, "Empty"),
+            TrieSegment::Hash => write!(f, "Hash"),
+            TrieSegment::Branch(nib) => write!(f, "Branch({})", nib),
+            TrieSegment::Extension(nibs) => write!(f, "Extension({})", nibs),
+            TrieSegment::Leaf(nibs) => write!(f, "Leaf({})", nibs),
         }
     }
 }
 
-impl PathSegment {
-    /// Get the node type of the [`PathSegment`].
+impl TrieSegment {
+    /// Get the node type of the [`TrieSegment`].
     pub 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,
+            TrieSegment::Empty => TrieNodeType::Empty,
+            TrieSegment::Hash => TrieNodeType::Hash,
+            TrieSegment::Branch(_) => TrieNodeType::Branch,
+            TrieSegment::Extension(_) => TrieNodeType::Extension,
+            TrieSegment::Leaf(_) => TrieNodeType::Leaf,
         }
     }
 
     /// Extracts the key piece used by the segment (if applicable).
     pub fn get_key_piece_from_seg_if_present(&self) -> Option<Nibbles> {
         match self {
-            PathSegment::Empty | PathSegment::Hash => None,
-            PathSegment::Branch(nib) => Some(Nibbles::from_nibble(*nib)),
-            PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => Some(*nibs),
+            TrieSegment::Empty | TrieSegment::Hash => None,
+            TrieSegment::Branch(nib) => Some(Nibbles::from_nibble(*nib)),
+            TrieSegment::Extension(nibs) | TrieSegment::Leaf(nibs) => Some(*nibs),
         }
     }
 }
 
-/// Creates a [`PathSegment`] given a node and a key we are querying.
+/// Creates a [`TrieSegment`] given a node and a key we are querying.
 ///
 /// This function is intended to be used during a trie query as we are
 /// traversing down a trie. Depending on the current node, we pop off nibbles
-/// and use these to create `PathSegment`s.
+/// and use these to create `TrieSegment`s.
 pub fn get_segment_from_node_and_key_piece<T: PartialTrie>(
     n: &Node<T>,
     k_piece: &Nibbles,
-) -> PathSegment {
+) -> TrieSegment {
     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),
+        TrieNodeType::Empty => TrieSegment::Empty,
+        TrieNodeType::Hash => TrieSegment::Hash,
+        TrieNodeType::Branch => TrieSegment::Branch(k_piece.get_nibble(0)),
+        TrieNodeType::Extension => TrieSegment::Extension(*k_piece),
+        TrieNodeType::Leaf => TrieSegment::Leaf(*k_piece),
     }
 }
 
@@ -247,16 +247,16 @@ pub fn get_segment_from_node_and_key_piece<T: PartialTrie>(
 mod tests {
     use std::str::FromStr;
 
-    use super::{IntoTrieKey, PathSegment, TriePath};
+    use super::{IntoTrieKey, TriePath, TrieSegment};
     use crate::nibbles::Nibbles;
 
     #[test]
     fn path_from_query_works() {
         let query_path: TriePath = vec![
-            PathSegment::Branch(1),
-            PathSegment::Branch(2),
-            PathSegment::Extension(0x34.into()),
-            PathSegment::Leaf(0x567.into()),
+            TrieSegment::Branch(1),
+            TrieSegment::Branch(2),
+            TrieSegment::Extension(0x34.into()),
+            TrieSegment::Leaf(0x567.into()),
         ]
         .into();
 

From 0a8c16d41942cd63e990ed6aebaa63e35327f6f5 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Thu, 15 Feb 2024 13:41:15 -0700
Subject: [PATCH 15/16] Fixed two clippy errors

---
 mpt_trie/src/trie_subsets.rs | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
index 8573d105c..974b92afc 100644
--- a/mpt_trie/src/trie_subsets.rs
+++ b/mpt_trie/src/trie_subsets.rs
@@ -420,17 +420,6 @@ mod tests {
                 nibbles: nibbles.into(),
             }
         }
-
-        fn with_key(&self, k: &Nibbles) -> bool {
-            self.nibbles.reverse() == *k
-        }
-    }
-
-    fn get_node_full_nibbles_from_list<'a>(
-        nodes: &'a [NodeFullNibbles],
-        k: &'a Nibbles,
-    ) -> Option<&'a NodeFullNibbles> {
-        nodes.iter().find(|x| x.with_key(k))
     }
 
     fn get_all_nodes_in_trie(trie: &TrieType) -> Vec<NodeFullNibbles> {

From 3bd89a171f3056aa571992fe9f57e4b693e0c899 Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Tue, 20 Feb 2024 14:24:10 -1000
Subject: [PATCH 16/16] Apply suggestions from code review

Nashtare's suggested changes for #39.

Co-authored-by: Robin Salen <30937548+Nashtare@users.noreply.github.com>
---
 mpt_trie/src/special_query.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
index 329726772..503331aa0 100644
--- a/mpt_trie/src/special_query.rs
+++ b/mpt_trie/src/special_query.rs
@@ -81,8 +81,8 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
     }
 }
 
-/// Attempt to pop `n` nibbles from the given [`Nibbles`] and "clamp" the
-/// nibbles popped by not popping more nibbles than exist.
+/// Attempts to pop `n` nibbles from the given [`Nibbles`] and "clamp" the
+/// nibbles popped by not popping more nibbles than there are.
 fn pop_nibbles_clamped(nibbles: &mut Nibbles, n: usize) -> Nibbles {
     let n_nibs_to_pop = nibbles.count.min(n);
     nibbles.pop_nibbles_front(n_nibs_to_pop)