Use cheatable entrypoint in constructors

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - support for `keccak_syscall` syscall. It can be used directly in cairo tests
 - `l1_handler_execute` cheatcode
+- support for `roll`ing/`warp`ing/`prank`ing the constructor logic (precalculate address, prank, assert pranked state in constructor)
 
 ### Cast
 

crates/cheatnet/src/cheatcodes/deploy.rs

@@ -1,29 +1,24 @@
-use crate::constants::{
-    build_block_context, build_invoke_transaction, TEST_ACCOUNT_CONTRACT_ADDRESS,
-};
+use crate::constants::TEST_ACCOUNT_CONTRACT_ADDRESS;
 use crate::state::DictStateReader;
 use crate::{cheatcodes::EnhancedHintError, CheatnetState};
-use anyhow::{Context, Result};
+use anyhow::Result;
 use blockifier::abi::abi_utils::selector_from_name;
-use blockifier::execution::execution_utils::felt_to_stark_felt;
+use blockifier::execution::execution_utils::{felt_to_stark_felt, stark_felt_to_felt};
 
-use blockifier::execution::entry_point::CallInfo;
 use blockifier::state::cached_state::CachedState;
 use blockifier::state::state_api::StateReader;
-use blockifier::transaction::account_transaction::AccountTransaction;
-use blockifier::transaction::transactions::{ExecutableTransaction, InvokeTransaction};
 use cairo_felt::Felt252;
+use cairo_vm::vm::errors::hint_errors::HintError::CustomHint;
+use starknet::core::utils::get_selector_from_name;
 
 use starknet_api::core::{ClassHash, ContractAddress, EntryPointSelector, PatriciaKey};
 use starknet_api::hash::{StarkFelt, StarkHash};
-use starknet_api::transaction::{
-    Calldata, ContractAddressSalt, InvokeTransactionV1, TransactionHash,
-};
+use starknet_api::transaction::ContractAddressSalt;
 use starknet_api::{patricia_key, stark_felt};
 
 use super::CheatcodeError;
