starknet_os_flow_tests: state diff comparison (#8497)

Author: nimrod-starkware

crates/blockifier/src/state/cached_state.rs

@@ -368,6 +368,16 @@ impl StateMaps {
             compiled_class_hash_keys: self.compiled_class_hashes.keys().cloned().collect(),
         }
     }
+
+    /// Returns the set of keys that aliases were potentially allocated for.
+    #[cfg(any(test, feature = "testing"))]
+    pub fn alias_keys(&self) -> HashSet<crate::state::stateful_compression::AliasKey> {
+        let mut keys = HashSet::from_iter(
+            self.get_contract_addresses().into_iter().map(|address| StorageKey(address.0)),
+        );
+        keys.extend(self.storage.keys().map(|(_address, storage_key)| *storage_key));
+        keys
+    }
 }
 
 /// Caches read and write requests.

crates/starknet_os/src/io.rs

@@ -1,5 +1,5 @@
 pub mod os_input;
 pub mod os_output;
 pub mod os_output_types;
-#[cfg(test)]
+#[cfg(any(feature = "testing", test))]
 pub mod test_utils;

crates/starknet_os/src/io/test_utils.rs

@@ -132,13 +132,13 @@ fn to_state_maps<CO: ContractChangesGetter, CL: UpdateGetter<ClassHash, Compiled
 }
 
 impl FullOsStateDiff {
-    pub(crate) fn as_state_maps(&self) -> StateMaps {
+    pub fn as_state_maps(&self) -> StateMaps {
         to_state_maps(&self.contracts, &self.classes)
     }
 }
 
 impl PartialOsStateDiff {
-    pub(crate) fn _as_state_maps(&self) -> StateMaps {
+    pub fn as_state_maps(&self) -> StateMaps {
         to_state_maps(&self.contracts, &self.classes)
     }
 }

crates/starknet_os_flow_tests/Cargo.toml

@@ -14,7 +14,7 @@ rstest.workspace = true
 starknet-types-core.workspace = true
 starknet_api.workspace = true
 starknet_committer.workspace = true
-starknet_os = { workspace = true, features = ["testing"] }
+starknet_os = { workspace = true, features = ["include_program_output", "testing"] }
 starknet_patricia = { workspace = true, features = ["testing"] }
 starknet_patricia_storage = { workspace = true, features = ["testing"] }
 tokio.workspace = true

crates/starknet_os_flow_tests/src/test_manager.rs

@@ -1,11 +1,14 @@
 #![allow(dead_code)]
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::LazyLock;
 
 use blockifier::abi::constants::STORED_BLOCK_HASH_BUFFER;
 use blockifier::blockifier_versioned_constants::VersionedConstants;
 use blockifier::bouncer::BouncerConfig;
 use blockifier::context::{BlockContext, ChainInfo, FeeTokenAddresses};
+use blockifier::state::cached_state::StateMaps;
+use blockifier::state::stateful_compression_test_utils::decompress;
+use blockifier::test_utils::ALIAS_CONTRACT_ADDRESS;
 use blockifier::transaction::transaction_execution::Transaction as BlockifierTransaction;
 use starknet_api::block::{BlockHash, BlockInfo, BlockNumber};
 use starknet_api::contract_class::ContractClass;
@@ -26,7 +29,7 @@ use starknet_os::io::os_input::{
     OsHintsConfig,
     StarknetOsInput,
 };
-use starknet_os::io::os_output::StarknetOsRunnerOutput;
+use starknet_os::io::os_output::{OsStateDiff, StarknetOsRunnerOutput};
 use starknet_os::runner::{run_os_stateless, DEFAULT_OS_LAYOUT};
 use starknet_patricia_storage::map_storage::BorrowedMapStorage;
 use starknet_types_core::felt::Felt;
@@ -68,6 +71,11 @@ pub(crate) struct TestManager<S: FlowTestState> {
     per_block_transactions: Vec<Vec<BlockifierTransaction>>,
 }
 
+pub(crate) struct OsTestOutput {
+    pub(crate) os_output: StarknetOsRunnerOutput,
+    pub(crate) decompressed_state_diff: StateMaps,
+}
+
 impl<S: FlowTestState> TestManager<S> {
     /// Creates a new `TestManager` with the provided initial state data.
     pub(crate) fn new_with_initial_state_data(initial_state_data: InitialStateData<S>) -> Self {
@@ -154,7 +162,7 @@ impl<S: FlowTestState> TestManager<S> {
     pub(crate) async fn execute_test_with_default_block_contexts(
         self,
         initial_block_number: u64,
-    ) -> StarknetOsRunnerOutput {
+    ) -> OsTestOutput {
         let n_blocks = self.per_block_transactions.len();
         let block_contexts: Vec<BlockContext> = (0..n_blocks)
             .map(|i| {
@@ -171,7 +179,7 @@ impl<S: FlowTestState> TestManager<S> {
     pub(crate) async fn execute_test_with_block_contexts(
         self,
         block_contexts: Vec<BlockContext>,
-    ) -> StarknetOsRunnerOutput {
+    ) -> OsTestOutput {
         assert_eq!(
             block_contexts.len(),
             self.per_block_transactions.len(),
@@ -193,10 +201,7 @@ impl<S: FlowTestState> TestManager<S> {
     }
 
     // Private method which executes the flow test.
-    async fn execute_flow_test(
-        mut self,
-        block_contexts: Vec<BlockContext>,
-    ) -> StarknetOsRunnerOutput {
+    async fn execute_flow_test(mut self, block_contexts: Vec<BlockContext>) -> OsTestOutput {
         let per_block_txs = self.per_block_transactions;
         let mut os_block_inputs = vec![];
         let mut cached_state_inputs = vec![];
@@ -209,6 +214,7 @@ impl<S: FlowTestState> TestManager<S> {
             contracts_trie_root_hash: self.initial_state.contracts_trie_root_hash,
             classes_trie_root_hash: self.initial_state.classes_trie_root_hash,
         };
+        let mut alias_keys = HashSet::new();
         for (block_txs, block_context) in per_block_txs.into_iter().zip(block_contexts.into_iter())
         {
             // Clone the block info for later use.
@@ -220,6 +226,7 @@ impl<S: FlowTestState> TestManager<S> {
             // Update the wrapped state.
             let state_diff = final_state.to_state_diff().unwrap();
             state = final_state.state;
+            alias_keys.extend(state_diff.state_maps.alias_keys());
             state.apply_writes(&state_diff.state_maps, &final_state.class_hash_to_class.borrow());
             // Commit the state diff.
             let committer_state_diff = create_committer_state_diff(block_summary.state_diff);
@@ -290,10 +297,18 @@ impl<S: FlowTestState> TestManager<S> {
             chain_id: CHAIN_ID_FOR_TESTS.clone(),
             strk_fee_token_address: *STRK_FEE_TOKEN_ADDRESS,
         };
-        let os_hints_config = OsHintsConfig { full_output: true, chain_info, ..Default::default() };
+        let os_hints_config = OsHintsConfig { chain_info, ..Default::default() };
         let os_hints = OsHints { os_input: starknet_os_input, os_hints_config };
         let layout = DEFAULT_OS_LAYOUT;
-        run_os_stateless(layout, os_hints).unwrap()
+        let os_output = run_os_stateless(layout, os_hints).unwrap();
+        let OsStateDiff::Partial(ref partial_os_state_diff) = os_output.os_output.state_diff else {
+            panic!("Expected a partial state diff in the output because of the OS config.");
+        };
+        let compressed_state_diff = partial_os_state_diff.as_state_maps();
+        let decompressed_state_diff =
+            decompress(&compressed_state_diff, &state, *ALIAS_CONTRACT_ADDRESS, alias_keys);
+
+        OsTestOutput { os_output, decompressed_state_diff }
     }
 }
 

crates/starknet_os_flow_tests/src/tests.rs

@@ -6,6 +6,8 @@ use blockifier_test_utils::cairo_versions::{CairoVersion, RunnableCairo1};
 use blockifier_test_utils::calldata::create_calldata;
 use blockifier_test_utils::contracts::FeatureContract;
 use rstest::rstest;
+use starknet_api::abi::abi_utils::get_storage_var_address;
+use starknet_api::core::calculate_contract_address;
 use starknet_api::executable_transaction::{DeclareTransaction, InvokeTransaction};
 use starknet_api::execution_resources::GasAmount;
 use starknet_api::test_utils::declare::declare_tx;
@@ -18,12 +20,18 @@ use starknet_api::test_utils::{
     DEFAULT_STRK_L2_GAS_PRICE,
 };
 use starknet_api::transaction::constants::DEPLOY_CONTRACT_FUNCTION_ENTRY_POINT_NAME;
-use starknet_api::transaction::fields::{AllResourceBounds, ResourceBounds, ValidResourceBounds};
+use starknet_api::transaction::fields::{
+    AllResourceBounds,
+    Calldata,
+    ContractAddressSalt,
+    ResourceBounds,
+    ValidResourceBounds,
+};
 use starknet_api::{declare_tx_args, invoke_tx_args};
 use starknet_types_core::felt::Felt;
 
 use crate::initial_state::create_default_initial_state_data;
-use crate::test_manager::{TestManager, FUNDED_ACCOUNT_ADDRESS};
+use crate::test_manager::{OsTestOutput, TestManager, FUNDED_ACCOUNT_ADDRESS};
 use crate::utils::{divide_vec_into_n_parts, get_class_info_of_cairo_1_feature_contract};
 
 pub(crate) static NON_TRIVIAL_RESOURCE_BOUNDS: LazyLock<ValidResourceBounds> =
@@ -92,16 +100,17 @@ async fn declare_deploy_scenario(#[values(1, 2)] n_blocks: usize) {
         DeclareTransaction::create(account_declare_tx, class_info, &CHAIN_ID_FOR_TESTS).unwrap();
     // Add the transaction to the test manager.
     test_manager.add_cairo1_declare_tx(tx, &test_contract_sierra);
-
+    let arg1 = Felt::from(7);
+    let arg2 = Felt::from(90);
     // Deploy the test contract using the deploy contract syscall.
     let constructor_calldata = [
-        2.into(),  // constructor length
-        7.into(),  // arg1
-        90.into(), // arg2
+        2.into(), // constructor length
+        arg1,
+        arg2,
     ];
-    let contract_address_salt = Felt::ONE;
+    let contract_address_salt = ContractAddressSalt(Felt::ONE);
     let calldata: Vec<_> =
-        [class_hash.0, contract_address_salt].into_iter().chain(constructor_calldata).collect();
+        [class_hash.0, contract_address_salt.0].into_iter().chain(constructor_calldata).collect();
     let deploy_contract_calldata = create_calldata(
         *FUNDED_ACCOUNT_ADDRESS,
         DEPLOY_CONTRACT_FUNCTION_ENTRY_POINT_NAME,
@@ -113,14 +122,40 @@ async fn declare_deploy_scenario(#[values(1, 2)] n_blocks: usize) {
         calldata: deploy_contract_calldata,
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
     };
+    let expected_contract_address = calculate_contract_address(
+        contract_address_salt,
+        class_hash,
+        &Calldata(constructor_calldata[1..].to_vec().into()),
+        *FUNDED_ACCOUNT_ADDRESS,
+    )
+    .unwrap();
     let deploy_contract_tx = invoke_tx(invoke_tx_args);
     let deploy_contract_tx =
         InvokeTransaction::create(deploy_contract_tx, &CHAIN_ID_FOR_TESTS).unwrap();
     test_manager.add_invoke_tx(deploy_contract_tx);
     test_manager.divide_transactions_into_n_blocks(n_blocks);
     let initial_block_number = CURRENT_BLOCK_NUMBER + 1;
-    let _os_runner_output =
+    let OsTestOutput { decompressed_state_diff, .. } =
         test_manager.execute_test_with_default_block_contexts(initial_block_number).await;
 
-    // TODO(Nimrod): Validate the OS output.
+    // Verify that the OS output contains the entry in the classes changes from the declare tx.
+    assert!(decompressed_state_diff.compiled_class_hashes.get(&class_hash).is_some_and(
+        |actual_compiled_class_hash| *actual_compiled_class_hash == compiled_class_hash
+    ));
+
+    // Verify the deployed contract.
+    assert!(
+        decompressed_state_diff
+            .class_hashes
+            .get(&expected_contract_address)
+            .is_some_and(|actual_class_hash| *actual_class_hash == class_hash)
+    );
+
+    // Verify the storage update from the contract's constructor.
+    assert!(
+        decompressed_state_diff
+            .storage
+            .get(&(expected_contract_address, get_storage_var_address("my_storage_var", &[])))
+            .is_some_and(|actual_value| *actual_value == arg1 + arg2)
+    );
 }