Now prunes tries properly (and more aggressively!) without hack (0xPo…

…lygonZero#97) * Finished logic for detecting trie collapses * A few fixes related to branch collapse detection * Requested PR changes for 0xPolygonZero#39 * Extension nodes are now always not hashed if encountered - However it will only not hash the child if the key matches the key piece of the extension. * Mix of different cherry-picked fixes - Fixes overall deal with calling the new pruning query in `mpt_trie` and changing the order of calls when generating proof IR. * Impled additional logic to extract remaining child - Realized that we needed a bit more logic to get the key to the remaining child after a branch collapse. - Ended up adding a bool to `TriePathIter`. * Cleanup from messy cherry picks - Also includes one core fix that got left out. * Apply suggestions from code review Robin's suggested changes for PR 0xPolygonZero#97 Co-authored-by: Robin Salen <[email protected]> * Requsted PR changes for 0xPolygonZero#97 --------- Co-authored-by: Ben <[email protected]> Co-authored-by: Robin Salen <[email protected]>
topos-protocol · Mar 11, 2024 · 5c8d516 · 5c8d516
1 parent 05f8722
commit 5c8d516
Show file tree

Hide file tree

Showing 9 changed files with 331 additions and 118 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix generation inputs logging pre-transaction execution ([#89](https://github.com/0xPolygonZero/zk_evm/pull/89))
 - Reduce state trie size for dummy payloads ([#88](https://github.com/0xPolygonZero/zk_evm/pull/88))
 - Fix post-txn trie debugging output for multi-logs receipts ([#86](https://github.com/0xPolygonZero/zk_evm/pull/86))
+- Fixed *most* failing blocks caused by the merged in aggressive pruning changes ([#97](https://github.com/0xPolygonZero/zk_evm/pull/97))
 
 ## [0.1.1] - 2024-03-01
 

diff --git a/mpt_trie/src/debug_tools/query.rs b/mpt_trie/src/debug_tools/query.rs
@@ -159,7 +159,7 @@ pub struct DebugQueryOutput {
     k: Nibbles,
 
     /// The nodes hit during the query.
-    pub node_path: TriePath,
+    node_path: TriePath,
     extra_node_info: Vec<Option<ExtraNodeSegmentInfo>>,
     node_found: bool,
     params: DebugQueryParams,

diff --git a/mpt_trie/src/nibbles.rs b/mpt_trie/src/nibbles.rs
@@ -26,6 +26,11 @@ pub type Nibble = u8;
 /// Used for the internal representation of a sequence of nibbles.
 pub type NibblesIntern = U512;
 
+const MULTIPLE_NIBBLES_APPEND_ASSERT_ERR_MSG: &str =
+    "Attempted to create a nibbles sequence longer than 64!";
+const SINGLE_NIBBLE_APPEND_ASSERT_ERR_MSG: &str =
+    "Attempted to append a single nibble that was greater than 15!";
+
 /// Because there are two different ways to convert to `Nibbles`, we don't want
 /// to rely on `From`. Instead, we'll define a new trait that defines both
 /// conversions.
@@ -776,12 +781,20 @@ impl Nibbles {
     }
 
     fn nibble_append_safety_asserts(&self, n: Nibble) {
-        assert!(self.count < 64);
-        assert!(n < 16);
+        assert!(
+            self.count < 64,
+            "{}",
+            MULTIPLE_NIBBLES_APPEND_ASSERT_ERR_MSG
+        );
+        assert!(n < 16, "{}", SINGLE_NIBBLE_APPEND_ASSERT_ERR_MSG);
     }
 
     fn nibbles_append_safety_asserts(&self, new_count: usize) {
-        assert!(new_count <= 64);
+        assert!(
+            new_count <= 64,
+            "{}",
+            MULTIPLE_NIBBLES_APPEND_ASSERT_ERR_MSG
+        );
     }
 
     // TODO: REMOVE BEFORE NEXT CRATE VERSION! THIS IS A TEMP HACK!
@@ -806,8 +819,12 @@ mod tests {
     use super::{Nibble, Nibbles, ToNibbles};
     use crate::nibbles::FromHexPrefixError;
 
-    const LONG_ZERO_NIBS_STR_LEN_63: &str =
-        "0x000000000000000000000000000000000000000000000000000000000000000";
+    const ZERO_NIBS_63: &str = "0x000000000000000000000000000000000000000000000000000000000000000";
+    const ZERO_NIBS_64: &str = "0x0000000000000000000000000000000000000000000000000000000000000000";
+    const ZERO_NIBS_64_LEADING_1: &str =
+        "0x1000000000000000000000000000000000000000000000000000000000000000";
+    const ZERO_NIBS_64_TRAILING_1: &str =
+        "0x0000000000000000000000000000000000000000000000000000000000000001";
 
     #[test]
     fn get_nibble_works() {
@@ -925,9 +942,8 @@ mod tests {
         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(),
+            Nibbles::from_str(ZERO_NIBS_63).unwrap(),
+            Nibbles::from_str(ZERO_NIBS_64_LEADING_1).unwrap(),
             |n| n.push_nibble_front(0x1),
         );
     }
@@ -937,9 +953,8 @@ mod tests {
         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(),
+            Nibbles::from_str(ZERO_NIBS_63).unwrap(),
+            Nibbles::from_str(ZERO_NIBS_64_TRAILING_1).unwrap(),
             |n| n.push_nibble_back(0x1),
         );
     }
@@ -951,9 +966,8 @@ mod tests {
         });
         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(),
+            Nibbles::from_str(ZERO_NIBS_63).unwrap(),
+            Nibbles::from_str(ZERO_NIBS_64_LEADING_1).unwrap(),
             |n| n.push_nibbles_front(&0x1.into()),
         );
     }
@@ -965,9 +979,8 @@ mod tests {
         });
         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(),
+            Nibbles::from_str(ZERO_NIBS_63).unwrap(),
+            Nibbles::from_str(ZERO_NIBS_64_TRAILING_1).unwrap(),
             |n| n.push_nibbles_back(&0x1.into()),
         );
     }
@@ -1264,15 +1277,15 @@ mod tests {
     fn nibbles_from_h256_works() {
         assert_eq!(
             format!("{:x}", Nibbles::from_h256_be(H256::from_low_u64_be(0))),
-            "0x0000000000000000000000000000000000000000000000000000000000000000",
+            ZERO_NIBS_64,
         );
         assert_eq!(
             format!("{:x}", Nibbles::from_h256_be(H256::from_low_u64_be(2048))),
             "0x0000000000000000000000000000000000000000000000000000000000000800"
         );
         assert_eq!(
             format!("{:x}", Nibbles::from_h256_le(H256::from_low_u64_be(0))),
-            "0x0000000000000000000000000000000000000000000000000000000000000000"
+            ZERO_NIBS_64
         );
 
         // Note that the first bit of the `Nibbles` changes if the count is odd.

diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs
@@ -19,6 +19,10 @@ pub struct TriePathIter<N: PartialTrie> {
     // Although wrapping `curr_node` in an option might be more "Rust like", the logic is a lot
     // cleaner with a bool.
     terminated: bool,
+
+    /// Always include the final node we encounter even if the key does not
+    /// match.
+    always_include_final_node_if_possible: bool,
 }
 
 impl<T: PartialTrie> Iterator for TriePathIter<T> {
@@ -39,9 +43,12 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
                 Some(TrieSegment::Hash)
             }
             Node::Branch { children, .. } => {
-                // Our query key has ended. Stop here.
                 if self.curr_key.is_empty() {
+                    // Our query key has ended. Stop here.
                     self.terminated = true;
+
+                    // In this case even if `always_include_final_node` is set, we have no
+                    // information on which branch to take, so we can't add in the last node.
                     return None;
                 }
 
@@ -58,7 +65,8 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
                     false => {
                         // Only a partial match. Stop.
                         self.terminated = true;
-                        None
+                        self.always_include_final_node_if_possible
+                            .then_some(TrieSegment::Extension(*nibbles))
                     }
                     true => {
                         pop_nibbles_clamped(&mut self.curr_key, nibbles.count);
@@ -72,7 +80,7 @@ impl<T: PartialTrie> Iterator for TriePathIter<T> {
             Node::Leaf { nibbles, .. } => {
                 self.terminated = true;
 
-                match self.curr_key == *nibbles {
+                match self.curr_key == *nibbles || self.always_include_final_node_if_possible {
                     false => None,
                     true => Some(TrieSegment::Leaf(*nibbles)),
                 }
@@ -93,14 +101,19 @@ 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) -> TriePathIter<T>
+pub fn path_for_query<K, T: PartialTrie>(
+    trie: &Node<T>,
+    k: K,
+    always_include_final_node_if_possible: bool,
+) -> TriePathIter<T>
 where
     K: Into<Nibbles>,
 {
     TriePathIter {
         curr_node: trie.clone().into(),
         curr_key: k.into(),
         terminated: false,
+        always_include_final_node_if_possible,
     }
 }
 
@@ -109,10 +122,15 @@ mod test {
     use std::str::FromStr;
 
     use super::path_for_query;
-    use crate::{nibbles::Nibbles, testing_utils::handmade_trie_1, utils::TrieSegment};
+    use crate::{
+        nibbles::Nibbles,
+        testing_utils::{common_setup, handmade_trie_1},
+        utils::TrieSegment,
+    };
 
     #[test]
-    fn query_iter_works() {
+    fn query_iter_works_no_last_node() {
+        common_setup();
         let (trie, ks) = handmade_trie_1();
 
         // ks --> vec![0x1234, 0x1324, 0x132400005_u64, 0x2001, 0x2002];
@@ -149,8 +167,50 @@ mod test {
         ];
 
         for (q, expected) in ks.into_iter().zip(res.into_iter()) {
-            let res: Vec<_> = path_for_query(&trie.node, q).collect();
+            let res: Vec<_> = path_for_query(&trie.node, q, false).collect();
             assert_eq!(res, expected)
         }
     }
+
+    #[test]
+    fn query_iter_works_with_last_node() {
+        common_setup();
+        let (trie, _) = handmade_trie_1();
+
+        let extension_expected = vec![
+            TrieSegment::Branch(1),
+            TrieSegment::Branch(3),
+            TrieSegment::Extension(0x24.into()),
+        ];
+
+        assert_eq!(
+            path_for_query(&trie, 0x13, true).collect::<Vec<_>>(),
+            extension_expected
+        );
+        assert_eq!(
+            path_for_query(&trie, 0x132, true).collect::<Vec<_>>(),
+            extension_expected
+        );
+
+        // The last node is a branch here, but we don't include it because a TrieSegment
+        // requires us to state which nibble we took in the branch, and we don't have
+        // this information.
+        assert_eq!(
+            path_for_query(&trie, 0x1324, true).collect::<Vec<_>>(),
+            extension_expected
+        );
+
+        let mut leaf_expected = extension_expected.clone();
+        leaf_expected.push(TrieSegment::Branch(0));
+        leaf_expected.push(TrieSegment::Leaf(Nibbles::from_str("0x0005").unwrap()));
+
+        assert_eq!(
+            path_for_query(&trie, 0x13240, true).collect::<Vec<_>>(),
+            leaf_expected
+        );
+        assert_eq!(
+            path_for_query(&trie, 0x132400, true).collect::<Vec<_>>(),
+            leaf_expected
+        );
+    }
 }
diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs
@@ -307,11 +307,17 @@ fn mark_nodes_that_are_needed<N: PartialTrie>(
             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);
+            // If we hit an extension node, we always must include it unless we exhausted
+            // our key.
+            if curr_nibbles.is_empty() {
+                return Ok(());
+            }
 
-            if r.nibbles_are_identical_up_to_smallest_count(nibbles) {
-                trie.info.touched = true;
+            trie.info.touched = true;
+
+            let nibbles: &Nibbles = trie.info.get_nibbles_expected();
+            if curr_nibbles.nibbles_are_identical_up_to_smallest_count(nibbles) {
+                curr_nibbles.pop_nibbles_front(nibbles.count);
                 return mark_nodes_that_are_needed(child, curr_nibbles);
             }
         }

diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs
@@ -216,7 +216,7 @@ impl TrieSegment {
     }
 
     /// Extracts the key piece used by the segment (if applicable).
-    pub fn get_key_piece_from_seg_if_present(&self) -> Option<Nibbles> {
+    pub(crate) fn get_key_piece_from_seg_if_present(&self) -> Option<Nibbles> {
         match self {
             TrieSegment::Empty | TrieSegment::Hash => None,
             TrieSegment::Branch(nib) => Some(Nibbles::from_nibble(*nib)),
@@ -230,7 +230,7 @@ impl TrieSegment {
 /// 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 `TrieSegment`s.
-pub fn get_segment_from_node_and_key_piece<T: PartialTrie>(
+pub(crate) fn get_segment_from_node_and_key_piece<T: PartialTrie>(
     n: &Node<T>,
     k_piece: &Nibbles,
 ) -> TrieSegment {