Add get_current_vm_step

CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Forge
 
+#### Added
+
+- `get_current_vm_step` function to get the current step from Cairo VM during test execution. For more see [docs](https://foundry-rs.github.io/starknet-foundry/snforge-library/testing/get_current_vm_step.html)
+
 #### Changed
 
 - Gas values in fuzzing test output are now displayed as whole numbers without fractional parts

crates/cheatnet/src/runtime_extensions/forge_config_extension.rs

@@ -1,3 +1,4 @@
+use cairo_vm::vm::vm_core::VirtualMachine;
 use config::RawForgeConfig;
 use conversions::serde::deserialize::BufferReader;
 use runtime::{CheatcodeHandlingResult, EnhancedHintError, ExtensionLogic, StarknetRuntime};
@@ -17,6 +18,7 @@ impl<'a> ExtensionLogic for ForgeConfigExtension<'a> {
         selector: &str,
         mut input_reader: BufferReader<'_>,
         _extended_runtime: &mut Self::Runtime,
+        _vm: &VirtualMachine,
     ) -> Result<CheatcodeHandlingResult, EnhancedHintError> {
         macro_rules! config_cheatcode {
             ( $prop:ident) => {{

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

@@ -19,6 +19,7 @@ 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::{
@@ -81,6 +82,7 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
         selector: &str,
         mut input_reader: BufferReader<'_>,
         extended_runtime: &mut Self::Runtime,
+        vm: &VirtualMachine,
     ) -> Result<CheatcodeHandlingResult, EnhancedHintError> {
         if let Some(oracle_selector) = self
             .oracle_hint_service
@@ -552,6 +554,33 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                     .cheat_block_hash(block_number, operation);
                 Ok(CheatcodeHandlingResult::from_serializable(()))
             }
+            "get_current_vm_step" => {
+                // Each contract call is executed in separate VM, hence all VM steps
+                // are calculated as sum of steps from calls + current VM steps.
+                // Since syscalls are added to VM resources after the execution, we need
+                // to include them manually here.
+                let top_call = extended_runtime
+                    .extended_runtime
+                    .extension
+                    .cheatnet_state
+                    .trace_data
+                    .current_call_stack
+                    .top();
+                let vm_steps_from_calls = calculate_vm_steps_from_calls(&top_call);
+                let top_call_syscalls = &extended_runtime
+                    .extended_runtime
+                    .extended_runtime
+                    .hint_handler
+                    .base
+                    .syscalls_usage;
+                let vm_steps_from_syscalls = &VersionedConstants::latest_constants()
+                    .get_additional_os_syscall_resources(top_call_syscalls)
+                    .n_steps;
+                let total_vm_steps =
+                    vm_steps_from_calls + vm_steps_from_syscalls + vm.get_current_step();
+
+                Ok(CheatcodeHandlingResult::from_serializable(total_vm_steps))
+            }
             _ => Ok(CheatcodeHandlingResult::Forwarded),
         }
     }
@@ -858,3 +887,19 @@ pub fn get_all_used_resources(
         l1_handler_payload_lengths,
     }
 }
+
+fn calculate_vm_steps_from_calls(top_call: &Rc<RefCell<CallTrace>>) -> usize {
+    // Resources from inner calls already include syscall resources used in them
+    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,
+            });
+    used_resources.n_steps
+}

crates/forge/tests/integration/get_current_vm_step.rs

@@ -0,0 +1,39 @@
+use crate::utils::runner::{Contract, assert_passed};
+use crate::utils::running_tests::run_test_case;
+use crate::utils::test_case;
+use forge_runner::forge_config::ForgeTrackedResource;
+use indoc::indoc;
+use std::path::Path;
+
+#[test]
+fn test_get_current_vm_step() {
+    let test = test_case!(
+        indoc!(
+            r#"
+            use snforge_std::testing::get_current_vm_step;
+            use snforge_std::{ContractClassTrait, DeclareResultTrait, declare};
+
+            #[test]
+            fn check_current_vm_step() {
+                let step_start = get_current_vm_step();
+
+                let contract = declare("HelloStarknet").unwrap().contract_class().clone();
+                let _ = contract.deploy(@ArrayTrait::new()).unwrap();
+
+                let step_end = get_current_vm_step();
+
+                assert!(step_end > step_start);
+            }
+        "#
+        ),
+        Contract::from_code_path(
+            "HelloStarknet".to_string(),
+            Path::new("tests/data/contracts/hello_starknet.cairo"),
+        )
+        .unwrap()
+    );
+
+    let result = run_test_case(&test, ForgeTrackedResource::CairoSteps);
+
+    assert_passed(&result);
+}

crates/forge/tests/integration/mod.rs

@@ -19,6 +19,7 @@ mod gas;
 mod generate_random_felt;
 mod get_available_gas;
 mod get_class_hash;
+mod get_current_vm_step;
 mod interact_with_state;
 mod l1_handler_executor;
 mod message_to_l1;

crates/runtime/src/lib.rs

