From 1970a4425bcd0acc830f9311084db028fd0d7ec3 Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Wed, 12 Feb 2025 15:27:07 +0000
Subject: [PATCH] fix(trie): reveal blinded sparse trie when calculating root
 (#14449)

---
 crates/engine/tree/src/tree/root.rs | 20 +++++---
 crates/trie/common/src/proofs.rs    |  8 ++++
 crates/trie/sparse/src/state.rs     | 71 ++++++++++++++++++++++-------
 crates/trie/sparse/src/trie.rs      |  5 --
 4 files changed, 76 insertions(+), 28 deletions(-)

diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
index 2e1a948bd..60fe5fe3c 100644
--- a/crates/engine/tree/src/tree/root.rs
+++ b/crates/engine/tree/src/tree/root.rs
@@ -74,6 +74,11 @@ pub struct SparseTrieUpdate {
 }
 
 impl SparseTrieUpdate {
+    /// Returns true if the update is empty.
+    pub fn is_empty(&self) -> bool {
+        self.state.is_empty() && self.multiproof.is_empty()
+    }
+
     /// Construct update from multiproof.
     pub fn from_multiproof(multiproof: MultiProof) -> Self {
         Self { multiproof, ..Default::default() }
@@ -673,15 +678,15 @@ where
     ) -> Option<SparseTrieUpdate> {
         let ready_proofs = self.proof_sequencer.add_proof(sequence_number, update);
 
-        if ready_proofs.is_empty() {
-            None
-        } else {
+        ready_proofs
+            .into_iter()
             // Merge all ready proofs and state updates
-            ready_proofs.into_iter().reduce(|mut acc_update, update| {
+            .reduce(|mut acc_update, update| {
                 acc_update.extend(update);
                 acc_update
             })
-        }
+            // Return None if the resulting proof is empty
+            .filter(|proof| !proof.is_empty())
     }
 
     /// Starts the main loop that handles all incoming messages, fetches proofs, applies them to the
@@ -1017,8 +1022,9 @@ where
     debug!(target: "engine::root", num_iterations, "All proofs processed, ending calculation");
 
     let start = Instant::now();
-    let (state_root, trie_updates) =
-        trie.root_with_updates().expect("sparse trie should be revealed");
+    let (state_root, trie_updates) = trie.root_with_updates().map_err(|e| {
+        ParallelStateRootError::Other(format!("could not calculate state root: {e:?}"))
+    })?;
     let elapsed = start.elapsed();
     metrics.sparse_trie_final_update_duration_histogram.record(elapsed);
 
diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs
index 62d76debd..1f37e639d 100644
--- a/crates/trie/common/src/proofs.rs
+++ b/crates/trie/common/src/proofs.rs
@@ -36,6 +36,14 @@ pub struct MultiProof {
 }
 
 impl MultiProof {
+    /// Returns true if the multiproof is empty.
+    pub fn is_empty(&self) -> bool {
+        self.account_subtree.is_empty() &&
+            self.branch_node_hash_masks.is_empty() &&
+            self.branch_node_tree_masks.is_empty() &&
+            self.storages.is_empty()
+    }
+
     /// Return the account proof nodes for the given account path.
     pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, Bytes)> {
         self.account_subtree.matching_nodes_sorted(path)
diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs
index ff2d4c553..cc149b8e4 100644
--- a/crates/trie/sparse/src/state.rs
+++ b/crates/trie/sparse/src/state.rs
@@ -1,5 +1,5 @@
 use crate::{
-    blinded::{BlindedProviderFactory, DefaultBlindedProviderFactory},
+    blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory},
     RevealedSparseTrie, SparseTrie,
 };
 use alloy_primitives::{
@@ -468,8 +468,12 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
     }
 
     /// Calculates the hashes of the nodes below the provided level.
+    ///
+    /// If the trie has not been revealed, this function does nothing.
     pub fn calculate_below_level(&mut self, level: usize) {
-        self.state.calculate_below_level(level);
+        if let SparseTrie::Revealed(trie) = &mut self.state {
+            trie.update_rlp_node_level(level);
+        }
     }
 
     /// Returns storage sparse trie root if the trie has been revealed.
@@ -477,22 +481,57 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         self.storages.get_mut(&account).and_then(|trie| trie.root())
     }
 
-    /// Returns sparse trie root if the trie has been revealed.
-    pub fn root(&mut self) -> Option<B256> {
-        self.state.root()
+    /// Returns mutable reference to the revealed sparse trie.
+    ///
+    /// If the trie is not revealed yet, its root will be revealed using the blinded node provider.
+    fn revealed_trie_mut(
+        &mut self,
+    ) -> SparseStateTrieResult<&mut RevealedSparseTrie<F::AccountNodeProvider>> {
+        match self.state {
+            SparseTrie::Blind => {
+                let (root_node, hash_mask, tree_mask) = self
+                    .provider_factory
+                    .account_node_provider()
+                    .blinded_node(&Nibbles::default())?
+                    .map(|node| {
+                        TrieNode::decode(&mut &node.node[..])
+                            .map(|decoded| (decoded, node.hash_mask, node.tree_mask))
+                    })
+                    .transpose()?
+                    .unwrap_or((TrieNode::EmptyRoot, None, None));
+                self.state
+                    .reveal_root_with_provider(
+                        self.provider_factory.account_node_provider(),
+                        root_node,
+                        hash_mask,
+                        tree_mask,
+                        self.retain_updates,
+                    )
+                    .map_err(Into::into)
+            }
+            SparseTrie::Revealed(ref mut trie) => Ok(trie),
+        }
+    }
+
+    /// Returns sparse trie root.
+    ///
+    /// If the trie has not been revealed, this function reveals the root node and returns its hash.
+    pub fn root(&mut self) -> SparseStateTrieResult<B256> {
+        Ok(self.revealed_trie_mut()?.root())
     }
 
     /// Returns sparse trie root and trie updates if the trie has been revealed.
-    pub fn root_with_updates(&mut self) -> Option<(B256, TrieUpdates)> {
+    pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> {
         let storage_tries = self.storage_trie_updates();
-        self.state.root_with_updates().map(|(root, updates)| {
-            let updates = TrieUpdates {
-                account_nodes: updates.updated_nodes,
-                removed_nodes: updates.removed_nodes,
-                storage_tries,
-            };
-            (root, updates)
-        })
+        let revealed = self.revealed_trie_mut()?;
+
+        let (root, updates) = (revealed.root(), revealed.take_updates());
+        let updates = TrieUpdates {
+            account_nodes: updates.updated_nodes,
+            removed_nodes: updates.removed_nodes,
+            storage_tries,
+        };
+        Ok((root, updates))
     }
 
     /// Returns storage trie updates for tries that have been revealed.
@@ -923,7 +962,7 @@ mod tests {
             })
             .unwrap();
 
-        assert_eq!(sparse.root(), Some(root));
+        assert_eq!(sparse.root().unwrap(), root);
 
         let address_3 = b256!("2000000000000000000000000000000000000000000000000000000000000000");
         let address_path_3 = Nibbles::unpack(address_3);
@@ -940,7 +979,7 @@ mod tests {
         trie_account_2.storage_root = sparse.storage_root(address_2).unwrap();
         sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap();
 
-        sparse.root();
+        sparse.root().unwrap();
 
         let sparse_updates = sparse.take_trie_updates().unwrap();
         // TODO(alexey): assert against real state root calculation updates
diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index 4636c15c9..63afde437 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -139,11 +139,6 @@ impl<P> SparseTrie<P> {
         let revealed = self.as_revealed_mut()?;
         Some((revealed.root(), revealed.take_updates()))
     }
-
-    /// Calculates the hashes of the nodes below the provided level.
-    pub fn calculate_below_level(&mut self, level: usize) {
-        self.as_revealed_mut().unwrap().update_rlp_node_level(level);
-    }
 }
 
 impl<P: BlindedProvider> SparseTrie<P> {