From dd06e06660500fd57bff8f0da577fe04d97cd0d6 Mon Sep 17 00:00:00 2001
From: AvivYossef-starkware <aviv.y@starkware.co>
Date: Wed, 30 Oct 2024 11:31:08 +0200
Subject: [PATCH] refactor(blockifier_reexecution): serialize StateMaps with
 costum serializer and output better error format

---
 Cargo.lock                                  |  1 +
 crates/blockifier/src/state/cached_state.rs | 42 ++++++++++++++++++++-
 crates/blockifier_reexecution/Cargo.toml    |  1 +
 crates/blockifier_reexecution/src/main.rs   | 15 ++++++--
 4 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1ff50d1206..8211d18ee8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1550,6 +1550,7 @@ dependencies = [
  "starknet-types-core",
  "starknet_api",
  "starknet_gateway",
+ "tempfile",
  "thiserror",
 ]
 
diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs
index d9a4a598f6..6f9b22de20 100644
--- a/crates/blockifier/src/state/cached_state.rs
+++ b/crates/blockifier/src/state/cached_state.rs
@@ -1,7 +1,8 @@
 use std::cell::RefCell;
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 
 use indexmap::IndexMap;
+use serde::{Serialize, Serializer};
 use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
@@ -545,14 +546,53 @@ type StorageDiff = IndexMap<ContractAddress, IndexMap<StorageKey, Felt>>;
 #[derive(Debug, Default, Eq, PartialEq)]
 pub struct CommitmentStateDiff {
     // Contract instance attributes (per address).
+    #[serde(serialize_with = "serialize_sorted")]
     pub address_to_class_hash: IndexMap<ContractAddress, ClassHash>,
+    #[serde(serialize_with = "serialize_sorted")]
     pub address_to_nonce: IndexMap<ContractAddress, Nonce>,
+    #[serde(serialize_with = "serialize_nested_sorted")]
     pub storage_updates: IndexMap<ContractAddress, IndexMap<StorageKey, Felt>>,
 
     // Global attributes.
+    #[serde(serialize_with = "serialize_sorted")]
     pub class_hash_to_compiled_class_hash: IndexMap<ClassHash, CompiledClassHash>,
 }
 
+// Define custom serializer to sort by keys.
+fn serialize_sorted<S, K, V>(map: &IndexMap<K, V>, serializer: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+    K: Ord + Serialize,
+    V: Serialize,
+{
+    // Convert the IndexMap to BTreeMap, which automatically sorts keys.
+    let sorted_map: BTreeMap<_, _> = map.iter().collect();
+    sorted_map.serialize(serializer)
+}
+
+// Custom serializer for nested `IndexMap` with different key types.
+fn serialize_nested_sorted<S, OuterK, InnerK, V>(
+    map: &IndexMap<OuterK, IndexMap<InnerK, V>>,
+    serializer: S,
+) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+    OuterK: Ord + Serialize,
+    InnerK: Ord + Serialize,
+    V: Serialize,
+{
+    // Sort each inner `IndexMap` and store in a `BTreeMap` for serialization.
+    let sorted_map: BTreeMap<_, BTreeMap<_, _>> = map
+        .iter()
+        .map(|(outer_key, inner_map)| {
+            let sorted_inner_map: BTreeMap<_, _> = inner_map.iter().collect();
+            (outer_key, sorted_inner_map)
+        })
+        .collect();
+
+    sorted_map.serialize(serializer)
+}
+
 impl From<StateMaps> for CommitmentStateDiff {
     fn from(diff: StateMaps) -> Self {
         Self {
diff --git a/crates/blockifier_reexecution/Cargo.toml b/crates/blockifier_reexecution/Cargo.toml
index dcf234df12..589b291698 100644
--- a/crates/blockifier_reexecution/Cargo.toml
+++ b/crates/blockifier_reexecution/Cargo.toml
@@ -24,6 +24,7 @@ starknet-core.workspace = true
 starknet-types-core.workspace = true
 starknet_api.workspace = true
 starknet_gateway.workspace = true
+tempfile.workspace = true
 thiserror.workspace = true
 
 [dev-dependencies]
diff --git a/crates/blockifier_reexecution/src/main.rs b/crates/blockifier_reexecution/src/main.rs
index fabd5de409..927f9d637f 100644
--- a/crates/blockifier_reexecution/src/main.rs
+++ b/crates/blockifier_reexecution/src/main.rs
@@ -4,11 +4,9 @@ use blockifier_reexecution::state_reader::test_state_reader::{
 };
 use blockifier_reexecution::state_reader::utils::JSON_RPC_VERSION;
 use clap::{Args, Parser, Subcommand};
-use pretty_assertions::assert_eq;
 use starknet_api::block::BlockNumber;
 use starknet_api::core::ContractAddress;
 use starknet_gateway::config::RpcStateReaderConfig;
-
 /// BlockifierReexecution CLI.
 #[derive(Debug, Parser)]
 #[clap(name = "blockifier-reexecution-cli", version)]
@@ -89,9 +87,18 @@ fn main() {
                 transaction_executor.finalize().expect("Couldn't finalize block");
             // TODO(Aner): compute correct block hash at storage slot 0x1 instead of removing it.
             expected_state_diff.storage_updates.shift_remove(&ContractAddress(1_u128.into()));
-            assert_eq!(expected_state_diff, actual_state_diff);
 
-            println!("RPC test passed successfully.");
+            if expected_state_diff != actual_state_diff {
+                let expected_json = serde_json::to_value(&expected_state_diff)
+                    .expect("Failed to serialize expected_state_diff");
+                let actual_json = serde_json::to_value(&actual_state_diff)
+                    .expect("Failed to serialize actual_state_diff");
+                eprintln!("Test failed! Differences in state_diff");
+                eprintln!("Expected: {expected_json}");
+                eprintln!("Actual: {actual_json}");
+            } else {
+                println!("RPC test passed successfully.");
+            }
         }
         Command::WriteRpcRepliesToJson { .. } => todo!(),
     }