starknet_os_flow_tests: initial state with extra contracts (#9422)

Author: dorimedini-starkware

crates/starknet_os_flow_tests/src/initial_state.rs

@@ -1,5 +1,5 @@
 #![allow(dead_code)]
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use blockifier::transaction::transaction_execution::Transaction;
 use blockifier_test_utils::cairo_versions::{CairoVersion, RunnableCairo1};
@@ -40,11 +40,14 @@ use crate::test_manager::{
     FUNDED_ACCOUNT_ADDRESS,
     STRK_FEE_TOKEN_ADDRESS,
 };
+use crate::tests::NON_TRIVIAL_RESOURCE_BOUNDS;
 use crate::utils::{
     commit_state_diff,
     create_cairo1_bootstrap_declare_tx,
     create_committer_state_diff,
+    create_declare_tx,
     execute_transactions,
+    get_class_hash_of_feature_contract,
     CommitmentOutput,
     ExecutionOutput,
 };
@@ -121,13 +124,18 @@ pub(crate) struct InitialState<S: FlowTestState> {
 /// Creates the initial state for the flow test which includes:
 /// Declares token and account contracts.
 /// Deploys both contracts and funds the account.
-pub(crate) async fn create_default_initial_state_data<S: FlowTestState>()
--> (InitialStateData<S>, NonceManager) {
-    let InitialTransactionsData {
-        transactions: default_initial_state_txs,
-        execution_contracts,
-        nonce_manager,
-    } = create_default_initial_state_txs_and_contracts();
+/// Also deploys extra contracts as requested (and declares them if they are not already declared).
+pub(crate) async fn create_default_initial_state_data<S: FlowTestState, const N: usize>(
+    extra_contracts: [(FeatureContract, Calldata); N],
+) -> (InitialStateData<S>, NonceManager, [ContractAddress; N]) {
+    let (
+        InitialTransactionsData {
+            transactions: default_initial_state_txs,
+            execution_contracts,
+            nonce_manager,
+        },
+        extra_contracts_addresses,
+    ) = create_default_initial_state_txs_and_contracts(extra_contracts);
     // Execute these 4 txs.
     let initial_state_reader = S::create_empty_state();
     let use_kzg_da = false;
@@ -170,7 +178,11 @@ pub(crate) async fn create_default_initial_state_data<S: FlowTestState>()
         classes_trie_root_hash: commitment_output.classes_trie_root_hash,
     };
 
-    (InitialStateData { initial_state, execution_contracts }, nonce_manager)
+    (
+        InitialStateData { initial_state, execution_contracts },
+        nonce_manager,
+        extra_contracts_addresses,
+    )
 }
 
 struct InitialTransactionsData {
@@ -179,7 +191,9 @@ struct InitialTransactionsData {
     pub(crate) nonce_manager: NonceManager,
 }
 
-fn create_default_initial_state_txs_and_contracts() -> InitialTransactionsData {
+fn create_default_initial_state_txs_and_contracts<const N: usize>(
+    extra_contracts: [(FeatureContract, Calldata); N],
+) -> (InitialTransactionsData, [ContractAddress; N]) {
     let mut os_execution_contracts = OsExecutionContracts::default();
     // Declare account and ERC20 contracts.
     let account_contract =
@@ -212,11 +226,38 @@ fn create_default_initial_state_txs_and_contracts() -> InitialTransactionsData {
     let nonce = nonce_manager.next(account_contract_address);
     let (deploy_contract_tx, _) = get_deploy_fee_token_tx_and_address(nonce);
     txs.push(deploy_contract_tx);
-    InitialTransactionsData {
-        transactions: txs,
-        execution_contracts: os_execution_contracts,
-        nonce_manager,
+
+    // Deploy extra contracts. Declare contracts that are not already declared.
+    let mut declared_contracts = HashSet::from([account_contract, erc20_contract]);
+    let mut extra_addresses = Vec::new();
+    for (contract, calldata) in extra_contracts {
+        if !declared_contracts.contains(&contract) {
+            // Add a declare transaction for the contract.
+            // No need for bootstrap mode: funded account already exists at this point.
+            txs.push(Transaction::new_for_sequencing(StarknetAPITransaction::Account(
+                create_declare_tx(contract, &mut nonce_manager, &mut os_execution_contracts, false),
+            )));
+            declared_contracts.insert(contract);
+        }
+        // Deploy.
+        let (deploy_tx, address) = get_deploy_contract_tx_and_address(
+            get_class_hash_of_feature_contract(contract),
+            calldata,
+            nonce_manager.next(account_contract_address),
+            *NON_TRIVIAL_RESOURCE_BOUNDS,
+        );
+        txs.push(deploy_tx);
+        extra_addresses.push(address);
     }
+
+    (
+        InitialTransactionsData {
+            transactions: txs,
+            execution_contracts: os_execution_contracts,
+            nonce_manager,
+        },
+        extra_addresses.try_into().unwrap(),
+    )
 }
 
 pub(crate) async fn commit_initial_state_diff(

crates/starknet_os_flow_tests/src/test_manager.rs

@@ -10,6 +10,7 @@ use blockifier::state::cached_state::{CommitmentStateDiff, 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 blockifier_test_utils::contracts::FeatureContract;
 use itertools::Itertools;
 use starknet_api::block::{BlockHash, BlockInfo, BlockNumber, PreviousBlockNumber};
 use starknet_api::contract_class::compiled_class_hash::{HashVersion, HashableCompiledClass};
@@ -24,6 +25,7 @@ use starknet_api::executable_transaction::{
 };
 use starknet_api::state::{SierraContractClass, StorageKey};
 use starknet_api::test_utils::{NonceManager, CHAIN_ID_FOR_TESTS};
+use starknet_api::transaction::fields::Calldata;
 use starknet_api::transaction::MessageToL1;
 use starknet_committer::block_committer::input::{IsSubset, StateDiff};
 use starknet_os::io::os_input::{
@@ -228,10 +230,18 @@ impl<S: FlowTestState> TestManager<S> {
 
     /// Creates a new `TestManager` with the default initial state.
     /// Returns the manager and a nonce manager to help keep track of nonces.
-    pub(crate) async fn new_with_default_initial_state() -> (Self, NonceManager) {
-        let (default_initial_state_data, nonce_manager) =
-            create_default_initial_state_data::<S>().await;
-        (Self::new_with_initial_state_data(default_initial_state_data), nonce_manager)
+    /// Optionally provide an array of extra contracts to declare and deploy - the addresses of
+    /// these contracts will be returned as an array of the same length.
+    pub(crate) async fn new_with_default_initial_state<const N: usize>(
+        extra_contracts: [(FeatureContract, Calldata); N],
+    ) -> (Self, NonceManager, [ContractAddress; N]) {
+        let (default_initial_state_data, nonce_manager, extra_addresses) =
+            create_default_initial_state_data::<S, N>(extra_contracts).await;
+        (
+            Self::new_with_initial_state_data(default_initial_state_data),
+            nonce_manager,
+            extra_addresses,
+        )
     }
 
     /// Advances the manager to the next block when adding new transactions.

crates/starknet_os_flow_tests/src/tests.rs

@@ -28,7 +28,7 @@ use starknet_api::transaction::fields::{
     ResourceBounds,
     ValidResourceBounds,
 };
-use starknet_api::{declare_tx_args, invoke_tx_args};
+use starknet_api::{calldata, declare_tx_args, invoke_tx_args};
 use starknet_committer::block_committer::input::{
     StarknetStorageKey,
     StarknetStorageValue,
@@ -61,7 +61,7 @@ pub(crate) static NON_TRIVIAL_RESOURCE_BOUNDS: LazyLock<ValidResourceBounds> =
 
 #[tokio::test]
 async fn test_initial_state_creation() {
-    let _initial_state = create_default_initial_state_data::<DictStateReader>().await;
+    let _initial_state = create_default_initial_state_data::<DictStateReader, 0>([]).await;
 }
 
 #[rstest]
@@ -91,8 +91,8 @@ async fn declare_deploy_scenario(
     // Initialize the test manager with a default initial state and get the nonce manager to help
     // keep track of nonces.
 
-    let (mut test_manager, mut nonce_manager) =
-        TestManager::<DictStateReader>::new_with_default_initial_state().await;
+    let (mut test_manager, mut nonce_manager, _) =
+        TestManager::<DictStateReader>::new_with_default_initial_state([]).await;
 
     // Declare a test contract.
     let test_contract = FeatureContract::TestContract(CairoVersion::Cairo1(RunnableCairo1::Casm));
@@ -184,66 +184,15 @@ async fn trivial_diff_scenario(
     #[values(false, true)] use_kzg_da: bool,
     #[values(false, true)] full_output: bool,
 ) {
-    // TODO(Dori): Create a util to define an initial state with the test contract declared and
-    //   deployed. The address of the test contract should be part of the initial state data.
-
     // Initialize the test manager with a default initial state and get the nonce manager to help
     // keep track of nonces.
 
-    let (mut test_manager, mut nonce_manager) =
-        TestManager::<DictStateReader>::new_with_default_initial_state().await;
-
-    // Declare and deploy a test contract.
-    let test_contract = FeatureContract::TestContract(CairoVersion::Cairo1(RunnableCairo1::Casm));
-    let test_contract_sierra = test_contract.get_sierra();
-    let class_hash = test_contract_sierra.calculate_class_hash();
-    let compiled_class_hash = test_contract.get_compiled_class_hash(&HashVersion::V2);
-    let declare_tx_args = declare_tx_args! {
-        sender_address: *FUNDED_ACCOUNT_ADDRESS,
-        class_hash,
-        compiled_class_hash,
-        resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
-        nonce: nonce_manager.next(*FUNDED_ACCOUNT_ADDRESS),
-    };
-    let account_declare_tx = declare_tx(declare_tx_args);
-    let class_info = get_class_info_of_feature_contract(test_contract);
-    let tx =
-        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
-        arg1,
-        arg2,
-    ];
-    let contract_address_salt = ContractAddressSalt(Felt::ONE);
-    let calldata: Vec<_> =
-        [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,
-        &calldata,
-    );
-    let invoke_tx_args = invoke_tx_args! {
-        sender_address: *FUNDED_ACCOUNT_ADDRESS,
-        nonce: nonce_manager.next(*FUNDED_ACCOUNT_ADDRESS),
-        calldata: deploy_contract_calldata,
-        resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
-    };
-    let test_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);
+    let (mut test_manager, mut nonce_manager, [test_contract_address]) =
+        TestManager::<DictStateReader>::new_with_default_initial_state([(
+            FeatureContract::TestContract(CairoVersion::Cairo1(RunnableCairo1::Casm)),
+            calldata![Felt::ONE, Felt::TWO],
+        )])
+        .await;
 
     let key = Felt::from(10u8);
     let value = Felt::from(11u8);
@@ -280,15 +229,9 @@ async fn trivial_diff_scenario(
         )
         .await;
 
-    // Explicitly check the storage updates (not just inclusion) to make sure the key has no diff.
-    // Only storage update expected is due to the constructor call.
-    let expected_test_contract_storage_updates = HashMap::from([(
-        StarknetStorageKey(get_storage_var_address("my_storage_var", &[])),
-        StarknetStorageValue(arg1 + arg2),
-    )]);
-    assert_eq!(
-        test_output.decompressed_state_diff.storage_updates.get(&test_contract_address).unwrap(),
-        &expected_test_contract_storage_updates
+    // Explicitly check the test contract has no storage update.
+    assert!(
+        !test_output.decompressed_state_diff.storage_updates.contains_key(&test_contract_address)
     );
 
     let perform_global_validations = true;

crates/starknet_os_flow_tests/src/utils.rs

@@ -425,3 +425,10 @@ pub(crate) fn get_class_info_of_feature_contract(feature_contract: FeatureContra
         }
     }
 }
+
+pub(crate) fn get_class_hash_of_feature_contract(feature_contract: FeatureContract) -> ClassHash {
+    match feature_contract.get_class() {
+        ContractClass::V0(class) => ClassHash(compute_deprecated_class_hash(&class).unwrap()),
+        ContractClass::V1(_) => feature_contract.get_sierra().calculate_class_hash(),
+    }
+}