Added resource usage estimation (#605)

<!-- Reference any GitHub issues resolved by this PR -->

Closes #249
Closes #362 

## Introduced changes

<!-- A brief description of the changes -->

- added gas estimation

## Breaking changes

<!-- List of all breaking changes, if applicable -->

- `call_contract`, `deploy`, and `deploy_at` Cheatnet API changed

## Checklist

<!-- Make sure all of these are complete -->

- [x] Linked relevant issue
- [ ] Updated relevant documentation
- [x] Added relevant tests
- [x] Performed self-review of the code
- [ ] Added changes to `CHANGELOG.md`

---------

Co-authored-by: Arcticae <tomekgsd@gmail.com>

crates/cheatnet/src/cheatcodes/deploy.rs

@@ -3,6 +3,7 @@ use crate::{cheatcodes::EnhancedHintError, CheatnetState};
 use anyhow::Result;
 use blockifier::abi::abi_utils::selector_from_name;
 use blockifier::execution::execution_utils::{felt_to_stark_felt, stark_felt_to_felt};
+use blockifier::transaction::constants::EXECUTE_ENTRY_POINT_NAME;
 
 use blockifier::state::state_api::{State, StateReader};
 use cairo_felt::Felt252;
@@ -17,15 +18,22 @@ use starknet_api::transaction::ContractAddressSalt;
 use starknet_api::{patricia_key, stark_felt};
 
 use super::CheatcodeError;
-use crate::rpc::{call_contract, CallContractOutput};
+use crate::rpc::{call_contract, CallContractOutput, ResourceReport};
+
+#[allow(clippy::module_name_repetitions)]
+#[derive(Debug, Clone, PartialEq)]
+pub struct DeployPayload {
+    pub contract_address: ContractAddress,
+    pub resource_report: ResourceReport,
+}
 
 impl CheatnetState {
     pub fn deploy_at(
         &mut self,
         class_hash: &ClassHash,
         calldata: &[Felt252],
         contract_address: ContractAddress,
-    ) -> Result<ContractAddress, CheatcodeError> {
+    ) -> Result<DeployPayload, CheatcodeError> {
         let salt = self.get_salt();
         self.increment_deploy_salt_base();
 
@@ -51,23 +59,37 @@ impl CheatnetState {
 
         let call_result = call_contract(
             &account_address,
-            &get_selector_from_name("__execute__").unwrap().to_felt252(),
+            &get_selector_from_name(EXECUTE_ENTRY_POINT_NAME)
+                .unwrap()
+                .to_felt252(),
             execute_calldata.as_slice(),
             self,
         )
         .unwrap_or_else(|err| panic!("Deploy txn failed: {err}"));
 
         match call_result {
-            CallContractOutput::Success { .. } => self
-                .blockifier_state
-                .set_class_hash_at(contract_address, *class_hash)
-                .map(|_| contract_address)
-                .map_err(|msg| {
-                    CheatcodeError::Unrecoverable(EnhancedHintError::from(CustomHint(Box::from(
-                        msg.to_string(),
-                    ))))
-                }),
-            CallContractOutput::Panic { panic_data } => {
+            CallContractOutput::Success {
+                resource_report, ..
+            } => {
+                let result = self
+                    .blockifier_state
+                    .set_class_hash_at(contract_address, *class_hash)
+                    .map(|_| contract_address)
+                    .map_err(|msg| {
+                        CheatcodeError::Unrecoverable(EnhancedHintError::from(CustomHint(
+                            Box::from(msg.to_string()),
+                        )))
+                    });
+
+                match result {
+                    Ok(contract_address) => Ok(DeployPayload {
+                        contract_address,
+                        resource_report,
+                    }),
+                    Err(cheatcode_error) => Err(cheatcode_error),
+                }
+            }
+            CallContractOutput::Panic { panic_data, .. } => {
                 let panic_data_str = panic_data
                     .iter()
                     .map(|x| as_cairo_short_string(x).unwrap())
@@ -87,7 +109,7 @@ impl CheatnetState {
 
                 Err(CheatcodeError::Recoverable(panic_data))
             }
-            CallContractOutput::Error { msg } => Err(CheatcodeError::Unrecoverable(
+            CallContractOutput::Error { msg, .. } => Err(CheatcodeError::Unrecoverable(
                 EnhancedHintError::from(CustomHint(Box::from(msg))),
             )),
         }
@@ -97,7 +119,7 @@ impl CheatnetState {
         &mut self,
         class_hash: &ClassHash,
         calldata: &[Felt252],
-    ) -> Result<ContractAddress, CheatcodeError> {
+    ) -> Result<DeployPayload, CheatcodeError> {
         let contract_address = self.precalculate_address(class_hash, calldata);
 
         self.deploy_at(class_hash, calldata, contract_address)

crates/cheatnet/src/constants.rs

@@ -13,8 +13,8 @@ use blockifier::{
 };
 use cairo_felt::Felt252;
 use cairo_vm::vm::runners::builtin_runner::{
-    BITWISE_BUILTIN_NAME, EC_OP_BUILTIN_NAME, HASH_BUILTIN_NAME, OUTPUT_BUILTIN_NAME,
-    POSEIDON_BUILTIN_NAME, RANGE_CHECK_BUILTIN_NAME, SIGNATURE_BUILTIN_NAME,
+    BITWISE_BUILTIN_NAME, EC_OP_BUILTIN_NAME, HASH_BUILTIN_NAME, KECCAK_BUILTIN_NAME,
+    OUTPUT_BUILTIN_NAME, POSEIDON_BUILTIN_NAME, RANGE_CHECK_BUILTIN_NAME, SIGNATURE_BUILTIN_NAME,
 };
 use camino::Utf8PathBuf;
 use starknet_api::{
@@ -45,18 +45,48 @@ pub const TEST_FAULTY_ACCOUNT_CONTRACT_CLASS_HASH: &str = "0x113";
 pub const SECURITY_TEST_CLASS_HASH: &str = "0x114";
 pub const TEST_ERC20_CONTRACT_CLASS_HASH: &str = "0x1010";
 
+pub const STEP_RESOURCE_COST: f64 = 0.01_f64;
+
+// HOW TO FIND:
+// 1. https://docs.starknet.io/documentation/architecture_and_concepts/Network_Architecture/fee-mechanism/#general_case
+// 2. src/starkware/cairo/lang/instances.py::starknet_with_keccak_instance
 #[must_use]
 pub fn build_block_context() -> BlockContext {
     // blockifier::test_utils::create_for_account_testing
     let vm_resource_fee_cost = Arc::new(HashMap::from([
-        (constants::N_STEPS_RESOURCE.to_string(), 1_f64),
-        (HASH_BUILTIN_NAME.to_string(), 1_f64),
-        (RANGE_CHECK_BUILTIN_NAME.to_string(), 1_f64),
-        (SIGNATURE_BUILTIN_NAME.to_string(), 1_f64),
-        (BITWISE_BUILTIN_NAME.to_string(), 1_f64),
-        (POSEIDON_BUILTIN_NAME.to_string(), 1_f64),
-        (OUTPUT_BUILTIN_NAME.to_string(), 1_f64),
-        (EC_OP_BUILTIN_NAME.to_string(), 1_f64),
+        (constants::N_STEPS_RESOURCE.to_string(), STEP_RESOURCE_COST),
+        (HASH_BUILTIN_NAME.to_string(), 32_f64 * STEP_RESOURCE_COST),
+        (
+            RANGE_CHECK_BUILTIN_NAME.to_string(),
+            16_f64 * STEP_RESOURCE_COST,
+        ),
+        (
+            SIGNATURE_BUILTIN_NAME.to_string(),
+            2048_f64 * STEP_RESOURCE_COST,
+        ), // ECDSA
+        (
+            BITWISE_BUILTIN_NAME.to_string(),
+            64_f64 * STEP_RESOURCE_COST,
+        ),
+        (
+            POSEIDON_BUILTIN_NAME.to_string(),
+            32_f64 * STEP_RESOURCE_COST,
+        ),
+        (OUTPUT_BUILTIN_NAME.to_string(), 0_f64 * STEP_RESOURCE_COST),
+        (
+            EC_OP_BUILTIN_NAME.to_string(),
+            1024_f64 * STEP_RESOURCE_COST,
+        ),
+        (
+            KECCAK_BUILTIN_NAME.to_string(),
+            2048_f64 * STEP_RESOURCE_COST, // 2**11
+        ),
+        // The gas estimation should panic in case it encounters a builtin that doesn't have a cost
+        // This builtin seems to be unused for cost estimation
+        // (
+        //     SEGMENT_ARENA_BUILTIN_NAME.to_string(),
+        //     0_f64 * STEP_RESOURCE_COST,
+        // ), // BUILTIN COST NOT FOUND
     ]));
 
     BlockContext {

crates/cheatnet/src/execution/gas.rs

@@ -0,0 +1,30 @@
+use std::collections::HashMap;
+
+use blockifier::{
+    abi::constants, block_context::BlockContext, execution::entry_point::ExecutionResources,
+    fee::fee_utils::calculate_l1_gas_by_vm_usage, transaction::objects::ResourcesMapping,
+};
+use cairo_vm::vm::runners::cairo_runner::ExecutionResources as VmExecutionResources;
+
+#[allow(clippy::module_name_repetitions)]
+#[must_use]
+pub fn gas_from_execution_resources(
+    block_context: &BlockContext,
+    resources: &ExecutionResources,
+) -> f64 {
+    let resource_mapping = vm_execution_resources_to_resource_mapping(&resources.vm_resources);
+    calculate_l1_gas_by_vm_usage(block_context, &resource_mapping)
+        .expect("Calculating gas failed, some resources were not included.")
+}
+
+#[must_use]
+fn vm_execution_resources_to_resource_mapping(
+    execution_resources: &VmExecutionResources,
+) -> ResourcesMapping {
+    let mut map = HashMap::from([(
+        constants::N_STEPS_RESOURCE.to_string(),
+        execution_resources.n_steps,
+    )]);
+    map.extend(execution_resources.builtin_instance_counter.clone());
+    ResourcesMapping(map)
+}

crates/cheatnet/src/execution/mod.rs

@@ -6,4 +6,5 @@ pub mod entry_point;
 pub mod events;
 #[allow(clippy::module_name_repetitions)]
 pub mod execution_info;
+pub mod gas;
 pub mod syscalls;

crates/cheatnet/src/rpc.rs

@@ -1,11 +1,13 @@
 use anyhow::Result;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use crate::panic_data::try_extract_panic_data;
 use crate::{
     constants::{build_block_context, build_transaction_context},
     execution::{
         entry_point::execute_call_entry_point, events::collect_emitted_events_from_spied_contracts,
+        gas::gas_from_execution_resources,
     },
     CheatnetState,
 };
@@ -22,14 +24,42 @@ use starknet_api::{
     transaction::Calldata,
 };
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct ResourceReport {
+    pub gas: f64,
+    pub steps: usize,
+    pub bultins: HashMap<String, usize>,
+}
+
+impl ResourceReport {
+    fn new(gas: f64, resources: &ExecutionResources) -> Self {
+        Self {
+            gas,
+            steps: resources.vm_resources.n_steps,
+            bultins: resources.vm_resources.builtin_instance_counter.clone(),
+        }
+    }
+}
+
+#[derive(Debug)]
 pub enum CallContractOutput {
-    Success { ret_data: Vec<Felt252> },
-    Panic { panic_data: Vec<Felt252> },
-    Error { msg: String },
+    Success {
+        ret_data: Vec<Felt252>,
+        resource_report: ResourceReport,
+    },
+    Panic {
+        panic_data: Vec<Felt252>,
+        resource_report: ResourceReport,
+    },
+    Error {
+        msg: String,
+        resource_report: ResourceReport,
+    },
 }
 
 // This does contract call without the transaction layer. This way `call_contract` can return data and modify state.
 // `call` and `invoke` on the transactional layer use such method under the hood.
+#[allow(clippy::too_many_lines)]
 pub fn call_contract(
     contract_address: &ContractAddress,
     entry_point_selector: &Felt252,
@@ -78,6 +108,9 @@ pub fn call_contract(
         &mut context,
     );
 
+    let gas = gas_from_execution_resources(&block_context, &resources);
+    let resource_report = ResourceReport::new(gas, &resources);
+
     match exec_result {
         Ok(call_info) => {
             if !cheatcode_state.spies.is_empty() {
@@ -95,6 +128,7 @@ pub fn call_contract(
 
             Ok(CallContractOutput::Success {
                 ret_data: return_data,
+                resource_report,
             })
         }
         Err(EntryPointExecutionError::ExecutionFailed { error_data }) => {
@@ -105,13 +139,20 @@ pub fn call_contract(
 
             Ok(CallContractOutput::Panic {
                 panic_data: err_data,
+                resource_report,
             })
         }
         Err(EntryPointExecutionError::VirtualMachineExecutionErrorWithTrace { trace, .. }) => {
             if let Some(panic_data) = try_extract_panic_data(&trace) {
-                Ok(CallContractOutput::Panic { panic_data })
+                Ok(CallContractOutput::Panic {
+                    panic_data,
+                    resource_report,
+                })
             } else {
-                Ok(CallContractOutput::Error { msg: trace })
+                Ok(CallContractOutput::Error {
+                    msg: trace,
+                    resource_report,
+                })
             }
         }
         Err(EntryPointExecutionError::PreExecutionError(
@@ -122,17 +163,26 @@ pub fn call_contract(
             let msg = format!(
                 "Entry point selector {selector_hash} not found in contract {contract_addr}"
             );
-            Ok(CallContractOutput::Error { msg })
+            Ok(CallContractOutput::Error {
+                msg,
+                resource_report,
+            })
         }
         Err(EntryPointExecutionError::PreExecutionError(
             PreExecutionError::UninitializedStorageAddress(contract_address),
         )) => {
             let address = contract_address.0.key().to_string();
             let msg = format!("Contract not deployed at address: {address}");
-            Ok(CallContractOutput::Error { msg })
+            Ok(CallContractOutput::Error {
+                msg,
+                resource_report,
+            })
         }
         Err(EntryPointExecutionError::StateError(StateError::StateReadError(msg))) => {
-            Ok(CallContractOutput::Error { msg })
+            Ok(CallContractOutput::Error {
+                msg,
+                resource_report,
+            })
         }
         result => panic!("Unparseable result: {result:?}"),
     }

crates/cheatnet/tests/cheatcodes/deploy.rs

@@ -18,7 +18,8 @@ fn deploy_at_predefined_address() {
     let class_hash = state.declare(&contract, &contracts).unwrap();
     let contract_address = state
         .deploy_at(&class_hash, &[], ContractAddress::from(1_u8))
-        .unwrap();
+        .unwrap()
+        .contract_address;
 
     assert_eq!(contract_address, ContractAddress::from(1_u8));
 
@@ -59,7 +60,8 @@ fn call_predefined_contract_from_proxy_contract() {
     let class_hash = state.declare(&contract, &contracts).unwrap();
     let prank_checker_address = state
         .deploy_at(&class_hash, &[], ContractAddress::from(1_u8))
-        .unwrap();
+        .unwrap()
+        .contract_address;
 
     assert_eq!(prank_checker_address, ContractAddress::from(1_u8));
 
@@ -92,7 +94,7 @@ fn deploy_contract_on_predefined_address_after_its_usage() {
     .unwrap();
 
     assert!(match output {
-        CallContractOutput::Error { msg } =>
+        CallContractOutput::Error { msg, .. } =>
             msg.contains("Requested contract address") && msg.contains("is not deployed"),
         _ => false,
     });
@@ -163,8 +165,6 @@ fn deploy_missing_arguments_in_constructor() {
 
     let output = state.deploy(&class_hash, &[Felt252::from(123_321)]);
 
-    dbg!(&output);
-
     assert!(match output {
         Err(CheatcodeError::Unrecoverable(EnhancedHintError::Hint(HintError::CustomHint(msg)))) =>
             msg.as_ref() == "Failed to deserialize param #2",

crates/cheatnet/tests/cheatcodes/get_class_hash.rs

@@ -9,7 +9,7 @@ fn get_class_hash_simple() {
     let contracts = get_contracts();
     let contract_name = "HelloStarknet".to_owned().to_felt252();
     let class_hash = state.declare(&contract_name, &contracts).unwrap();
-    let contract_address = state.deploy(&class_hash, &[]).unwrap();
+    let contract_address = state.deploy(&class_hash, &[]).unwrap().contract_address;
 
     assert_eq!(class_hash, state.get_class_hash(contract_address).unwrap());
 }
@@ -21,7 +21,7 @@ fn get_class_hash_upgrade() {
     let contracts = get_contracts();
     let contract_name = "GetClassHashCheckerUpg".to_owned().to_felt252();
     let class_hash = state.declare(&contract_name, &contracts).unwrap();
-    let contract_address = state.deploy(&class_hash, &[]).unwrap();
+    let contract_address = state.deploy(&class_hash, &[]).unwrap().contract_address;
 
     assert_eq!(class_hash, state.get_class_hash(contract_address).unwrap());