From 12d3fbefffa54e3b157e8f60b658caa6959a5e35 Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:49:07 +0000
Subject: [PATCH] fix(trie): check branch node masks if `store_in_db_trie` is
 `None` (#13828)

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

diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index 1c5453cb6330..4a4f0f185ceb 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -694,7 +694,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();
@@ -711,7 +714,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,
@@ -734,7 +737,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
                     }
@@ -769,17 +772,19 @@ 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
+                                    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);
                                 }
@@ -882,7 +887,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));
@@ -1227,12 +1235,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>,
     },
 }
 
@@ -1245,12 +1253,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,
         }
     }
 }