refactor(blockifier_reexecution): add chain ID to TestStateReader

starkware-libs · Nov 12, 2024 · 3d59ffc · 3d59ffc
1 parent dbf7b4a
commit 3d59ffc
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/crates/blockifier_reexecution/src/main.rs b/crates/blockifier_reexecution/src/main.rs
@@ -10,6 +10,7 @@ use blockifier_reexecution::state_reader::utils::{
 };
 use clap::{Args, Parser, Subcommand};
 use starknet_api::block::BlockNumber;
+use starknet_api::core::ChainId;
 use starknet_gateway::config::RpcStateReaderConfig;
 
 /// BlockifierReexecution CLI.
@@ -86,6 +87,7 @@ fn main() {
             reexecute_and_verify_correctness(ConsecutiveTestStateReaders::new(
                 BlockNumber(block_number - 1),
                 Some(config),
+                ChainId::Mainnet,
                 false,
             ));
 
@@ -106,8 +108,12 @@ fn main() {
                 json_rpc_version: JSON_RPC_VERSION.to_string(),
             };
 
-            let consecutive_state_readers =
-                ConsecutiveTestStateReaders::new(BlockNumber(block_number - 1), Some(config), true);
+            let consecutive_state_readers = ConsecutiveTestStateReaders::new(
+                BlockNumber(block_number - 1),
+                Some(config),
+                ChainId::Mainnet,
+                true,
+            );
 
             let serializable_data_next_block =
                 consecutive_state_readers.get_serializable_data_next_block().unwrap();

diff --git a/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs b/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
@@ -169,6 +169,7 @@ impl Default for RetryConfig {
 pub struct TestStateReader {
     pub(crate) rpc_state_reader: RpcStateReader,
     pub(crate) retry_config: RetryConfig,
+    pub(crate) chain_id: ChainId,
     #[allow(dead_code)]
     pub(crate) contract_class_mapping_dumper: Arc<Mutex<Option<StarknetContractClassMapping>>>,
 }
@@ -178,6 +179,7 @@ impl Default for TestStateReader {
         Self {
             rpc_state_reader: RpcStateReader::from_latest(&get_rpc_state_reader_config()),
             retry_config: RetryConfig::default(),
+            chain_id: ChainId::Mainnet,
             contract_class_mapping_dumper: Arc::new(Mutex::new(None)),
         }
     }
@@ -229,20 +231,26 @@ impl StateReader for TestStateReader {
 }
 
 impl TestStateReader {
-    pub fn new(config: &RpcStateReaderConfig, block_number: BlockNumber, dump_mode: bool) -> Self {
+    pub fn new(
+        config: &RpcStateReaderConfig,
+        chain_id: ChainId,
+        block_number: BlockNumber,
+        dump_mode: bool,
+    ) -> Self {
         let contract_class_mapping_dumper = Arc::new(Mutex::new(match dump_mode {
             true => Some(HashMap::new()),
             false => None,
         }));
         Self {
             rpc_state_reader: RpcStateReader::from_number(config, block_number),
-            contract_class_mapping_dumper,
             retry_config: RetryConfig::default(),
+            chain_id,
+            contract_class_mapping_dumper,
         }
     }
 
     pub fn new_for_testing(block_number: BlockNumber) -> Self {
-        TestStateReader::new(&get_rpc_state_reader_config(), block_number, false)
+        TestStateReader::new(&get_rpc_state_reader_config(), ChainId::Mainnet, block_number, false)
     }
 
     /// Get the block info of the current block.
@@ -327,7 +335,7 @@ impl TestStateReader {
     pub fn get_block_context(&self) -> ReexecutionResult<BlockContext> {
         Ok(BlockContext::new(
             self.get_block_info()?,
-            get_chain_info(&ChainId::Mainnet),
+            get_chain_info(&self.chain_id),
             self.get_versioned_constants()?.clone(),
             BouncerConfig::max(),
         ))
@@ -461,17 +469,20 @@ impl ConsecutiveTestStateReaders {
     pub fn new(
         last_constructed_block_number: BlockNumber,
         config: Option<RpcStateReaderConfig>,
+        chain_id: ChainId,
         dump_mode: bool,
     ) -> Self {
         let config = config.unwrap_or(get_rpc_state_reader_config());
         Self {
             last_block_state_reader: TestStateReader::new(
                 &config,
+                chain_id.clone(),
                 last_constructed_block_number,
                 dump_mode,
             ),
             next_block_state_reader: TestStateReader::new(
                 &config,
+                chain_id,
                 last_constructed_block_number.next().expect("Overflow in block number"),
                 false,
             ),