Skip to content

Commit 240bd5e

Browse files
Add flag to save call traces to files (#1549)
Closes software-mansion/cairo-profiler#3 ## Introduced changes - call traces are saved in `.snfoundry_trace` directory (maybe we should add a flag to clean this directory too?) ## 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]>
1 parent d100392 commit 240bd5e

File tree

14 files changed

+269
-7
lines changed

14 files changed

+269
-7
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1212
#### Added
1313

1414
- `store` and `load` cheatcodes
15+
- `--save-trace-data` flag to `snforge test` command. Traces can be used for profiling purposes.
1516

1617
#### Changed
1718

crates/cheatnet/src/state.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ pub enum CheatStatus<T> {
140140
}
141141

142142
/// Tree structure representing trace of a call.
143+
#[derive(Clone)]
143144
pub struct CallTrace {
144145
pub entry_point: CallEntryPoint,
145146
pub nested_calls: Vec<Rc<RefCell<CallTrace>>>,

crates/forge-runner/src/lib.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use futures::StreamExt;
1919
use once_cell::sync::Lazy;
2020
use scarb_api::StarknetContractArtifacts;
2121
use smol_str::SmolStr;
22+
use trace_data::save_trace_data;
2223

2324
use std::collections::HashMap;
2425
use std::default::Default;
@@ -31,6 +32,7 @@ pub mod compiled_runnable;
3132
pub mod expected_result;
3233
pub mod test_case_summary;
3334
pub mod test_crate_summary;
35+
pub mod trace_data;
3436

3537
mod fuzzer;
3638
mod gas;
@@ -62,6 +64,7 @@ pub struct RunnerConfig {
6264
pub exit_first: bool,
6365
pub fuzzer_runs: u32,
6466
pub fuzzer_seed: u64,
67+
pub save_trace_data: bool,
6568
}
6669

6770
impl RunnerConfig {
@@ -72,12 +75,14 @@ impl RunnerConfig {
7275
exit_first: bool,
7376
fuzzer_runs: u32,
7477
fuzzer_seed: u64,
78+
save_trace_data: bool,
7579
) -> Self {
7680
Self {
7781
workspace_root,
7882
exit_first,
7983
fuzzer_runs,
8084
fuzzer_seed,
85+
save_trace_data,
8186
}
8287
}
8388
}
@@ -186,6 +191,12 @@ pub async fn run_tests_from_crate(
186191

187192
print_test_result(&result);
188193

194+
if runner_config.save_trace_data {
195+
if let AnyTestCaseSummary::Single(result) = &result {
196+
save_trace_data(result);
197+
}
198+
}
199+
189200
if result.is_failed() && runner_config.exit_first {
190201
interrupted = true;
191202
rec.close();

crates/forge-runner/src/running.rs

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1+
use std::cell::RefCell;
12
use std::collections::HashMap;
23
use std::default::Default;
34
use std::marker::PhantomData;
5+
use std::rc::Rc;
46
use std::sync::Arc;
57

68
use crate::compiled_runnable::ValidatedForkConfig;
@@ -33,7 +35,7 @@ use cheatnet::runtime_extensions::forge_runtime_extension::{
3335
get_all_execution_resources, ForgeExtension, ForgeRuntime,
3436
};
3537
use cheatnet::runtime_extensions::io_runtime_extension::IORuntimeExtension;
36-
use cheatnet::state::{BlockInfoReader, CheatnetState, ExtendedStateReader};
38+
use cheatnet::state::{BlockInfoReader, CallTrace, CheatnetState, ExtendedStateReader};
3739
use itertools::chain;
3840
use runtime::starknet::context;
3941
use runtime::starknet::context::BlockInfo;
@@ -163,6 +165,7 @@ fn build_syscall_handler<'a>(
163165

164166
pub struct RunResultWithInfo {
165167
pub(crate) run_result: Result<RunResult, RunnerError>,
168+
pub(crate) call_trace: Rc<RefCell<CallTrace>>,
166169
pub(crate) gas_used: u128,
167170
}
168171

@@ -259,13 +262,15 @@ pub fn run_test_case(
259262
);
260263

261264
let block_context = get_context(&forge_runtime).block_context.clone();
265+
let call_trace_ref = get_call_trace_ref(&mut forge_runtime);
262266
let execution_resources = get_all_execution_resources(forge_runtime);
263267

264268
let gas = calculate_used_gas(&block_context, &mut blockifier_state, &execution_resources);
265269

266270
Ok(RunResultWithInfo {
267271
run_result,
268272
gas_used: gas,
273+
call_trace: call_trace_ref,
269274
})
270275
}
271276

@@ -282,6 +287,7 @@ fn extract_test_case_summary(
282287
case,
283288
args,
284289
result_with_info.gas_used,
290+
&result_with_info.call_trace,
285291
)),
286292
// CairoRunError comes from VirtualMachineError which may come from HintException that originates in TestExecutionSyscallHandler
287293
Err(RunnerError::CairoRunError(error)) => Ok(TestCaseSummary::Failed {
@@ -331,3 +337,15 @@ fn get_context<'a>(runtime: &'a ForgeRuntime) -> &'a EntryPointExecutionContext
331337
.hint_handler
332338
.context
333339
}
340+
341+
fn get_call_trace_ref(runtime: &mut ForgeRuntime) -> Rc<RefCell<CallTrace>> {
342+
runtime
343+
.extended_runtime
344+
.extended_runtime
345+
.extended_runtime
346+
.extension
347+
.cheatnet_state
348+
.trace_data
349+
.current_call_stack
350+
.top()
351+
}

crates/forge-runner/src/test_case_summary.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
use crate::compiled_runnable::TestCaseRunnable;
22
use crate::expected_result::{ExpectedPanicValue, ExpectedTestResult};
33
use crate::gas::check_available_gas;
4+
use crate::trace_data::CallTrace;
45
use cairo_felt::Felt252;
56
use cairo_lang_runner::short_string::as_cairo_short_string;
67
use cairo_lang_runner::{RunResult, RunResultValue};
8+
use cheatnet::state::CallTrace as InternalCallTrace;
79
use num_traits::Pow;
10+
use std::cell::RefCell;
811
use std::option::Option;
12+
use std::rc::Rc;
913

1014
#[derive(Debug, PartialEq, Clone, Default)]
1115
pub struct GasStatistics {
@@ -52,20 +56,23 @@ pub struct FuzzingStatistics {
5256
pub trait TestType {
5357
type GasInfo: std::fmt::Debug + Clone;
5458
type TestStatistics: std::fmt::Debug + Clone;
59+
type TraceData: std::fmt::Debug + Clone;
5560
}
5661

5762
#[derive(Debug, PartialEq, Clone)]
5863
pub struct Fuzzing;
5964
impl TestType for Fuzzing {
6065
type GasInfo = GasStatistics;
6166
type TestStatistics = FuzzingStatistics;
67+
type TraceData = ();
6268
}
6369

6470
#[derive(Debug, PartialEq, Clone)]
6571
pub struct Single;
6672
impl TestType for Single {
6773
type GasInfo = u128;
6874
type TestStatistics = ();
75+
type TraceData = CallTrace;
6976
}
7077

7178
/// Summary of running a single test case
@@ -83,6 +90,8 @@ pub enum TestCaseSummary<T: TestType> {
8390
gas_info: <T as TestType>::GasInfo,
8491
/// Statistics of the test run
8592
test_statistics: <T as TestType>::TestStatistics,
93+
/// Test trace data
94+
trace_data: <T as TestType>::TraceData,
8695
},
8796
/// Test case failed
8897
Failed {
@@ -147,6 +156,7 @@ impl TestCaseSummary<Fuzzing> {
147156
arguments,
148157
gas_info: _,
149158
test_statistics: (),
159+
trace_data: _,
150160
} => {
151161
let runs = results.len();
152162
let gas_usages: Vec<u128> = results
@@ -163,6 +173,7 @@ impl TestCaseSummary<Fuzzing> {
163173
gas_info: GasStatistics::new(&gas_usages),
164174
arguments,
165175
test_statistics: FuzzingStatistics { runs },
176+
trace_data: (),
166177
}
167178
}
168179
TestCaseSummary::Failed {
@@ -191,6 +202,7 @@ impl TestCaseSummary<Single> {
191202
test_case: &TestCaseRunnable,
192203
arguments: Vec<Felt252>,
193204
gas: u128,
205+
call_trace: &Rc<RefCell<InternalCallTrace>>,
194206
) -> Self {
195207
let name = test_case.name.to_string();
196208
let msg = extract_result_data(&run_result, &test_case.expected_result);
@@ -203,6 +215,7 @@ impl TestCaseSummary<Single> {
203215
arguments,
204216
test_statistics: (),
205217
gas_info: gas,
218+
trace_data: CallTrace::from(call_trace.borrow().clone()),
206219
};
207220
check_available_gas(&test_case.available_gas, summary)
208221
}
@@ -235,6 +248,7 @@ impl TestCaseSummary<Single> {
235248
arguments,
236249
test_statistics: (),
237250
gas_info: gas,
251+
trace_data: CallTrace::from(call_trace.borrow().clone()),
238252
},
239253
},
240254
},
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
use std::fs;
2+
use std::path::PathBuf;
3+
4+
// Will be provided by profiler crate in the future
5+
// This module will be removed!
6+
use cheatnet::state::CallTrace as InternalCallTrace;
7+
use serde::{Deserialize, Serialize};
8+
use starknet_api::core::{ClassHash, ContractAddress, EntryPointSelector};
9+
use starknet_api::deprecated_contract_class::EntryPointType;
10+
use starknet_api::transaction::Calldata;
11+
12+
use crate::test_case_summary::{Single, TestCaseSummary};
13+
14+
pub const TRACE_DIR: &str = "snfoundry_trace";
15+
16+
/// Tree structure representing trace of a call.
17+
#[derive(Debug, Clone, Deserialize, Serialize)]
18+
pub struct CallTrace {
19+
pub entry_point: CallEntryPoint,
20+
pub nested_calls: Vec<CallTrace>,
21+
}
22+
23+
impl From<InternalCallTrace> for CallTrace {
24+
fn from(value: InternalCallTrace) -> Self {
25+
CallTrace {
26+
entry_point: CallEntryPoint::from(value.entry_point),
27+
nested_calls: value
28+
.nested_calls
29+
.into_iter()
30+
.map(|c| CallTrace::from(c.borrow().clone()))
31+
.collect(),
32+
}
33+
}
34+
}
35+
36+
#[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
37+
pub struct CallEntryPoint {
38+
pub class_hash: Option<ClassHash>,
39+
pub code_address: Option<ContractAddress>,
40+
pub entry_point_type: EntryPointType,
41+
pub entry_point_selector: EntryPointSelector,
42+
pub calldata: Calldata,
43+
pub storage_address: ContractAddress,
44+
pub caller_address: ContractAddress,
45+
pub call_type: CallType,
46+
pub initial_gas: u64,
47+
}
48+
49+
impl From<blockifier::execution::entry_point::CallEntryPoint> for CallEntryPoint {
50+
fn from(value: blockifier::execution::entry_point::CallEntryPoint) -> Self {
51+
let blockifier::execution::entry_point::CallEntryPoint {
52+
class_hash,
53+
code_address,
54+
entry_point_type,
55+
entry_point_selector,
56+
calldata,
57+
storage_address,
58+
caller_address,
59+
call_type,
60+
initial_gas,
61+
} = value;
62+
63+
CallEntryPoint {
64+
class_hash,
65+
code_address,
66+
entry_point_type,
67+
entry_point_selector,
68+
calldata,
69+
storage_address,
70+
caller_address,
71+
call_type: call_type.into(),
72+
initial_gas,
73+
}
74+
}
75+
}
76+
77+
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq, Deserialize, Serialize)]
78+
pub enum CallType {
79+
#[default]
80+
Call = 0,
81+
Delegate = 1,
82+
}
83+
84+
impl From<blockifier::execution::entry_point::CallType> for CallType {
85+
fn from(value: blockifier::execution::entry_point::CallType) -> Self {
86+
match value {
87+
blockifier::execution::entry_point::CallType::Call => CallType::Call,
88+
blockifier::execution::entry_point::CallType::Delegate => CallType::Delegate,
89+
}
90+
}
91+
}
92+
93+
pub fn save_trace_data(summary: &TestCaseSummary<Single>) {
94+
if let TestCaseSummary::Passed {
95+
name, trace_data, ..
96+
} = summary
97+
{
98+
let serialized_trace =
99+
serde_json::to_string(trace_data).expect("Failed to serialize call trace");
100+
let dir_to_save_trace = PathBuf::from(TRACE_DIR);
101+
fs::create_dir_all(&dir_to_save_trace)
102+
.expect("Failed to create a file to save call trace to");
103+
104+
let filename = format!("{name}.json");
105+
fs::write(dir_to_save_trace.join(filename), serialized_trace)
106+
.expect("Failed to write call trace to a file");
107+
}
108+
}

0 commit comments

Comments
 (0)