Fix gas calculation (#1623)

## Introduced changes

- now gas used for constructor execution and l1 handler execution is
included in total gas cost
- `DeployPayload` and `CallOutput` struct were removed (no longer
necessary)

## Checklist

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

Author: piotmag769

CHANGELOG.md

@@ -13,6 +13,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Gas estimation is now aligned with the Starknet v0.13
 
+### Fixed
+
+- Gas used in constructors and handling of L1 messages is now properly included in total gas cost
+
 ## [0.16.0] - 2024-01-26
 
 ### Forge

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/mod.rs

@@ -30,7 +30,7 @@ use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{
 };
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::{
     execution::cheated_syscalls::SingleSegmentResponse,
-    rpc::{CallFailure, CallOutput, CallResult},
+    rpc::{CallFailure, CallResult},
 };
 
 use super::io_runtime_extension::IORuntime;
@@ -81,15 +81,15 @@ where
         self,
         blockifier_state: &mut BlockifierState,
         cheatnet_state: &mut CheatnetState,
-    ) -> CallOutput;
+    ) -> CallResult;
 }
 
 impl ExecuteCall for CallContractRequest {
     fn execute_call(
         self: CallContractRequest,
         blockifier_state: &mut BlockifierState,
         cheatnet_state: &mut CheatnetState,
-    ) -> CallOutput {
+    ) -> CallResult {
         let contract_address = self.contract_address;
 
         let entry_point = CallEntryPoint {
@@ -110,7 +110,6 @@ impl ExecuteCall for CallContractRequest {
             entry_point,
             &AddressOrClassHash::ContractAddress(contract_address),
         )
-        .unwrap_or_else(|err| panic!("Transaction execution error: {err}"))
     }
 }
 
@@ -119,7 +118,7 @@ impl ExecuteCall for LibraryCallRequest {
         self: LibraryCallRequest,
         blockifier_state: &mut BlockifierState,
         cheatnet_state: &mut CheatnetState,
-    ) -> CallOutput {
+    ) -> CallResult {
         let class_hash = self.class_hash;
 
         let entry_point = CallEntryPoint {
@@ -140,7 +139,6 @@ impl ExecuteCall for LibraryCallRequest {
             entry_point,
             &AddressOrClassHash::ClassHash(class_hash),
         )
-        .unwrap_or_else(|err| panic!("Transaction execution error: {err}"))
     }
 }
 
@@ -161,34 +159,22 @@ fn execute_syscall<Request: ExecuteCall + SyscallRequest>(
     let mut blockifier_state = BlockifierState::from(syscall_handler.state);
 
     let call_result = request.execute_call(&mut blockifier_state, cheatnet_state);
-    write_call_response(
-        syscall_handler,
-        cheatnet_state,
-        vm,
-        gas_counter,
-        call_result,
-    )?;
+    write_call_response(syscall_handler, vm, gas_counter, call_result)?;
     Ok(())
 }
 
 fn write_call_response(
     syscall_handler: &mut SyscallHintProcessor<'_>,
-    cheatnet_state: &mut CheatnetState,
     vm: &mut VirtualMachine,
     gas_counter: u64,
-    call_output: CallOutput,
+    call_result: CallResult,
 ) -> Result<(), HintError> {
-    let response_wrapper: SyscallResponseWrapper<SingleSegmentResponse> = match call_output.result {
+    let response_wrapper: SyscallResponseWrapper<SingleSegmentResponse> = match call_result {
         CallResult::Success { ret_data, .. } => {
             let memory_segment_start_ptr = syscall_handler
                 .read_only_segments
                 .allocate(vm, &ret_data.iter().map(Into::into).collect())?;
 
-            // add execution resources used by call to all used resources
-            cheatnet_state
-                .used_resources
-                .extend(&call_output.used_resources);
-
             SyscallResponseWrapper::Success {
                 gas_counter,
                 response: SingleSegmentResponse {

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/rpc.rs

@@ -1,4 +1,3 @@
-use anyhow::Result;
 use blockifier::execution::execution_utils::stark_felt_to_felt;
 use cairo_lang_runner::casm_run::format_next_item;
 
@@ -39,13 +38,6 @@ impl UsedResources {
     }
 }
 
-/// Represents call output, along with the data and the resources consumed during execution
-#[derive(Debug)]
-pub struct CallOutput {
-    pub result: CallResult,
-    pub used_resources: UsedResources,
-}
-
 /// Enum representing possible call execution result, along with the data
 #[derive(Debug)]
 pub enum CallResult {
@@ -178,7 +170,7 @@ pub fn call_l1_handler(
     contract_address: &ContractAddress,
     entry_point_selector: &Felt252,
     calldata: &[Felt252],
-) -> Result<CallOutput> {
+) -> CallResult {
     let entry_point_selector = create_entry_point_selector(entry_point_selector);
     let calldata = create_execute_calldata(calldata);
 
@@ -207,7 +199,7 @@ pub fn call_entry_point(
     cheatnet_state: &mut CheatnetState,
     mut entry_point: CallEntryPoint,
     starknet_identifier: &AddressOrClassHash,
-) -> Result<CallOutput> {
+) -> CallResult {
     let mut resources = ExecutionResources::default();
     let account_context = build_transaction_context();
     let block_context = build_block_context(cheatnet_state.block_info);
@@ -228,15 +220,16 @@ pub fn call_entry_point(
         &mut context,
     );
 
-    let result = CallResult::from_execution_result(&exec_result, starknet_identifier);
-
-    Ok(CallOutput {
-        result,
-        used_resources: UsedResources {
-            execution_resources: resources,
-            l2_to_l1_payloads_length: exec_result.map_or(vec![], |call_info| {
-                call_info.get_sorted_l2_to_l1_payloads_length().unwrap()
-            }),
-        },
-    })
+    let call_result = CallResult::from_execution_result(&exec_result, starknet_identifier);
+
+    let used_resources = UsedResources {
+        execution_resources: resources,
+        l2_to_l1_payloads_length: exec_result.map_or(vec![], |call_info| {
+            call_info.get_sorted_l2_to_l1_payloads_length().unwrap()
+        }),
+    };
+    // add execution resources used by call contract, library call or l1 handler execution to all used resources
+    cheatnet_state.used_resources.extend(&used_resources);
+
+    call_result
 }

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/deploy.rs

@@ -1,6 +1,6 @@
 use crate::constants::TEST_ADDRESS;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{
-    AddressOrClassHash, CallFailure,
+    AddressOrClassHash, CallFailure, UsedResources,
 };
 use crate::state::BlockifierState;
 use crate::CheatnetState;
@@ -27,19 +27,13 @@ use starknet_api::transaction::Calldata;
 
 use super::CheatcodeError;
 
-#[derive(Debug)]
-pub struct DeployCallPayload {
-    pub used_resources: ExecutionResources,
-    pub contract_address: ContractAddress,
-}
-
 pub fn deploy_at(
     blockifier_state: &mut BlockifierState,
     cheatnet_state: &mut CheatnetState,
     class_hash: &ClassHash,
     calldata: &[Felt252],
     contract_address: ContractAddress,
-) -> Result<DeployCallPayload, CheatcodeError> {
+) -> Result<ContractAddress, CheatcodeError> {
     let blockifier_state_raw: &mut dyn State = blockifier_state.blockifier_state;
 
     if let Ok(class_hash) = blockifier_state_raw.get_class_hash_at(contract_address) {
@@ -69,20 +63,17 @@ pub fn deploy_at(
         calldata.to_vec().iter().map(felt_to_stark_felt).collect(),
     ));
 
-    let resources = &mut ExecutionResources::default();
+    let mut resources = ExecutionResources::default();
     let result = cheated_syscalls::execute_deployment(
         blockifier_state_raw,
-        resources,
+        &mut resources,
         entry_point_execution_ctx,
         ctor_context,
         calldata,
         u64::MAX,
         cheatnet_state,
     )
-    .map(|_call_info| DeployCallPayload {
-        used_resources: resources.clone(),
-        contract_address,
-    })
+    .map(|_call_info| contract_address)
     .map_err(|err| {
         let call_contract_failure = CallFailure::from_execution_error(
             &err,
@@ -91,6 +82,13 @@ pub fn deploy_at(
         CheatcodeError::from(call_contract_failure)
     });
     cheatnet_state.increment_deploy_salt_base();
+
+    // add execution resources used by deploy to all used resources
+    cheatnet_state.used_resources.extend(&UsedResources {
+        execution_resources: resources,
+        l2_to_l1_payloads_length: vec![],
+    });
+
     result
 }
 
@@ -99,7 +97,7 @@ pub fn deploy(
     cheatnet_state: &mut CheatnetState,
     class_hash: &ClassHash,
     calldata: &[Felt252],
-) -> Result<DeployCallPayload, CheatcodeError> {
+) -> Result<ContractAddress, CheatcodeError> {
     let contract_address = cheatnet_state.precalculate_address(class_hash, calldata);
 
     deploy_at(

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/l1_handler_execute.rs

@@ -1,5 +1,5 @@
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{
-    call_l1_handler, CallOutput,
+    call_l1_handler, CallResult,
 };
 use crate::state::{BlockifierState, CheatnetState};
 use blockifier::abi::abi_utils::starknet_keccak;
@@ -14,7 +14,7 @@ impl BlockifierState<'_> {
         function_name: &Felt252,
         from_address: &Felt252,
         payload: &[Felt252],
-    ) -> CallOutput {
+    ) -> CallResult {
         let selector = starknet_keccak(&function_name.to_bytes_be());
 
         let mut calldata = vec![from_address.clone()];
@@ -27,6 +27,5 @@ impl BlockifierState<'_> {
             &selector,
             calldata.as_slice(),
         )
-        .expect("Calling l1 handler failed")
     }
 }

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs

@@ -3,9 +3,7 @@ use std::collections::HashMap;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{
     CallFailure, CallResult, UsedResources,
 };
-use crate::runtime_extensions::forge_runtime_extension::cheatcodes::deploy::{
-    deploy, deploy_at, DeployCallPayload,
-};
+use crate::runtime_extensions::forge_runtime_extension::cheatcodes::deploy::{deploy, deploy_at};
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::CheatcodeError;
 use crate::state::{BlockifierState, CallTrace, CheatTarget};
 use anyhow::{Context, Result};
@@ -288,9 +286,9 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
             "deploy" => {
                 let class_hash = reader.read_felt().into_();
                 let calldata = reader.read_vec();
-                let cheatable_starknet_runtime = &mut extended_runtime.extended_runtime;
+                let cheatnet_runtime = &mut extended_runtime.extended_runtime;
                 let mut blockifier_state = BlockifierState::from(
-                    cheatable_starknet_runtime
+                    cheatnet_runtime
                         .extended_runtime
                         .extended_runtime
                         .hint_handler
@@ -299,10 +297,7 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
 
                 handle_deploy_result(deploy(
                     &mut blockifier_state,
-                    cheatable_starknet_runtime
-                        .extended_runtime
-                        .extension
-                        .cheatnet_state,
+                    cheatnet_runtime.extended_runtime.extension.cheatnet_state,
                     &class_hash,
                     &calldata,
                 ))
@@ -393,16 +388,13 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                 let mut blockifier_state =
                     BlockifierState::from(cheatnet_runtime.extended_runtime.hint_handler.state);
 
-                match blockifier_state
-                    .l1_handler_execute(
-                        cheatnet_runtime.extension.cheatnet_state,
-                        contract_address,
-                        &function_name,
-                        &from_address,
-                        &payload,
-                    )
-                    .result
-                {
+                match blockifier_state.l1_handler_execute(
+                    cheatnet_runtime.extension.cheatnet_state,
+                    contract_address,
+                    &function_name,
+                    &from_address,
+                    &payload,
+                ) {
                     CallResult::Success { .. } => {
                         Ok(CheatcodeHandlingResult::Handled(vec![Felt252::from(0)]))
                     }
@@ -667,11 +659,11 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
 }
 
 fn handle_deploy_result(
-    deploy_result: Result<DeployCallPayload, CheatcodeError>,
+    deploy_result: Result<ContractAddress, CheatcodeError>,
 ) -> Result<CheatcodeHandlingResult, EnhancedHintError> {
     match deploy_result {
-        Ok(deploy_payload) => {
-            let felt_contract_address: Felt252 = deploy_payload.contract_address.into_();
+        Ok(contract_address) => {
+            let felt_contract_address = contract_address.into_();
             let result = vec![Felt252::from(0), felt_contract_address];
             Ok(CheatcodeHandlingResult::Handled(result))
         }

crates/cheatnet/tests/builtins/panic_call.rs

@@ -22,8 +22,7 @@ fn call_contract_error() {
         &contract_address,
         &selector,
         &[Felt252::from(420)],
-    )
-    .unwrap();
+    );
 
     assert_error!(output, "0x496e70757420746f6f206c6f6e6720666f7220617267756d656e7473 ('Input too long for arguments')");
 }
@@ -44,8 +43,7 @@ fn call_contract_panic() {
         &contract_address,
         &selector,
         &[],
-    )
-    .unwrap();
+    );
 
     assert_panic!(
         output,

crates/cheatnet/tests/builtins/segment_arena.rs

@@ -20,8 +20,7 @@ fn segment_arena_simple() {
         &contract_address,
         &selector,
         &[],
-    )
-    .unwrap();
+    );
 
     assert_success!(output, &[]);
 }