From e6b4ab854a16a7dcc4c350fece504f600e5bca1c Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Thu, 16 Jan 2025 18:38:01 +0000
Subject: [PATCH 1/2] fix(trie): check cached branch node masks if store in db
 trie is none

---
 crates/trie/sparse/src/trie.rs | 41 ++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index 80b0e31985ca..12bda425fb98 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -680,7 +680,10 @@ impl<P> RevealedSparseTrie<P> {
                     if let Some((hash, store_in_db_trie)) =
                         hash.zip(*store_in_db_trie).filter(|_| !prefix_set_contains(&path))
                     {
-                        (RlpNode::word_rlp(&hash), SparseNodeType::Extension { store_in_db_trie })
+                        (
+                            RlpNode::word_rlp(&hash),
+                            SparseNodeType::Extension { store_in_db_trie: Some(store_in_db_trie) },
+                        )
                     } else if buffers.rlp_node_stack.last().is_some_and(|e| e.0 == child_path) {
                         let (_, child, child_node_type) = buffers.rlp_node_stack.pop().unwrap();
                         self.rlp_buf.clear();
@@ -697,7 +700,7 @@ impl<P> RevealedSparseTrie<P> {
                             "Extension node"
                         );
 
-                        *store_in_db_trie = Some(store_in_db_trie_value);
+                        *store_in_db_trie = store_in_db_trie_value;
 
                         (
                             rlp_node,
@@ -720,7 +723,7 @@ impl<P> RevealedSparseTrie<P> {
                         buffers.rlp_node_stack.push((
                             path,
                             RlpNode::word_rlp(&hash),
-                            SparseNodeType::Branch { store_in_db_trie },
+                            SparseNodeType::Branch { store_in_db_trie: Some(store_in_db_trie) },
                         ));
                         continue
                     }
@@ -755,17 +758,18 @@ impl<P> RevealedSparseTrie<P> {
                                 let last_child_nibble = child_path.last().unwrap();
 
                                 // Determine whether we need to set trie mask bit.
-                                let should_set_tree_mask_bit =
-                                    // A blinded node has the tree mask bit set
-                                    (
-                                        child_node_type.is_hash() &&
-                                        self.branch_node_tree_masks
-                                            .get(&path)
-                                            .is_some_and(|mask| mask.is_bit_set(last_child_nibble))
-                                    ) ||
+                                let should_set_tree_mask_bit = if let Some(store_in_db_trie) =
+                                    child_node_type.store_in_db_trie()
+                                {
                                     // A branch or an extension node explicitly set the
                                     // `store_in_db_trie` flag
-                                    child_node_type.store_in_db_trie();
+                                    store_in_db_trie
+                                } else {
+                                    // A blinded node has the tree mask bit set
+                                    self.branch_node_tree_masks
+                                        .get(&path)
+                                        .is_some_and(|mask| mask.is_bit_set(last_child_nibble))
+                                };
                                 if should_set_tree_mask_bit {
                                     tree_mask.set_bit(last_child_nibble);
                                 }
@@ -868,7 +872,10 @@ impl<P> RevealedSparseTrie<P> {
                     };
                     *store_in_db_trie = Some(store_in_db_trie_value);
 
-                    (rlp_node, SparseNodeType::Branch { store_in_db_trie: store_in_db_trie_value })
+                    (
+                        rlp_node,
+                        SparseNodeType::Branch { store_in_db_trie: Some(store_in_db_trie_value) },
+                    )
                 }
             };
             buffers.rlp_node_stack.push((path, rlp_node, node_type));
@@ -1191,12 +1198,12 @@ enum SparseNodeType {
     /// Sparse extension node.
     Extension {
         /// A flag indicating whether the extension node should be stored in the database.
-        store_in_db_trie: bool,
+        store_in_db_trie: Option<bool>,
     },
     /// Sparse branch node.
     Branch {
         /// A flag indicating whether the branch node should be stored in the database.
-        store_in_db_trie: bool,
+        store_in_db_trie: Option<bool>,
     },
 }
 
@@ -1209,12 +1216,12 @@ impl SparseNodeType {
         matches!(self, Self::Branch { .. })
     }
 
-    const fn store_in_db_trie(&self) -> bool {
+    const fn store_in_db_trie(&self) -> Option<bool> {
         match *self {
             Self::Extension { store_in_db_trie } | Self::Branch { store_in_db_trie } => {
                 store_in_db_trie
             }
-            _ => false,
+            _ => None,
         }
     }
 }

From f36bb0b7ad6fecb3bdcb828c53b666a1f0569471 Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Thu, 16 Jan 2025 23:48:57 +0000
Subject: [PATCH 2/2] check that blinded node is hash

---
 crates/trie/sparse/src/trie.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index 12bda425fb98..e25529f9b315 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -766,9 +766,10 @@ impl<P> RevealedSparseTrie<P> {
                                     store_in_db_trie
                                 } else {
                                     // A blinded node has the tree mask bit set
-                                    self.branch_node_tree_masks
-                                        .get(&path)
-                                        .is_some_and(|mask| mask.is_bit_set(last_child_nibble))
+                                    child_node_type.is_hash() &&
+                                        self.branch_node_tree_masks.get(&path).is_some_and(
+                                            |mask| mask.is_bit_set(last_child_nibble),
+                                        )
                                 };
                                 if should_set_tree_mask_bit {
                                     tree_mask.set_bit(last_child_nibble);