Trace as tree (#1539)

Issues that arises: #1553

## Introduced changes

- trace as recursive struct (tree)
- `get_call_trace` returning full call from test instead of
`get_last_call_trace` (due to ambiguity with e.g. for deploy without
constructor new trace isn't added)

## 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`

---------

Co-authored-by: Maksymilian Demitraszek <[email protected]>
Co-authored-by: Maksymilian Demitraszek <[email protected]>

Author: piotmag769

crates/cheatnet/src/constants.rs

@@ -6,15 +6,19 @@ use blockifier::execution::contract_class::{ContractClassV1, ContractClassV1Inne
 use blockifier::execution::contract_class::ContractClass;
 use cairo_vm::types::program::Program;
 
+use crate::runtime_extensions::common::create_entry_point_selector;
+use blockifier::execution::entry_point::{CallEntryPoint, CallType};
+use conversions::IntoConv;
+use starknet::core::utils::get_selector_from_name;
 use starknet_api::deprecated_contract_class::EntryPointType;
 
 use runtime::starknet::context::ERC20_CONTRACT_ADDRESS;
 use runtime::starknet::state::DictStateReader;
 use starknet_api::{
-    core::{ClassHash, ContractAddress, Nonce, PatriciaKey},
+    core::{ClassHash, ContractAddress, PatriciaKey},
     hash::{StarkFelt, StarkHash},
     patricia_key, stark_felt,
-    transaction::{Calldata, DeclareTransactionV2, InvokeTransactionV1},
+    transaction::Calldata,
 };
 
 pub const MAX_FEE: u128 = 1_000_000 * 100_000_000_000; // 1000000 * min_gas_price.
@@ -32,32 +36,6 @@ pub const TEST_CONTRACT_CLASS_HASH: &str = "0x117";
 // snforge_std/src/cheatcodes.cairo::test_address
 pub const TEST_ADDRESS: &str = "0x01724987234973219347210837402";
 
-#[must_use]
-pub fn build_declare_transaction(
-    nonce: Nonce,
-    class_hash: ClassHash,
-    sender_address: ContractAddress,
-) -> DeclareTransactionV2 {
-    DeclareTransactionV2 {
-        nonce,
-        class_hash,
-        sender_address,
-        ..Default::default()
-    }
-}
-
-#[must_use]
-pub fn build_invoke_transaction(
-    calldata: Calldata,
-    sender_address: ContractAddress,
-) -> InvokeTransactionV1 {
-    InvokeTransactionV1 {
-        sender_address,
-        calldata,
-        ..Default::default()
-    }
-}
-
 fn contract_class_no_entrypoints() -> ContractClass {
     let inner = ContractClassV1Inner {
         program: Program::default(),
@@ -99,3 +77,22 @@ pub fn build_testing_state() -> DictStateReader {
         ..Default::default()
     }
 }
+
+#[must_use]
+pub fn build_test_entry_point() -> CallEntryPoint {
+    let test_selector = get_selector_from_name("TEST_CONTRACT_SELECTOR")
+        .unwrap()
+        .into_();
+    let entry_point_selector = create_entry_point_selector(&test_selector);
+    CallEntryPoint {
+        class_hash: None,
+        code_address: Some(ContractAddress(patricia_key!(TEST_ADDRESS))),
+        entry_point_type: EntryPointType::External,
+        entry_point_selector,
+        calldata: Calldata(Arc::new(vec![])),
+        storage_address: ContractAddress(patricia_key!(TEST_ADDRESS)),
+        caller_address: ContractAddress::default(),
+        call_type: CallType::Call,
+        initial_gas: u64::MAX,
+    }
+}

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

@@ -34,7 +34,11 @@ pub fn execute_call_entry_point(
 ) -> EntryPointExecutionResult<CallInfo> {
     // region: Modified blockifier code
     // We skip recursion depth validation here.
-    cheatnet_state.trace_info.push(entry_point.clone());
+
+    cheatnet_state
+        .trace_data
+        .enter_nested_call(entry_point.clone());
+
     if let Some(ret_data) = get_ret_data_by_call_entry_point(entry_point, cheatnet_state) {
         return Ok(mocked_call_info(entry_point.clone(), ret_data.clone()));
     }
@@ -86,6 +90,8 @@ pub fn execute_call_entry_point(
         ),
     };
 
+    cheatnet_state.trace_data.exit_nested_call();
+
     result.map_err(|error| {
         // endregion
         match error {

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

@@ -54,22 +54,6 @@ impl<'a> ExtensionLogic for CallToBlockifierExtension<'a> {
         vm: &mut VirtualMachine,
         extended_runtime: &mut IORuntime<'a>,
     ) -> Result<SyscallHandlingResult, HintError> {
-        match selector {
-            // We clear it here to ensure that on each new contract call
-            // deploy syscall and library call made from test code
-            // we start with an empty trace
-            // Same case when handling l1_handler_execute and in `deploy_at`
-            DeprecatedSyscallSelector::CallContract
-            | DeprecatedSyscallSelector::LibraryCall
-            | DeprecatedSyscallSelector::LibraryCallL1Handler
-            | DeprecatedSyscallSelector::Deploy => extended_runtime
-                .extended_runtime
-                .extension
-                .cheatnet_state
-                .trace_info
-                .clear(),
-            _ => (),
-        }
         match selector {
             // We execute contract calls and library calls with modified blockifier
             // This is redirected to drop ForgeRuntimeExtension

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

@@ -40,7 +40,6 @@ pub fn deploy_at(
     calldata: &[Felt252],
     contract_address: ContractAddress,
 ) -> Result<DeployCallPayload, CheatcodeError> {
-    cheatnet_state.trace_info.clear();
     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) {

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

@@ -7,7 +7,7 @@ use crate::runtime_extensions::forge_runtime_extension::cheatcodes::deploy::{
     deploy, deploy_at, DeployCallPayload,
 };
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::CheatcodeError;
-use crate::state::{BlockifierState, CheatTarget};
+use crate::state::{BlockifierState, CallTrace, CheatTarget};
 use anyhow::{Context, Result};
 use blockifier::execution::call_info::{CallExecution, CallInfo};
 use blockifier::execution::deprecated_syscalls::DeprecatedSyscallSelector;
@@ -393,7 +393,6 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                 let mut blockifier_state =
                     BlockifierState::from(cheatnet_runtime.extended_runtime.hint_handler.state);
 
-                cheatnet_runtime.extension.cheatnet_state.trace_info.clear();
                 match blockifier_state
                     .l1_handler_execute(
                         cheatnet_runtime.extension.cheatnet_state,
@@ -592,18 +591,17 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                     Felt252::from_bytes_be(&s_bytes[0..16]),  // 16 high bytes of s
                 ]))
             }
-            "get_last_call_trace" => {
-                let trace_info = &extended_runtime
+            "get_call_trace" => {
+                let call_trace = &extended_runtime
                     .extended_runtime
                     .extended_runtime
                     .extension
                     .cheatnet_state
-                    .trace_info;
-                let mut output = vec![Felt252::from(trace_info.len())];
+                    .trace_data
+                    .current_call_stack
+                    .borrow_full_trace();
+                let output = serialize_call_trace(call_trace);
 
-                for call_entry_point in trace_info {
-                    output.append(&mut serialize_call_entry_point(call_entry_point));
-                }
                 Ok(CheatcodeHandlingResult::Handled(output))
             }
             "store" => {
@@ -684,6 +682,19 @@ fn handle_deploy_result(
     }
 }
 
+fn serialize_call_trace(call_trace: &CallTrace) -> Vec<Felt252> {
+    let mut output = vec![];
+    output.append(&mut serialize_call_entry_point(&call_trace.entry_point));
+
+    output.push(Felt252::from(call_trace.nested_calls.len()));
+
+    for call_trace in &call_trace.nested_calls {
+        output.append(&mut serialize_call_trace(&call_trace.borrow()));
+    }
+
+    output
+}
+
 fn serialize_call_entry_point(call_entry_point: &CallEntryPoint) -> Vec<Felt252> {
     let entry_point_type = match call_entry_point.entry_point_type {
         EntryPointType::Constructor => 0,

crates/cheatnet/src/state.rs

@@ -19,14 +19,17 @@ use runtime::starknet::state::DictStateReader;
 
 use starknet_api::core::EntryPointSelector;
 
+use crate::constants::build_test_entry_point;
 use starknet_api::transaction::ContractAddressSalt;
 use starknet_api::{
     core::{ClassHash, CompiledClassHash, ContractAddress, Nonce},
     hash::StarkFelt,
     state::StorageKey,
 };
+use std::cell::{Ref, RefCell};
 use std::collections::HashMap;
 use std::hash::BuildHasher;
+use std::rc::Rc;
 
 // Specifies which contracts to target
 // with a cheatcode function
@@ -136,7 +139,45 @@ pub enum CheatStatus<T> {
     Uncheated,
 }
 
-#[derive(Default)]
+/// Tree structure representing trace of a call.
+pub struct CallTrace {
+    pub entry_point: CallEntryPoint,
+    pub nested_calls: Vec<Rc<RefCell<CallTrace>>>,
+}
+
+pub struct NotEmptyCallStack(Vec<Rc<RefCell<CallTrace>>>);
+
+impl NotEmptyCallStack {
+    pub fn from(elem: Rc<RefCell<CallTrace>>) -> Self {
+        NotEmptyCallStack(vec![elem])
+    }
+
+    pub fn push(&mut self, elem: Rc<RefCell<CallTrace>>) {
+        self.0.push(elem);
+    }
+
+    pub fn top(&mut self) -> Rc<RefCell<CallTrace>> {
+        let top_val = self.0.pop().unwrap();
+        let borrowed_ref = top_val.clone();
+        self.0.push(top_val);
+        borrowed_ref
+    }
+
+    pub fn pop(&mut self) -> Rc<RefCell<CallTrace>> {
+        assert!(self.0.len() > 1, "You cannot make NotEmptyCallStack empty");
+        self.0.pop().unwrap()
+    }
+
+    #[must_use]
+    pub fn borrow_full_trace(&self) -> Ref<'_, CallTrace> {
+        self.0.first().unwrap().borrow()
+    }
+}
+
+pub struct TraceData {
+    pub current_call_stack: NotEmptyCallStack,
+}
+
 pub struct CheatnetState {
     pub rolled_contracts: HashMap<ContractAddress, CheatStatus<Felt252>>,
     pub global_roll: Option<Felt252>,
@@ -155,8 +196,38 @@ pub struct CheatnetState {
     pub block_info: BlockInfo,
     // execution resources used by all contract calls
     pub used_resources: UsedResources,
-    // trace info of the last call
-    pub trace_info: Vec<CallEntryPoint>,
+
+    pub trace_data: TraceData,
+}
+
+impl Default for CheatnetState {
+    fn default() -> Self {
+        let test_call = Rc::new(RefCell::new(CallTrace {
+            entry_point: build_test_entry_point(),
+            nested_calls: vec![],
+        }));
+        Self {
+            rolled_contracts: Default::default(),
+            global_roll: None,
+            pranked_contracts: Default::default(),
+            global_prank: None,
+            warped_contracts: Default::default(),
+            global_warp: None,
+            elected_contracts: Default::default(),
+            global_elect: None,
+            mocked_functions: Default::default(),
+            spoofed_contracts: Default::default(),
+            global_spoof: None,
+            spies: vec![],
+            detected_events: vec![],
+            deploy_salt_base: 0,
+            block_info: Default::default(),
+            used_resources: Default::default(),
+            trace_data: TraceData {
+                current_call_stack: NotEmptyCallStack::from(test_call),
+            },
+        }
+    }
 }
 
 impl CheatnetState {
@@ -224,6 +295,27 @@ impl CheatnetState {
     }
 }
 
+impl TraceData {
+    pub fn enter_nested_call(&mut self, entry_point: CallEntryPoint) {
+        let new_call = Rc::new(RefCell::new(CallTrace {
+            entry_point,
+            nested_calls: vec![],
+        }));
+        let current_call = self.current_call_stack.top();
+
+        current_call
+            .borrow_mut()
+            .nested_calls
+            .push(new_call.clone());
+
+        self.current_call_stack.push(new_call);
+    }
+
+    pub fn exit_nested_call(&mut self) {
+        self.current_call_stack.pop();
+    }
+}
+
 fn match_node_response<T: Default>(result: StateResult<T>) -> StateResult<T> {
     match result {
         Ok(class_hash) => Ok(class_hash),

crates/forge-runner/src/running.rs

@@ -10,9 +10,7 @@ use crate::test_case_summary::{Single, TestCaseSummary};
 use crate::{RunnerConfig, RunnerParams, TestCaseRunnable, CACHE_DIR};
 use anyhow::{bail, ensure, Result};
 use blockifier::execution::common_hints::ExecutionMode;
-use blockifier::execution::entry_point::{
-    CallEntryPoint, CallType, EntryPointExecutionContext, ExecutionResources,
-};
+use blockifier::execution::entry_point::{EntryPointExecutionContext, ExecutionResources};
 use blockifier::execution::execution_utils::ReadOnlySegments;
 use blockifier::execution::syscalls::hint_processor::SyscallHintProcessor;
 use blockifier::state::cached_state::CachedState;
@@ -27,6 +25,7 @@ use cairo_vm::types::relocatable::Relocatable;
 use cairo_vm::vm::vm_core::VirtualMachine;
 use camino::Utf8Path;
 use cheatnet::constants as cheatnet_constants;
+use cheatnet::constants::build_test_entry_point;
 use cheatnet::forking::state::ForkStateReader;
 use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::CallToBlockifierExtension;
 use cheatnet::runtime_extensions::cheatable_starknet_runtime_extension::CheatableStarknetRuntimeExtension;
@@ -39,13 +38,6 @@ use itertools::chain;
 use runtime::starknet::context;
 use runtime::starknet::context::BlockInfo;
 use runtime::{ExtendedRuntime, StarknetRuntime};
-use starknet::core::utils::get_selector_from_name;
-use starknet_api::core::PatriciaKey;
-use starknet_api::core::{ContractAddress, EntryPointSelector};
-use starknet_api::deprecated_contract_class::EntryPointType;
-use starknet_api::hash::StarkHash;
-use starknet_api::patricia_key;
-use starknet_api::transaction::Calldata;
 use tokio::sync::mpsc::Sender;
 use tokio::task::JoinHandle;
 
@@ -152,22 +144,7 @@ fn build_syscall_handler<'a>(
     execution_resources: &'a mut ExecutionResources,
     context: &'a mut EntryPointExecutionContext,
 ) -> SyscallHintProcessor<'a> {
-    let test_selector = get_selector_from_name("TEST_CONTRACT_SELECTOR").unwrap();
-    let entry_point_selector =
-        EntryPointSelector(StarkHash::new(test_selector.to_bytes_be()).unwrap());
-    let entry_point = CallEntryPoint {
-        class_hash: None,
-        code_address: Some(ContractAddress(patricia_key!(
-            cheatnet_constants::TEST_ADDRESS
-        ))),
-        entry_point_type: EntryPointType::External,
-        entry_point_selector,
-        calldata: Calldata(Arc::new(vec![])),
-        storage_address: ContractAddress(patricia_key!(cheatnet_constants::TEST_ADDRESS)),
-        caller_address: ContractAddress::default(),
-        call_type: CallType::Call,
-        initial_gas: u64::MAX,
-    };
+    let entry_point = build_test_entry_point();
 
     SyscallHintProcessor::new(
         blockifier_state,