Apply code review suggestions

Author: franciszekjob

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/entry_point.rs

@@ -225,16 +225,6 @@ pub fn execute_call_entry_point(
                 cheatnet_state.trace_data.exit_nested_call();
             }
 
-            let resources_from_inner_calls = &res
-                .call_info
-                .inner_calls
-                .iter()
-                .fold(ExecutionResources::default(), |acc, call_info| {
-                    &acc + &call_info.resources
-                });
-            let resources_from_current_call = &res.call_info.resources - resources_from_inner_calls;
-            cheatnet_state.add_used_resources(&resources_from_current_call);
-
             Ok(res.call_info)
         }
         Err(EntryPointExecutionError::PreExecutionError(err))

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

@@ -192,8 +192,6 @@ fn execute_syscall<Request: ExecuteCall + SyscallRequest>(
     // Increment, since the selector was peeked into before
     syscall_handler.syscall_ptr += 1;
     syscall_handler.increment_syscall_count_by(&selector, 1);
-
-    cheatnet_state.add_used_syscall(&selector, 1);
     // endregion
 
     let syscall_gas_cost = syscall_handler

crates/cheatnet/src/runtime_extensions/cheatable_starknet_runtime_extension.rs

@@ -195,8 +195,6 @@ impl CheatableStarknetRuntimeExtension<'_> {
         syscall_handler.syscall_ptr += 1;
         syscall_handler.increment_syscall_count_by(&selector, 1);
 
-        self.cheatnet_state.add_used_syscall(&selector, 1);
-
         let syscall_gas_cost = syscall_handler
             .get_gas_cost_from_selector(&selector)
             .map_err(|error| SyscallExecutorBaseError::GasCost { error, selector })?;

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

@@ -20,6 +20,7 @@ pub mod replace_bytecode;
 pub mod spy_events;
 pub mod spy_messages_to_l1;
 pub mod storage;
+pub mod testing;
 
 /// A structure used for returning cheatcode errors in tests
 #[derive(Debug)]

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

@@ -0,0 +1,52 @@
+use std::{cell::RefCell, rc::Rc};
+
+use blockifier::{
+    blockifier_versioned_constants::VersionedConstants,
+    execution::syscalls::vm_syscall_utils::SyscallUsageMap,
+};
+use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
+
+use crate::{
+    runtime_extensions::common::sum_syscall_usage,
+    trace_data::{CallTrace, CallTraceNode},
+};
+
+pub fn calculate_steps_from_calls(
+    top_call: &Rc<RefCell<CallTrace>>,
+    top_call_syscalls: &SyscallUsageMap,
+) -> usize {
+    let used_resources =
+        &top_call
+            .borrow()
+            .nested_calls
+            .iter()
+            .fold(ExecutionResources::default(), |acc, node| match node {
+                CallTraceNode::EntryPointCall(call_trace) => {
+                    &acc + &call_trace.borrow().used_execution_resources
+                }
+                CallTraceNode::DeployWithoutConstructor => acc,
+            });
+
+    let inner_calls_syscalls = &top_call.borrow().nested_calls.iter().fold(
+        SyscallUsageMap::new(),
+        |acc, node| match node {
+            CallTraceNode::EntryPointCall(call_trace) => sum_syscall_usage(
+                acc,
+                &sum_syscall_usage(
+                    call_trace.borrow().used_syscalls_sierra_gas.clone(),
+                    &call_trace.borrow().used_syscalls_vm_resources,
+                ),
+            ),
+            CallTraceNode::DeployWithoutConstructor => acc,
+        },
+    );
+
+    let total_syscalls = sum_syscall_usage(inner_calls_syscalls.clone(), top_call_syscalls);
+
+    let total_syscalls_exeucution_resources = &VersionedConstants::latest_constants()
+        .get_additional_os_syscall_resources(&total_syscalls);
+
+    let resources_from_calls = used_resources + total_syscalls_exeucution_resources;
+
+    resources_from_calls.n_steps
+}

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

@@ -2,6 +2,7 @@ use self::contracts_data::ContractsData;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::UsedResources;
 use crate::runtime_extensions::common::sum_syscall_usage;
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::replace_bytecode::ReplaceBytecodeError;
+use crate::runtime_extensions::forge_runtime_extension::cheatcodes::testing::calculate_steps_from_calls;
 use crate::runtime_extensions::{
     call_to_blockifier_runtime_extension::{
         CallToBlockifierRuntime,
@@ -19,7 +20,6 @@ use crate::runtime_extensions::{
 };
 use crate::trace_data::{CallTrace, CallTraceNode, GasReportData};
 use anyhow::{Context, Result, anyhow};
-use blockifier::blockifier_versioned_constants::VersionedConstants;
 use blockifier::bouncer::vm_resources_to_sierra_gas;
 use blockifier::context::TransactionContext;
 use blockifier::execution::call_info::{
@@ -555,21 +555,23 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                 Ok(CheatcodeHandlingResult::from_serializable(()))
             }
             "get_current_step" => {
-                let used_resources = &extended_runtime
+                let top_call = extended_runtime
                     .extended_runtime
                     .extension
                     .cheatnet_state
-                    .used_resources;
-                let execution_resources_from_used_syscalls =
-                    &VersionedConstants::latest_constants().get_additional_os_syscall_resources(
-                        &extended_runtime
-                            .extended_runtime
-                            .extension
-                            .cheatnet_state
-                            .used_syscalls,
-                    );
-                let resources_from_calls = used_resources + execution_resources_from_used_syscalls;
-                let total_steps = resources_from_calls.n_steps + vm.get_current_step();
+                    .trace_data
+                    .current_call_stack
+                    .top();
+
+                let top_call_syscalls = &extended_runtime
+                    .extended_runtime
+                    .extended_runtime
+                    .hint_handler
+                    .base
+                    .syscalls_usage;
+
+                let steps_from_calls = calculate_steps_from_calls(&top_call, top_call_syscalls);
+                let total_steps = steps_from_calls + vm.get_current_step();
 
                 Ok(CheatcodeHandlingResult::from_serializable(total_steps))
             }

crates/cheatnet/src/state.rs

@@ -10,11 +10,9 @@ use crate::runtime_extensions::forge_runtime_extension::cheatcodes::spy_events::
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::spy_messages_to_l1::MessageToL1;
 use crate::trace_data::{CallTrace, NotEmptyCallStack, TraceData};
 use blockifier::execution::contract_class::RunnableCompiledClass;
-use blockifier::execution::syscalls::vm_syscall_utils::{SyscallSelector, SyscallUsageMap};
 use blockifier::state::errors::StateError::UndeclaredClassHash;
 use blockifier::state::state_api::{StateReader, StateResult};
 use cairo_vm::Felt252;
-use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
 use conversions::serde::deserialize::CairoDeserialize;
 use conversions::string::TryFromHexStr;
 use indexmap::IndexMap;
@@ -248,8 +246,6 @@ pub struct CheatnetState {
     pub fuzzer_args: Vec<String>,
     pub block_hash_contracts: HashMap<(ContractAddress, u64), (CheatSpan, Felt)>,
     pub global_block_hash: HashMap<u64, (Felt, Vec<ContractAddress>)>,
-    pub used_resources: ExecutionResources,
-    pub used_syscalls: SyscallUsageMap,
 }
 
 pub type EncounteredErrors = IndexMap<ClassHash, Vec<usize>>;
@@ -281,8 +277,6 @@ impl Default for CheatnetState {
             fuzzer_args: Vec::default(),
             block_hash_contracts: HashMap::default(),
             global_block_hash: HashMap::default(),
-            used_resources: ExecutionResources::default(),
-            used_syscalls: SyscallUsageMap::default(),
         }
     }
 }
@@ -406,13 +400,4 @@ impl CheatnetState {
     pub fn clear_error(&mut self, class_hash: ClassHash) {
         self.encountered_errors.shift_remove(&class_hash);
     }
-
-    pub fn add_used_resources(&mut self, resources: &ExecutionResources) {
-        self.used_resources += resources;
-    }
-
-    pub fn add_used_syscall(&mut self, syscall: &SyscallSelector, count: usize) {
-        let usage = self.used_syscalls.entry(*syscall).or_default();
-        usage.call_count += count;
-    }
 }