Move trace data code to separate file

Author: ddoktorski

crates/cheatnet/src/lib.rs

@@ -6,3 +6,4 @@ pub mod predeployment;
 pub mod runtime_extensions;
 pub mod state;
 pub mod sync_client;
+pub mod trace_data;

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

@@ -17,7 +17,7 @@ use crate::runtime_extensions::{
         storage::{calculate_variable_address, load, store},
     },
 };
-use crate::state::{CallTrace, CallTraceNode, GasReportData};
+use crate::trace_data::{CallTrace, CallTraceNode, GasReportData};
 use anyhow::{Context, Result, anyhow};
 use blockifier::bouncer::vm_resources_to_sierra_gas;
 use blockifier::context::TransactionContext;

crates/cheatnet/src/state.rs

@@ -3,42 +3,31 @@ use crate::forking::state::ForkStateReader;
 use crate::predeployment::erc20::eth::eth_predeployed_contract;
 use crate::predeployment::erc20::strk::strk_predeployed_contract;
 use crate::predeployment::predeployed_contract::PredeployedContract;
-use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::CallResult;
-use crate::runtime_extensions::common::sum_syscall_usage;
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::cheat_execution_info::{
     ExecutionInfoMock, ResourceBounds,
 };
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::spy_events::Event;
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::spy_messages_to_l1::MessageToL1;
-use blockifier::blockifier_versioned_constants::VersionedConstants;
-use blockifier::execution::call_info::{ExecutionSummary, OrderedEvent, OrderedL2ToL1Message};
+use crate::trace_data::{CallTrace, NotEmptyCallStack, TraceData};
 use blockifier::execution::contract_class::RunnableCompiledClass;
-use blockifier::execution::entry_point::CallEntryPoint;
-use blockifier::execution::syscalls::vm_syscall_utils::SyscallUsageMap;
 use blockifier::state::errors::StateError::UndeclaredClassHash;
 use blockifier::state::state_api::{StateReader, StateResult};
-use cairo_annotations::trace_data::L1Resources;
 use cairo_vm::Felt252;
-use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
-use cairo_vm::vm::trace::trace_entry::RelocatedTraceEntry;
 use conversions::serde::deserialize::CairoDeserialize;
-use conversions::serde::serialize::{BufferWriter, CairoSerialize};
 use conversions::string::TryFromHexStr;
 use indexmap::IndexMap;
 use runtime::starknet::constants::TEST_CONTRACT_CLASS_HASH;
 use runtime::starknet::context::SerializableBlockInfo;
 use runtime::starknet::state::DictStateReader;
 use starknet_api::block::BlockInfo;
 use starknet_api::core::{ChainId, EntryPointSelector};
-use starknet_api::execution_resources::GasVector;
 use starknet_api::transaction::fields::ContractAddressSalt;
-use starknet_api::transaction::fields::GasVectorComputationMode;
 use starknet_api::{
     core::{ClassHash, CompiledClassHash, ContractAddress, Nonce},
     state::StorageKey,
 };
 use starknet_types_core::felt::Felt;
-use std::cell::{OnceCell, Ref, RefCell};
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::num::NonZeroUsize;
 use std::rc::Rc;
@@ -206,165 +195,6 @@ impl<T> CheatStatus<T> {
     }
 }
 