@@ -279,6 +279,7 @@ impl<Extension: ExtensionLogic> ExtendedRuntime<Extension> {
             &selector,
             BufferReader::new(&inputs),
             &mut self.extended_runtime,
+            vm,
         );
 
         let res = match result {
@@ -423,6 +424,7 @@ pub trait ExtensionLogic {
         _selector: &str,
         _input_reader: BufferReader,
         _extended_runtime: &mut Self::Runtime,
+        _vm: &VirtualMachine,
     ) -> Result<CheatcodeHandlingResult, EnhancedHintError> {
         Ok(CheatcodeHandlingResult::Forwarded)
     }

crates/sncast/src/starknet_commands/script/run.rs

@@ -104,6 +104,7 @@ impl<'a> ExtensionLogic for CastScriptExtension<'a> {
         selector: &str,
         mut input_reader: BufferReader,
         _extended_runtime: &mut Self::Runtime,
+        _vm: &VirtualMachine,
     ) -> Result<CheatcodeHandlingResult, EnhancedHintError> {
         match selector {
             "call" => {

docs/listings/testing_reference/Scarb.toml

@@ -0,0 +1,12 @@
+[package]
+name = "testing_reference"
+version = "0.1.0"
+edition = "2024_07"
+
+[dependencies]
+starknet = "2.12.0"
+
+[dev-dependencies]
+snforge_std = { path = "../../../snforge_std" }
+
+[[target.starknet-contract]]

docs/listings/testing_reference/src/lib.cairo

@@ -0,0 +1,27 @@
+#[starknet::interface]
+pub trait ICounter<TContractState> {
+    fn increment(ref self: TContractState);
+}
+
+#[starknet::contract]
+pub mod Counter {
+    use starknet::storage::{StoragePointerReadAccess, StoragePointerWriteAccess};
+
+    #[storage]
+    struct Storage {
+        i: felt252,
+    }
+
+    #[constructor]
+    fn constructor(ref self: ContractState) {
+        self.i.write(0);
+    }
+
+    #[abi(embed_v0)]
+    impl CounterImpl of super::ICounter<ContractState> {
+        fn increment(ref self: ContractState) {
+            let current_value = self.i.read();
+            self.i.write(current_value + 1);
+        }
+    }
+}

docs/listings/testing_reference/tests/tests.cairo

@@ -0,0 +1,30 @@
+use snforge_std::testing::get_current_vm_step;
+use snforge_std::{ContractClassTrait, DeclareResultTrait, declare};
+use testing_reference::{ICounterSafeDispatcher, ICounterSafeDispatcherTrait};
+
+#[feature("safe_dispatcher")]
+fn setup() {
+    // Deploy contract
+    let (contract_address, _) = declare("Counter")
+        .unwrap()
+        .contract_class()
+        .deploy(@array![])
+        .unwrap();
+
+    let dispatcher = ICounterSafeDispatcher { contract_address };
+
+    // Increment counter a few times
+    dispatcher.increment();
+    dispatcher.increment();
+    dispatcher.increment();
+}
+
+#[test]
+fn test_setup_steps() {
+    let steps_start = get_current_vm_step();
+    setup();
+    let steps_end = get_current_vm_step();
+
+    // Assert that setup used no more than 100 steps
+    assert!(steps_end - steps_start <= 100);
+}

docs/src/SUMMARY.md

@@ -129,6 +129,8 @@
     * [env](appendix/snforge-library/env.md)
     * [signature](appendix/snforge-library/signature.md)
     * [fuzzable](appendix/snforge-library/fuzzable.md)
+    * [testing](appendix/snforge-library/testing.md)
+        * [get_current_vm_step](appendix/snforge-library/testing/get_current_vm_step.md)
 * [`sncast` Commands](appendix/sncast.md)
     * [common flags](appendix/sncast/common.md)
     * [account](appendix/sncast/account/account.md)

docs/src/appendix/snforge-library/testing.md

@@ -0,0 +1,7 @@
+# `testing` Module
+
+Module containing functions useful for testing.
+
+## Functions
+
+* [`get_current_vm_step`](./testing/get_current_vm_step.md)

docs/src/appendix/snforge-library/testing/get_current_vm_step.md

@@ -0,0 +1,49 @@
+# `get_current_vm_step`
+
+Gets the current step from Cairo VM during test execution.
+
+```rust
+fn get_current_vm_step() -> u32;
+```
+
+## Example
+
+Let's consider a simple counter contract that increments a value stored in its storage.
+
+```rust
+{{#include ../../../../listings/testing_reference/src/lib.cairo}}
+```
+
+Now, let's define `setup` function which deploys this contract and increments the counter a few times and assert that `setup` function does not exceed a certain number of steps during its execution. This is particularly useful for performance testing and ensuring that our setup logic remains efficient.
+
+```rust
+{{#include ../../../../listings/testing_reference/tests/tests.cairo}}
+```
+
+<!-- { "package_name": "testing_reference" } -->
+Let's run the test:
+```shell
+$ snforge test test_setup_steps
+```
+
+<details>
+<summary>Output:</summary>
+
+```shell
+Collected 1 test(s) from testing_reference package
+Running 0 test(s) from src/
+Running 1 test(s) from tests/
+[FAIL] testing_reference_integrationtest::tests::test_setup_steps
+
+Failure data:
+    "assertion failed: `steps_end - steps_start <= 100`."
+
+Tests: 0 passed, 1 failed, 0 ignored, 0 filtered out
+
+Failures:
+    testing_reference_integrationtest::tests::test_setup_steps
+```
+</details>
+<br>
+
+As we can see, the test fails because the `setup` function exceeded the allowed number of steps (100 in this case).

snforge_std/src/lib.cairo

@@ -126,6 +126,8 @@ pub mod fuzzable;
 
 pub mod signature;
 
+pub mod testing;
+
 pub mod trace;
 
 #[doc(hidden)]

snforge_std/src/testing.cairo

@@ -0,0 +1,6 @@
+use crate::cheatcode::execute_cheatcode_and_deserialize;
+
+/// Gets the current step from Cairo VM during test execution
+pub fn get_current_vm_step() -> u32 {
+    execute_cheatcode_and_deserialize::<'get_current_vm_step', u32>(array![].span())
+}