fix(stage): handle state root computation when class root is zero (#335)

Author: kariy

crates/sync/stage/src/trie.rs

@@ -6,7 +6,7 @@ use katana_provider::api::state_update::StateUpdateProvider;
 use katana_provider::api::trie::TrieWriter;
 use starknet::macros::short_string;
 use starknet_types_core::hash::{Poseidon, StarkHash};
-use tracing::{error, trace, trace_span};
+use tracing::{debug, debug_span, error};
 
 use crate::{Stage, StageExecutionInput, StageExecutionOutput, StageResult};
 
@@ -41,7 +41,7 @@ where
     fn execute<'a>(&'a mut self, input: &'a StageExecutionInput) -> BoxFuture<'a, StageResult> {
         Box::pin(async move {
             for block_number in input.from()..=input.to() {
-                let span = trace_span!("compute_state_root", %block_number);
+                let span = debug_span!("state_trie.compute_state_root", %block_number);
                 let _enter = span.enter();
 
                 let expected_state_root = self
@@ -55,29 +55,41 @@ where
                     .state_update(block_number.into())?
                     .ok_or(Error::MissingStateUpdate(block_number))?;
 
-                let class_trie_root = self
+                let computed_contract_trie_root =
+                    self.provider.trie_insert_contract_updates(block_number, &state_update)?;
+
+                debug!(
+                    contract_trie_root = format!("{computed_contract_trie_root:#x}"),
+                    "Computed contract trie root."
+                );
+
+                let computed_class_trie_root = self
                     .provider
                     .trie_insert_declared_classes(block_number, &state_update.declared_classes)?;
 
-                let contract_trie_root =
-                    self.provider.trie_insert_contract_updates(block_number, &state_update)?;
+                debug!(
+                    classes_tri_root = format!("{computed_class_trie_root:#x}"),
+                    "Computed classes trie root."
+                );
 
-                // Compute the state root:
-                // hash("STARKNET_STATE_V0", contract_trie_root, class_trie_root)
-                let computed_state_root = Poseidon::hash_array(&[
-                    short_string!("STARKNET_STATE_V0"),
-                    contract_trie_root,
-                    class_trie_root,
-                ]);
+                let computed_state_root = if computed_class_trie_root == Felt::ZERO {
+                    computed_contract_trie_root
+                } else {
+                    Poseidon::hash_array(&[
+                        short_string!("STARKNET_STATE_V0"),
+                        computed_contract_trie_root,
+                        computed_class_trie_root,
+                    ])
+                };
 
                 // Verify that the computed state root matches the expected state root from the
                 // block header
                 if computed_state_root != expected_state_root {
                     error!(
-                        block = %block_number,
+                        target: "stage",
                         state_root = %format!("{computed_state_root:#x}"),
                         expected_state_root = %format!("{expected_state_root:#x}"),
-                        "Wrong trie root for block - computed state root does not match expected state root (from header)",
+                        "Bad state trie root for block - computed state root does not match expected state root (from header)",
                     );
 
                     return Err(Error::StateRootMismatch {
@@ -88,11 +100,7 @@ where
                     .into());
                 }
 
-                trace!(
-                    block = %block_number,
-                    state_root = %format!("{computed_state_root:#x}"),
-                    "State root verified successfully."
-                );
+                debug!("State root verified successfully.");
             }
 
             Ok(StageExecutionOutput { last_block_processed: input.to() })