-#[derive(Clone, Debug)]
-pub struct GasReportData {
-    execution_summary: ExecutionSummary,
-    gas: OnceCell<GasVector>,
-}
-
-impl GasReportData {
-    #[must_use]
-    pub fn new(execution_summary: ExecutionSummary) -> Self {
-        Self {
-            execution_summary,
-            gas: OnceCell::new(),
-        }
-    }
-
-    pub fn get_gas(&self) -> &GasVector {
-        self.gas.get_or_init(|| {
-            self.execution_summary.clone().to_partial_gas_vector(
-                VersionedConstants::latest_constants(),
-                &GasVectorComputationMode::All,
-            )
-        })
-    }
-}
-
-/// Tree structure representing trace of a call.
-#[derive(Debug)]
-pub struct CallTrace {
-    // only these are serialized
-    pub entry_point: CallEntryPoint,
-    pub nested_calls: Vec<CallTraceNode>,
-    pub result: CallResult,
-    // serialize end
-
-    // These also include resources used by internal calls
-    pub used_execution_resources: ExecutionResources,
-    pub used_l1_resources: L1Resources,
-    pub used_syscalls_vm_resources: SyscallUsageMap,
-    pub used_syscalls_sierra_gas: SyscallUsageMap,
-    pub vm_trace: Option<Vec<RelocatedTraceEntry>>,
-    pub gas_consumed: u64,
-    pub events: Vec<OrderedEvent>,
-    pub signature: Vec<Felt>,
-
-    // This is updated only once after the entire test execution.
-    pub gas_report_data: Option<GasReportData>,
-}
-
-impl CairoSerialize for CallTrace {
-    fn serialize(&self, output: &mut BufferWriter) {
-        self.entry_point.serialize(output);
-
-        let visible_calls: Vec<_> = self
-            .nested_calls
-            .iter()
-            .filter_map(CallTraceNode::extract_entry_point_call)
-            .collect();
-
-        visible_calls.serialize(output);
-
-        self.result.serialize(output);
-    }
-}
-
-impl CallTrace {
-    fn default_successful_call() -> Self {
-        Self {
-            entry_point: CallEntryPoint::default(),
-            used_execution_resources: ExecutionResources::default(),
-            used_l1_resources: L1Resources::default(),
-            used_syscalls_vm_resources: SyscallUsageMap::default(),
-            used_syscalls_sierra_gas: SyscallUsageMap::default(),
-            nested_calls: vec![],
-            result: CallResult::Success { ret_data: vec![] },
-            vm_trace: None,
-            gas_consumed: u64::default(),
-            events: vec![],
-            signature: vec![],
-            gas_report_data: None,
-        }
-    }
-
-    #[must_use]
-    pub fn get_total_used_syscalls(&self) -> SyscallUsageMap {
-        sum_syscall_usage(
-            self.used_syscalls_vm_resources.clone(),
-            &self.used_syscalls_sierra_gas,
-        )
-    }
-}
-
-/// Enum representing node of a trace of a call.
-#[derive(Clone, Debug)]
-pub enum CallTraceNode {
-    EntryPointCall(Rc<RefCell<CallTrace>>),
-    DeployWithoutConstructor,
-}
-
-impl CallTraceNode {
-    #[must_use]
-    pub fn extract_entry_point_call(&self) -> Option<&Rc<RefCell<CallTrace>>> {
-        if let CallTraceNode::EntryPointCall(trace) = self {
-            Some(trace)
-        } else {
-            None
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-struct CallStackElement {
-    call_trace: Rc<RefCell<CallTrace>>,
-    cheated_data: CheatedData,
-}
-
-#[derive(Debug)]
-pub struct NotEmptyCallStack(Vec<CallStackElement>);
-
-impl NotEmptyCallStack {
-    pub fn from(elem: Rc<RefCell<CallTrace>>) -> Self {
-        NotEmptyCallStack(vec![CallStackElement {
-            call_trace: elem,
-            cheated_data: CheatedData::default(),
-        }])
-    }
-
-    pub fn push(&mut self, elem: Rc<RefCell<CallTrace>>, cheated_data: CheatedData) {
-        self.0.push(CallStackElement {
-            call_trace: elem,
-            cheated_data,
-        });
-    }
-
-    pub fn top(&mut self) -> Rc<RefCell<CallTrace>> {
-        let top_val = self.0.last().unwrap();
-        top_val.call_trace.clone()
-    }
-
-    pub fn top_cheated_data(&mut self) -> CheatedData {
-        let top_val = self.0.last().unwrap();
-        top_val.cheated_data.clone()
-    }
-
-    fn pop(&mut self) -> CallStackElement {
-        assert!(self.0.len() > 1, "You cannot make NotEmptyCallStack empty");
-        self.0.pop().unwrap()
-    }
-
-    #[must_use]
-    pub fn size(&self) -> usize {
-        self.0.len()
-    }
-
-    #[must_use]
-    pub fn borrow_full_trace(&self) -> Ref<'_, CallTrace> {
-        self.0.first().unwrap().call_trace.borrow()
-    }
-}
-
 #[derive(Clone, Default, Debug, PartialEq, Eq)]
 pub struct CheatedTxInfo {
     pub version: Option<Felt>,
@@ -399,12 +229,6 @@ pub struct CheatedData {
     pub tx_info: CheatedTxInfo,
 }
 
-#[derive(Debug)]
-pub struct TraceData {
-    pub current_call_stack: NotEmptyCallStack,
-    pub is_vm_trace_needed: bool,
-}
-
 pub struct CheatnetState {
     pub cheated_execution_info_contracts: HashMap<ContractAddress, ExecutionInfoMock>,
     pub global_cheated_execution_info: ExecutionInfoMock,
@@ -577,88 +401,3 @@ impl CheatnetState {
         self.encountered_errors.shift_remove(&class_hash);
     }
 }
-
-impl TraceData {
-    pub fn enter_nested_call(&mut self, entry_point: CallEntryPoint, cheated_data: CheatedData) {
-        let new_call = Rc::new(RefCell::new(CallTrace {
-            entry_point,
-            ..CallTrace::default_successful_call()
-        }));
-        let current_call = self.current_call_stack.top();
-
-        current_call
-            .borrow_mut()
-            .nested_calls
-            .push(CallTraceNode::EntryPointCall(new_call.clone()));
-
-        self.current_call_stack.push(new_call, cheated_data);
-    }
-
-    pub fn set_class_hash_for_current_call(&mut self, class_hash: ClassHash) {
-        let current_call = self.current_call_stack.top();
-        current_call.borrow_mut().entry_point.class_hash = Some(class_hash);
-    }
-
-    pub fn set_vm_trace_for_current_call(&mut self, vm_trace: Vec<RelocatedTraceEntry>) {
-        let current_call = self.current_call_stack.top();
-        current_call.borrow_mut().vm_trace = Some(vm_trace);
-    }
-
-    pub fn update_current_call_result(&mut self, result: CallResult) {
-        let current_call = self.current_call_stack.top();
-        current_call.borrow_mut().result = result;
-    }
-
-    pub fn clear_current_call_events_and_messages(&mut self) {
-        let current_call = self.current_call_stack.top();
-        current_call.borrow_mut().events.clear();
-        current_call
-            .borrow_mut()
-            .used_l1_resources
-            .l2_l1_message_sizes
-            .clear();
-    }
-
-    #[expect(clippy::too_many_arguments)]
-    pub fn update_current_call(
-        &mut self,
-        execution_resources: ExecutionResources,
-        gas_consumed: u64,
-        used_syscalls_vm_resources: SyscallUsageMap,
-        used_syscalls_sierra_gas: SyscallUsageMap,
-        result: CallResult,
-        l2_to_l1_messages: &[OrderedL2ToL1Message],
-        signature: Vec<Felt>,
-        events: Vec<OrderedEvent>,
-    ) {
-        let current_call = self.current_call_stack.top();
-        let mut current_call = current_call.borrow_mut();
-
-        current_call.used_execution_resources = execution_resources;
-        current_call.gas_consumed = gas_consumed;
-        current_call.used_syscalls_vm_resources = used_syscalls_vm_resources;
-        current_call.used_syscalls_sierra_gas = used_syscalls_sierra_gas;
-
-        current_call.used_l1_resources.l2_l1_message_sizes = l2_to_l1_messages
-            .iter()
-            .map(|ordered_message| ordered_message.message.payload.0.len())
-            .collect();
-
-        current_call.result = result;
-        current_call.signature = signature;
-        current_call.events = events;
-    }
-
-    pub fn exit_nested_call(&mut self) {
-        self.current_call_stack.pop();
-    }
-
-    pub fn add_deploy_without_constructor_node(&mut self) {
-        let current_call = self.current_call_stack.top();
-
-        current_call
-            .borrow_mut()
-            .nested_calls
-            .push(CallTraceNode::DeployWithoutConstructor);
-    }
-}