-use crate::conversions::felt_from_short_string;
-use crate::panic_data::try_extract_panic_data;
+use crate::conversions::{felt_from_short_string, field_element_to_felt252};
+use crate::rpc::{call_contract, CallContractOutput};
 
 impl CheatnetState {
     pub fn deploy(
@@ -33,9 +28,9 @@ impl CheatnetState {
     ) -> Result<ContractAddress, CheatcodeError> {
         // Deploy a contract using syscall deploy.
         let account_address = ContractAddress(patricia_key!(TEST_ACCOUNT_CONTRACT_ADDRESS));
-        let block_context = build_block_context();
         let entry_point_selector = selector_from_name("deploy_contract");
         let salt = self.get_salt();
+        let contract_address = self.precalculate_address(class_hash, calldata);
         self.increment_deploy_salt_base();
 
         let blockifier_state: &mut CachedState<DictStateReader> = &mut self.blockifier_state;
@@ -57,41 +52,23 @@ impl CheatnetState {
             &salt,
         );
 
-        let nonce = blockifier_state
-            .get_nonce_at(account_address)
-            .context("Failed to get nonce")
-            .map_err::<EnhancedHintError, _>(From::from)?;
-        let tx = build_invoke_transaction(execute_calldata, account_address);
-        let tx = InvokeTransactionV1 { nonce, ..tx };
-        let account_tx = AccountTransaction::Invoke(InvokeTransaction {
-            tx: starknet_api::transaction::InvokeTransaction::V1(tx),
-            tx_hash: TransactionHash::default(), // TODO(#358): Check if this is legit
-        });
-
-        let tx_info = account_tx
-            .execute(blockifier_state, &block_context, true, true)
-            .unwrap_or_else(|e| panic!("Unparseable transaction error: {e:?}"));
-
-        if let Some(CallInfo { execution, .. }) = tx_info.execute_call_info {
-            let contract_address = execution
-                .retdata
-                .0
-                .get(0)
-                .expect("Failed to get contract_address from return_data");
-
-            let contract_address = ContractAddress::try_from(*contract_address)
-                .expect("Failed to cast contract address into the right struct");
-
-            return Ok(contract_address);
+        let call_result = call_contract(
+            &account_address,
+            &field_element_to_felt252(&get_selector_from_name("__execute__").unwrap()),
+            execute_calldata.as_slice(),
+            self,
+        )
+        .unwrap_or_else(|err| panic!("Deploy txn failed: {err}"));
+
+        match call_result {
+            CallContractOutput::Success { .. } => Ok(contract_address),
+            CallContractOutput::Panic { panic_data } => {
+                Err(CheatcodeError::Recoverable(panic_data))
+            }
+            CallContractOutput::Error { msg } => Err(CheatcodeError::Unrecoverable(
+                EnhancedHintError::from(CustomHint(Box::from(msg))),
+            )),
         }
-
-        let revert_error = tx_info
-            .revert_error
-            .expect("Unparseable tx info, {tx_info:?}");
-        let extracted_panic_data = try_extract_panic_data(&revert_error)
-            .expect("Unparseable error message, {revert_error}");
-
-        Err(CheatcodeError::Recoverable(extracted_panic_data))
     }
 }
 
@@ -101,7 +78,7 @@ fn create_execute_calldata(
     account_address: &ContractAddress,
     entry_point_selector: &EntryPointSelector,
     salt: &ContractAddressSalt,
-) -> Calldata {
+) -> Vec<Felt252> {
     let calldata_len = u128::try_from(calldata.len()).unwrap();
     let mut execute_calldata = vec![
         *account_address.0.key(),      // Contract address.
@@ -113,13 +90,15 @@ fn create_execute_calldata(
     ];
     let mut calldata: Vec<StarkFelt> = calldata.iter().map(felt_to_stark_felt).collect();
     execute_calldata.append(&mut calldata);
-    Calldata(execute_calldata.into())
+    return execute_calldata
+        .iter()
+        .map(|sf| stark_felt_to_felt(*sf))
+        .collect();
 }
 
 #[cfg(test)]
 mod test {
     use super::*;
-    use std::sync::Arc;
 
     #[test]
     fn execute_calldata() {
@@ -132,16 +111,16 @@ mod test {
         );
         assert_eq!(
             calldata,
-            Calldata(Arc::new(vec![
-                StarkFelt::from(111_u32),
-                StarkFelt::from(222_u32),
-                StarkFelt::from(5_u32),
-                StarkFelt::from(123_u32),
-                StarkFelt::from(333_u32),
-                StarkFelt::from(2_u32),
-                StarkFelt::from(100_u32),
-                StarkFelt::from(200_u32),
-            ]))
+            vec![
+                Felt252::from(111_u32),
+                Felt252::from(222_u32),
+                Felt252::from(5_u32),
+                Felt252::from(123_u32),
+                Felt252::from(333_u32),
+                Felt252::from(2_u32),
+                Felt252::from(100_u32),
+                Felt252::from(200_u32),
+            ]
         );
     }
 
@@ -156,14 +135,14 @@ mod test {
         );
         assert_eq!(
             calldata,
-            Calldata(Arc::new(vec![
-                StarkFelt::from(111_u32),
-                StarkFelt::from(222_u32),
-                StarkFelt::from(3_u32),
-                StarkFelt::from(123_u32),
-                StarkFelt::from(333_u32),
-                StarkFelt::from(0_u32),
-            ]))
+            vec![
+                Felt252::from(111_u32),
+                Felt252::from(222_u32),
+                Felt252::from(3_u32),
+                Felt252::from(123_u32),
+                Felt252::from(333_u32),
+                Felt252::from(0_u32),
+            ]
         );
     }
 }

crates/cheatnet/src/constants.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 use std::{collections::HashMap, fs, path::PathBuf};
 
+use blockifier::execution::contract_class::ContractClassV1;
 use blockifier::state::cached_state::GlobalContractCache;
 use blockifier::{
     abi::{abi_utils::get_storage_var_address, constants},
@@ -112,16 +113,30 @@ pub fn build_invoke_transaction(
     }
 }
 
-fn load_contract_class(
+fn read_predeployed_contract_file(
     predeployed_contracts: &Utf8PathBuf,
     contract_path: &str,
-) -> ContractClassV0 {
+) -> String {
     let full_contract_path: PathBuf = predeployed_contracts.join(contract_path).into();
-    let raw_contract_class =
-        fs::read_to_string(full_contract_path).expect("Failed to read predeployed contracts");
+    fs::read_to_string(full_contract_path).expect("Failed to read predeployed contracts")
+}
+
+fn load_v0_contract_class(
+    predeployed_contracts: &Utf8PathBuf,
+    contract_path: &str,
+) -> ContractClassV0 {
+    let raw_contract_class = read_predeployed_contract_file(predeployed_contracts, contract_path);
     ContractClassV0::try_from_json_string(&raw_contract_class).unwrap()
 }
 
+fn load_v1_contract_class(
+    predeployed_contracts: &Utf8PathBuf,
+    contract_path: &str,
+) -> ContractClassV1 {
+    let raw_contract_class = read_predeployed_contract_file(predeployed_contracts, contract_path);
+    ContractClassV1::try_from_json_string(&raw_contract_class).unwrap()
+}
+
 fn erc20_account_balance_key() -> StorageKey {
     get_storage_var_address(
         "ERC20_balances",
@@ -135,11 +150,8 @@ fn erc20_account_balance_key() -> StorageKey {
 // Account does not include validations
 #[must_use]
 pub fn build_testing_state(predeployed_contracts: &Utf8PathBuf) -> CachedState<DictStateReader> {
-    let account_class = load_contract_class(
-        predeployed_contracts,
-        "account_no_validations_contract.casm.json",
-    );
-    let erc20_class = load_contract_class(
+    let account_class = load_v1_contract_class(predeployed_contracts, "account_cairo1.casm.json");
+    let erc20_class = load_v0_contract_class(
         predeployed_contracts,
         "erc20_contract_without_some_syscalls_compiled.json",
     );
@@ -148,7 +160,7 @@ pub fn build_testing_state(predeployed_contracts: &Utf8PathBuf) -> CachedState<D
     let test_erc20_class_hash = ClassHash(stark_felt!(TEST_ERC20_CONTRACT_CLASS_HASH));
 
     let class_hash_to_class = HashMap::from([
-        (test_account_class_hash, ContractClass::V0(account_class)),
+        (test_account_class_hash, ContractClass::V1(account_class)),
         (test_erc20_class_hash, ContractClass::V0(erc20_class)),
     ]);
 

crates/cheatnet/src/conversions.rs

@@ -1,7 +1,9 @@
+use blockifier::execution::execution_utils::felt_to_stark_felt;
 use blockifier::execution::execution_utils::stark_felt_to_felt;
 use cairo_felt::Felt252;
+use starknet::core::types::FieldElement;
 use starknet::core::utils::get_selector_from_name;
-use starknet_api::core::{ClassHash, ContractAddress};
+use starknet_api::core::{ClassHash, ContractAddress, PatriciaKey};
 
 #[must_use]
 pub fn felt_selector_from_name(name: &str) -> Felt252 {
@@ -24,11 +26,28 @@ pub fn class_hash_to_felt(class_hash: ClassHash) -> Felt252 {
     stark_felt_to_felt(class_hash.0)
 }
 
+#[must_use]
+pub fn class_hash_from_felt(felt: &Felt252) -> ClassHash {
+    ClassHash(felt_to_stark_felt(felt))
+}
+
+#[must_use]
+pub fn contract_address_from_felt(felt: &Felt252) -> ContractAddress {
+    ContractAddress(
+        PatriciaKey::try_from(felt_to_stark_felt(felt))
+            .expect("StarkFelt to PatriciaKey conversion failed"),
+    )
+}
+
+#[must_use]
+pub fn field_element_to_felt252(field_element: &FieldElement) -> Felt252 {
+    Felt252::from_bytes_be(&field_element.to_bytes_be())
+}
+
 #[cfg(test)]
 mod test {
-    use starknet_api::hash::StarkFelt;
-
     use super::*;
+    use starknet_api::hash::StarkFelt;
 
     #[test]
     fn parsing_felt_from_short_string() {