Precalculate contract address function

benchmarks/data/forge.cairo

@@ -4,17 +4,17 @@ use option::OptionTrait;
 use traits::TryInto;
 use starknet::ContractAddress;
 use starknet::Felt252TryIntoContractAddress;
-use cheatcodes::PreparedContract;
 
 use project::IHelloStarknetSafeDispatcher;
 use project::IHelloStarknetSafeDispatcherTrait;
 
+use cheatcodes::ContractClass;
+use cheatcodes::ContractClassTrait;
+
 fn deploy_hello_starknet() -> ContractAddress {
-    let class_hash = declare('HelloStarknet').unwrap();
-    let prepared = PreparedContract {
-        contract_address: 1234, class_hash: class_hash, constructor_calldata: @ArrayTrait::new()
-    };
-    let contract_address = deploy(prepared).unwrap();
+    let contract = declare('HelloStarknet').unwrap();
+    let constructor_calldata = @ArrayTrait::new()
+    let contract_address = contract.deploy(constructor_calldata).unwrap();
 
     let contract_address: ContractAddress = contract_address.try_into().unwrap();
 

benchmarks/data/protostar.cairo

@@ -4,14 +4,11 @@ use option::OptionTrait;
 use traits::TryInto;
 use starknet::ContractAddress;
 use starknet::Felt252TryIntoContractAddress;
-use cheatcodes::PreparedContract;
 
 fn deploy_hello_starknet() -> felt252 {
-    let class_hash = declare('HelloStarknet').unwrap();
-    let prepared = PreparedContract {
-        contract_address: 1234, class_hash: class_hash, constructor_calldata: @ArrayTrait::new()
-    };
-    let contract_address = deploy(prepared).unwrap();
+    let contract = declare('HelloStarknet');
+
+    let contract_address = contract.deploy(@ArrayTrait::new()).unwrap();
 
     contract_address
 }

cheatcodes/src/cheatcodes.cairo

@@ -14,16 +14,20 @@ use starknet::Felt252TryIntoClassHash;
 use starknet::Felt252TryIntoContractAddress;
 
 #[derive(Drop, Clone)]
-struct PreparedContract {
+struct ContractClass {
     class_hash: ClassHash,
-    constructor_calldata: @Array::<felt252>,
 }
 
 #[derive(Drop, Clone)]
 struct RevertedTransaction {
     panic_data: Array::<felt252>,
 }
 
+trait ContractClassTrait {
+    fn precalculate_address(self: @ContractClass, constructor_calldata: @Array::<felt252>) -> ContractAddress;
+    fn deploy(self: @ContractClass, constructor_calldata: @Array::<felt252>) -> Result<ContractAddress, RevertedTransaction>;
+}
+
 trait RevertedTransactionTrait {
     fn first(self: @RevertedTransaction) -> felt252;
 }
@@ -34,54 +38,73 @@ impl RevertedTransactionImpl of RevertedTransactionTrait {
     }
 }
 
-fn declare(contract: felt252) -> ClassHash {
-    let span = cheatcode::<'declare'>(array![contract].span());
+impl ContractClassImpl of ContractClassTrait {
+    fn precalculate_address(self: @ContractClass, constructor_calldata: @Array::<felt252>) -> ContractAddress {
+        let mut inputs: Array::<felt252> = _prepare_calldata(self.class_hash, constructor_calldata);
 
-    let exit_code = *span[0];
-    let result = *span[1];
-    assert(exit_code == 0, 'declare should never fail');
-    result.try_into().unwrap()
-}
+        let outputs = cheatcode::<'precalculate_address'>(inputs.span());
+        (*outputs[0]).try_into().unwrap()
+    }
 
-fn deploy(prepared_contract: PreparedContract) -> Result::<ContractAddress, RevertedTransaction> {
-    let PreparedContract{class_hash, constructor_calldata } = prepared_contract;
-    let mut inputs = array![class_hash.into()];
+    fn deploy(self: @ContractClass, constructor_calldata: @Array::<felt252>) -> Result<ContractAddress, RevertedTransaction> {
+        let mut inputs = _prepare_calldata(self.class_hash, constructor_calldata);
+
+        let outputs = cheatcode::<'deploy'>(inputs.span());
+        let exit_code = *outputs[0];
+
+        if exit_code == 0 {
+            let result = *outputs[1];
+            Result::<ContractAddress, RevertedTransaction>::Ok(result.try_into().unwrap())
+        } else {
+            let panic_data_len_felt = *outputs[1];
+            let panic_data_len = panic_data_len_felt.try_into().unwrap();
+            let mut panic_data = array![];
+
+            let offset = 2;
+            let mut i = offset;
+            loop {
+                if panic_data_len + offset == i {
+                    break ();
+                }
+                panic_data.append(*outputs[i]);
+                i += 1;
+            };
+
+            Result::<ContractAddress, RevertedTransaction>::Err(RevertedTransaction { panic_data })
+        }
+    }
+}
 
+fn _prepare_calldata(class_hash: @ClassHash, constructor_calldata: @Array::<felt252>) -> Array::<felt252>  {
+    let class_hash: felt252 = class_hash.clone().into();
+    let mut inputs: Array::<felt252> = array![class_hash];
     let calldata_len_felt = constructor_calldata.len().into();
     inputs.append(calldata_len_felt);
 
     let calldata_len = constructor_calldata.len();
     let mut i = 0;
+
     loop {
-        if calldata_len == i {
-            break ();
+        if i == calldata_len {
+            break();
         }
         inputs.append(*constructor_calldata[i]);
         i += 1;
     };
 
-    let outputs = cheatcode::<'deploy'>(inputs.span());
-    let exit_code = *outputs[0];
-
-    if exit_code == 0 {
-        let result = *outputs[1];
-        Result::<ContractAddress, RevertedTransaction>::Ok(result.try_into().unwrap())
-    } else {
-        let panic_data_len_felt = *outputs[1];
-        let panic_data_len = panic_data_len_felt.try_into().unwrap();
-        let mut panic_data = array![];
-
-        let offset = 2;
-        let mut i = offset;
-        loop {
-            if panic_data_len + offset == i {
-                break ();
-            }
-            panic_data.append(*outputs[i]);
-            i += 1;
-        };
-
-        Result::<ContractAddress, RevertedTransaction>::Err(RevertedTransaction { panic_data })
+    inputs
+}
+
+fn declare(contract: felt252) -> ContractClass {
+    let span = cheatcode::<'declare'>(array![contract].span());
+
+    let exit_code = *span[0];
+    let result = *span[1];
+    assert(exit_code == 0, 'declare should never fail');
+    let class_hash = result.try_into().unwrap();
+
+    ContractClass {
+        class_hash
     }
 }
 

cheatcodes/src/lib.cairo

@@ -1,11 +1,11 @@
 mod cheatcodes;
 
 use cheatcodes::declare;
-use cheatcodes::deploy;
 
-use cheatcodes::PreparedContract;
 use cheatcodes::RevertedTransaction;
 use cheatcodes::RevertedTransactionTrait;
+use cheatcodes::ContractClass;
+use cheatcodes::ContractClassTrait;
 
 use cheatcodes::start_prank;
 use cheatcodes::stop_prank;

crates/cheatable-starknet/src/cheatcodes.rs

@@ -12,6 +12,7 @@ use thiserror::Error;
 pub mod declare;
 pub mod deploy;
 pub mod prank;
+pub mod precalculate_address;
 pub mod roll;
 pub mod warp;
 

crates/cheatable-starknet/src/cheatcodes/deploy.rs

@@ -8,6 +8,7 @@ use crate::{
 };
 use anyhow::{Context, Result};
 use blockifier::abi::abi_utils::selector_from_name;
+use blockifier::execution::execution_utils::felt_to_stark_felt;
 use num_traits::cast::ToPrimitive;
 
 use blockifier::execution::entry_point::CallInfo;
@@ -28,7 +29,7 @@ use cairo_lang_runner::casm_run::MemBuffer;
 
 impl CheatedState {
     pub fn deploy(
-        &self,
+        &mut self,
         buffer: &mut MemBuffer,
         blockifier_state: &mut CachedState<DictStateReader>,
         inputs: &[Felt252],
@@ -38,18 +39,16 @@ impl CheatedState {
         let class_hash = inputs[0].clone();
 
         let calldata_length = inputs[1].to_usize().unwrap();
-        let mut calldata = vec![];
-        for felt in inputs.iter().skip(2).take(calldata_length) {
-            calldata.push(felt.clone());
-        }
+
+        let calldata = Vec::from(&inputs[2..(2 + calldata_length)]);
 
         // 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 = ContractAddressSalt::default();
+        let salt = self.get_salt();
         let class_hash = ClassHash(StarkFelt::new(class_hash.to_be_bytes()).unwrap());
-
+        self.increment_deploy_salt_base();
         let contract_class = blockifier_state.get_compiled_contract_class(&class_hash)?;
         if contract_class.constructor_selector().is_none() && !calldata.is_empty() {
             write_cheatcode_panic(
@@ -126,7 +125,7 @@ fn create_execute_calldata(
     ];
     let mut calldata: Vec<StarkFelt> = calldata
         .iter()
-        .map(|data| StarkFelt::new(data.to_be_bytes()).unwrap())
+        .map(|data| felt_to_stark_felt(data))
         .collect();
     execute_calldata.append(&mut calldata);
     Calldata(execute_calldata.into())

crates/cheatable-starknet/src/cheatcodes/precalculate_address.rs

@@ -0,0 +1,52 @@
+use crate::constants::TEST_ACCOUNT_CONTRACT_ADDRESS;
+
+use crate::{cheatcodes::EnhancedHintError, CheatedState};
+use blockifier::execution::execution_utils::felt_to_stark_felt;
+use cairo_felt::Felt252;
+use cairo_lang_runner::casm_run::MemBuffer;
+use num_traits::cast::ToPrimitive;
+use starknet_api::core::PatriciaKey;
+use starknet_api::core::{calculate_contract_address, ClassHash, ContractAddress};
+use starknet_api::hash::{StarkFelt, StarkHash};
+use starknet_api::transaction::Calldata;
+
+use starknet_api::patricia_key;
+
+impl CheatedState {
+    pub fn precalculate_address(
+        &self,
+        buffer: &mut MemBuffer,
+        inputs: &[Felt252],
+    ) -> Result<(), EnhancedHintError> {
+        let class_hash = inputs[0].clone();
+
+        let account_address = ContractAddress(patricia_key!(TEST_ACCOUNT_CONTRACT_ADDRESS));
+        let class_hash = ClassHash(StarkFelt::new(class_hash.to_be_bytes()).unwrap());
+        let salt = self.get_salt();
+        let calldata_length = inputs[1].to_usize().unwrap();
+
+        let calldata = Vec::from(&inputs[2..(2 + calldata_length)]);
+
+        let execute_calldata = create_execute_calldata(&calldata);
+        let contract_address =
+            calculate_contract_address(salt, class_hash, &execute_calldata, account_address)
+                .unwrap();
+
+        let contract_address = Felt252::from_bytes_be((*contract_address.0.key()).bytes());
+
+        buffer
+            .write(contract_address)
+            .expect("Failed to insert declared contract class hash");
+        Ok(())
+    }
+}
+
+fn create_execute_calldata(calldata: &[Felt252]) -> Calldata {
+    let mut execute_calldata = vec![];
+    let mut calldata: Vec<StarkFelt> = calldata
+        .iter()
+        .map(|data| felt_to_stark_felt(data))
+        .collect();
+    execute_calldata.append(&mut calldata);
+    Calldata(execute_calldata.into())
+}

crates/cheatable-starknet/src/lib.rs

@@ -1,5 +1,7 @@
 use cairo_felt::Felt252;
 use starknet_api::core::ContractAddress;
+use starknet_api::hash::StarkFelt;
+use starknet_api::transaction::ContractAddressSalt;
 use std::collections::HashMap;
 
 pub mod cheatcodes;
@@ -11,6 +13,7 @@ pub struct CheatedState {
     pub rolled_contracts: HashMap<ContractAddress, Felt252>,
     pub pranked_contracts: HashMap<ContractAddress, ContractAddress>,
     pub warped_contracts: HashMap<ContractAddress, Felt252>,
+    pub deploy_salt_base: u32,
 }
 
 impl CheatedState {
@@ -20,8 +23,17 @@ impl CheatedState {
             rolled_contracts: HashMap::new(),
             pranked_contracts: HashMap::new(),
             warped_contracts: HashMap::new(),
+            deploy_salt_base: 0,
         }
     }
+
+    pub fn increment_deploy_salt_base(&mut self) -> () {
+        self.deploy_salt_base += 1;
+    }
+
+    pub fn get_salt(&self) -> ContractAddressSalt {
+        ContractAddressSalt(StarkFelt::from(self.deploy_salt_base))
+    }
 }
 
 impl Default for CheatedState {

crates/forge/src/cheatcodes_hint_processor.rs

@@ -236,6 +236,9 @@ impl CairoHintProcessor<'_> {
                 print(inputs);
                 Ok(())
             }
+            "precalculate_address" => self
+                .cheated_state
+                .precalculate_address(&mut buffer, &inputs),
             _ => Err(anyhow!("Unknown cheatcode selector: {selector}")).map_err(Into::into),
         }?;
 

crates/forge/tests/data/contracts/constructor_simple.cairo

@@ -0,0 +1,25 @@
+#[starknet::contract]
+mod ConstructorSimple {
+    #[storage]
+    struct Storage {
+        number: felt252
+    }
+
+    #[constructor]
+    fn constructor(ref self: ContractState, number: felt252) {
+        self.number.write(number);
+    }
+
+    #[external(v0)]
+    fn add_to_number(ref self: ContractState, number: felt252) -> felt252 {
+        let new_number = self.number.read() + number;
+        self.number.write(new_number);
+        new_number
+    }
+
+    #[external(v0)]
+    fn get_number(self: @ContractState) -> felt252 {
+        self.number.read()
+    }
+}
+

crates/forge/tests/data/simple_package/tests/contract.cairo

@@ -5,16 +5,16 @@ use traits::TryInto;
 use starknet::ContractAddress;
 use starknet::Felt252TryIntoContractAddress;
 
-use cheatcodes::{ declare, PreparedContract, deploy };
+use cheatcodes::{ declare, ContractClassTrait };
 
 use simple_package::hello_starknet::IHelloStarknetDispatcher;
 use simple_package::hello_starknet::IHelloStarknetDispatcherTrait;
 
 #[test]
 fn call_and_invoke() {
-    let class_hash = declare('HelloStarknet');
-    let prepared = PreparedContract { class_hash: class_hash, constructor_calldata: @ArrayTrait::new() };
-    let contract_address = deploy(prepared).unwrap();
+    let contract = declare('HelloStarknet');
+    let constructor_calldata = @ArrayTrait::new();
+    let contract_address = contract.deploy(constructor_calldata).unwrap();
     let dispatcher = IHelloStarknetDispatcher { contract_address };
 
     let balance = dispatcher.get_balance();