Use context in debugging

crates/debugging/src/lib.rs

@@ -7,6 +7,5 @@ mod contracts_data_store;
 mod trace;
 mod tree;
 
-pub use contracts_data_store::ContractsDataStore;
 pub use trace::components::{Component, Components};
-pub use trace::types::Trace;
+pub use trace::{context::Context, types::Trace};

crates/debugging/src/trace/collect.rs

@@ -1,10 +1,9 @@
-use crate::Trace;
 use crate::contracts_data_store::ContractsDataStore;
-use crate::trace::components::Components;
 use crate::trace::types::{
     CallerAddress, ContractAddress, ContractName, ContractTrace, Selector, TestName, TraceInfo,
     TransformedCallResult, TransformedCalldata,
 };
+use crate::{Context, Trace};
 use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{
     CallFailure, CallResult as CheatnetCallResult,
 };
@@ -16,22 +15,16 @@ use starknet_api::execution_utils::format_panic_data;
 
 pub struct Collector<'a> {
     call_trace: &'a CallTrace,
-    contracts_data_store: &'a ContractsDataStore,
-    components: &'a Components,
+    context: &'a Context,
 }
 
 impl<'a> Collector<'a> {
     /// Creates a new [`Collector`] from a given `cheatnet` [`CallTrace`], [`ContractsDataStore`] and [`Verbosity`].
     #[must_use]
-    pub fn new(
-        call_trace: &'a CallTrace,
-        contracts_data_store: &'a ContractsDataStore,
-        components: &'a Components,
-    ) -> Collector<'a> {
+    pub fn new(call_trace: &'a CallTrace, context: &'a Context) -> Collector<'a> {
         Collector {
             call_trace,
-            contracts_data_store,
-            components,
+            context,
         }
     }
 
@@ -43,7 +36,7 @@ impl<'a> Collector<'a> {
     }
 
     fn collect_contract_trace(&self) -> ContractTrace {
-        let components = self.components;
+        let components = self.context.components();
         let entry_point = &self.call_trace.entry_point;
         let nested_calls = self.collect_nested_calls();
         let contract_name = self.collect_contract_name();
@@ -75,29 +68,28 @@ impl<'a> Collector<'a> {
             .map(|call_trace| {
                 Collector {
                     call_trace: &call_trace.borrow(),
-                    contracts_data_store: self.contracts_data_store,
-                    components: self.components,
+                    context: self.context,
                 }
                 .collect_contract_trace()
             })
             .collect()
     }
 
     fn collect_contract_name(&self) -> ContractName {
-        self.contracts_data_store
+        self.contracts_data_store()
             .get_contract_name(self.class_hash())
             .cloned()
             .unwrap_or_else(|| ContractName("forked contract".to_string()))
     }
 
     fn collect_selector(&self) -> &Selector {
-        self.contracts_data_store
+        self.contracts_data_store()
             .get_selector(&self.call_trace.entry_point.entry_point_selector)
             .expect("`Selector` should be present")
     }
 
     fn collect_abi(&self) -> &[AbiEntry] {
-        self.contracts_data_store
+        self.contracts_data_store()
             .get_abi(self.class_hash())
             .expect("`ABI` should be present")
     }
@@ -140,6 +132,10 @@ impl<'a> Collector<'a> {
             .as_ref()
             .expect("class_hash should be set in `fn execute_call_entry_point` in cheatnet")
     }
+
+    fn contracts_data_store(&self) -> &ContractsDataStore {
+        self.context.contracts_data_store()
+    }
 }
 
 fn format_result_message(tag: &str, message: &str) -> String {

crates/debugging/src/trace/context.rs

@@ -0,0 +1,38 @@
+use crate::Components;
+use crate::contracts_data_store::ContractsDataStore;
+use cheatnet::forking::data::ForkData;
+use cheatnet::runtime_extensions::forge_runtime_extension::contracts_data::ContractsData;
+
+/// Context is a structure that holds the necessary data for creating a [`Trace`](crate::Trace).
+pub struct Context {
+    contracts_data_store: ContractsDataStore,
+    components: Components,
+}
+
+impl Context {
+    /// Creates a new instance of [`Context`] from a given `cheatnet` [`ContractsData`], [`ForkData`] and [`Components`].
+    #[must_use]
+    pub fn new(
+        contracts_data: &ContractsData,
+        fork_data: &ForkData,
+        components: Components,
+    ) -> Self {
+        let contracts_data_store = ContractsDataStore::new(contracts_data, fork_data);
+        Self {
+            contracts_data_store,
+            components,
+        }
+    }
+
+    /// Returns a reference to the [`ContractsDataStore`].
+    #[must_use]
+    pub fn contracts_data_store(&self) -> &ContractsDataStore {
+        &self.contracts_data_store
+    }
+
+    /// Returns a reference to the [`Components`].
+    #[must_use]
+    pub fn components(&self) -> &Components {
+        &self.components
+    }
+}

crates/debugging/src/trace/mod.rs

@@ -1,3 +1,4 @@
 pub mod collect;
 pub mod components;
+pub mod context;
 pub mod types;

crates/debugging/src/trace/types.rs

@@ -1,7 +1,7 @@
-use crate::contracts_data_store::ContractsDataStore;
+use crate::Context;
 use crate::trace::collect::Collector;
 use crate::trace::components::{
-    CallResultContainer, CallTypeContainer, CalldataContainer, CallerAddressContainer, Components,
+    CallResultContainer, CallTypeContainer, CalldataContainer, CallerAddressContainer,
     ContractAddressContainer, ContractNameContainer, EntryPointTypeContainer,
 };
 use crate::tree::TreeSerialize;
@@ -56,15 +56,10 @@ pub struct ContractAddress(pub ApiContractAddress);
 pub struct CallerAddress(pub ApiContractAddress);
 
 impl Trace {
-    /// Creates a new [`Trace`] from a given `cheatnet` [`CallTrace`], [`ContractsDataStore`], [`Verbosity`] and a test name.
+    /// Creates a new [`Trace`] from a given [`Context`] and a test name.
     #[must_use]
-    pub fn new(
-        call_trace: &CallTrace,
-        contracts_data_store: &ContractsDataStore,
-        components: &Components,
-        test_name: String,
-    ) -> Self {
-        Collector::new(call_trace, contracts_data_store, components).collect_trace(test_name)
+    pub fn new(call_trace: &CallTrace, context: &Context, test_name: String) -> Self {
+        Collector::new(call_trace, context).collect_trace(test_name)
     }
 }
 

crates/forge-runner/src/debugging/mod.rs

@@ -18,11 +18,6 @@ pub fn build_debugging_trace(
     fork_data: &ForkData,
 ) -> Option<debugging::Trace> {
     let components = trace_args.to_components()?;
-    let contracts_data_store = debugging::ContractsDataStore::new(contracts_data, fork_data);
-    Some(debugging::Trace::new(
-        call_trace,
-        &contracts_data_store,
-        &components,
-        test_name,
-    ))
+    let context = debugging::Context::new(contracts_data, fork_data, components);
+    Some(debugging::Trace::new(call_trace, &context, test_name))
 }