Skip to content

Commit

Permalink
refactor logic
Browse files Browse the repository at this point in the history
  • Loading branch information
piotmag769 committed Jun 19, 2024
1 parent c80b055 commit 005fd0b
Show file tree
Hide file tree
Showing 3 changed files with 139 additions and 111 deletions.
85 changes: 46 additions & 39 deletions crates/cairo-profiler/src/trace_reader/function_trace_builder.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
use crate::profiler_config::FunctionLevelConfig;
use crate::sierra_loader::StatementsFunctionsMap;
use crate::trace_reader::function_name::FunctionName;
use crate::trace_reader::function_trace_builder::function_stack_trace::CallStack;
use crate::trace_reader::function_trace_builder::inlining::build_original_function_stack;
use crate::trace_reader::sample::{
FunctionCall, InternalFunctionCall, MeasurementUnit, MeasurementValue, Sample,
use crate::trace_reader::function_trace_builder::function_stack_trace::{
CallStack, VecWithLimitedCapacity,
};
use crate::trace_reader::function_trace_builder::inlining::build_original_call_stack_with_inlined_calls;
use crate::trace_reader::sample::{FunctionCall, MeasurementUnit, MeasurementValue, Sample};
use cairo_lang_sierra::extensions::core::{CoreConcreteLibfunc, CoreLibfunc, CoreType};
use cairo_lang_sierra::program::{GenStatement, Program, StatementIdx};
use cairo_lang_sierra::program_registry::ProgramRegistry;
use cairo_lang_sierra_to_casm::compiler::CairoProgramDebugInfo;
use itertools::{chain, Itertools};
use itertools::Itertools;
use std::collections::HashMap;
use std::ops::AddAssign;
use trace_data::TraceEntry;
Expand Down Expand Up @@ -68,9 +68,7 @@ pub fn collect_function_level_profiling_info(
1
};

// TODO: make sure the `max_function_stack_trace_depth` flag works with inlines
let mut call_stack = CallStack::new(function_level_config.max_function_stack_trace_depth);
let mut original_call_stacks_of_non_inlined_functions_calls = vec![];

// The value is the steps of the stack trace so far, not including the pending steps being
// tracked at the time. The key is a function stack trace.
Expand All @@ -83,8 +81,8 @@ pub fn collect_function_level_profiling_info(
let mut end_of_program_reached = false;

for step in trace {
// Skip the header. This only makes sense when a header was added to CASM program.
if step.pc < real_minimal_pc && run_with_call_header {
// Skip the header.
if step.pc < real_minimal_pc {
header_steps += 1;
continue;
}
Expand Down Expand Up @@ -113,30 +111,16 @@ pub fn collect_function_level_profiling_info(
function_level_config.split_generics,
);

// TODO: optimize clones
let original_function_stack_of_last_non_inlined_function_call = chain!(
original_call_stacks_of_non_inlined_functions_calls
.last()
.cloned()
.unwrap_or_default(),
[FunctionCall::InternalFunctionCall(
InternalFunctionCall::NonInlined(current_function_name.clone())
)]
)
.collect();

let original_function_stack = if function_level_config.show_inlined_functions {
build_original_function_stack(
sierra_statement_idx,
maybe_statements_functions_map.as_ref(),
original_function_stack_of_last_non_inlined_function_call,
)
} else {
original_function_stack_of_last_non_inlined_function_call
};
let current_call_stack = build_current_call_stack(
&call_stack,
current_function_name.clone(),
function_level_config.show_inlined_functions,
sierra_statement_idx,
maybe_statements_functions_map.as_ref(),
);

*functions_stack_traces
.entry(original_function_stack.clone())
.entry(current_call_stack.into())
.or_insert(Steps(0)) += 1;

let Some(gen_statement) = program.statements.get(sierra_statement_idx.0) else {
Expand All @@ -149,17 +133,18 @@ pub fn collect_function_level_profiling_info(
sierra_program_registry.get_libfunc(&invocation.libfunc_id),
Ok(CoreConcreteLibfunc::FunctionCall(_))
) {
// TODO: hide this logic in CallStack
// and test with recursive functions.
original_call_stacks_of_non_inlined_functions_calls
.push(original_function_stack);
call_stack.enter_function_call(current_function_name);
let current_call_stack = build_current_call_stack(
&call_stack,
current_function_name.clone(),
function_level_config.show_inlined_functions,
sierra_statement_idx,
maybe_statements_functions_map.as_ref(),
);

call_stack.enter_function_call(current_call_stack);
}
}
GenStatement::Return(_) => {
// TODO: hide this logic in CallStack
// and test with recursive functions
original_call_stacks_of_non_inlined_functions_calls.pop();
if call_stack.exit_function_call().is_none() {
end_of_program_reached = true;
}
Expand Down Expand Up @@ -212,3 +197,25 @@ fn maybe_sierra_statement_index_by_pc(
MaybeSierraStatementIndex::SierraStatementIndex(statement_index)
}
}

fn build_current_call_stack(
call_stack: &CallStack,
current_function_name: FunctionName,
show_inlined_functions: bool,
sierra_statement_idx: StatementIdx,
maybe_statements_functions_map: Option<&StatementsFunctionsMap>,
) -> VecWithLimitedCapacity<FunctionCall> {
let current_call_stack = call_stack.current_call_stack(current_function_name);

if show_inlined_functions {
let (current_call_stack, max_capacity) = current_call_stack.deconstruct();
let current_call_stack_with_inlined_calls = build_original_call_stack_with_inlined_calls(
sierra_statement_idx,
maybe_statements_functions_map,
current_call_stack,
);
VecWithLimitedCapacity::from(current_call_stack_with_inlined_calls, max_capacity)
} else {
current_call_stack
}
}
Original file line number Diff line number Diff line change
@@ -1,67 +1,93 @@
use crate::trace_reader::function_name::FunctionName;

struct CallStackElement {
pub function_name: FunctionName,
/// Consecutive recursive calls to this function that are currently on the stack.
recursive_calls_count: usize,
}
use crate::trace_reader::sample::{FunctionCall, InternalFunctionCall};

/// The function call stack of the current function, excluding the current function call.
pub(super) struct CallStack {
stack: Vec<CallStackElement>,
/// Tracks the depth of the function call stack, without limit. This is usually equal to
/// `stack.len()`, but if the actual stack is deeper than `max_function_trace_depth`,
/// this remains reliable while `stack` does not.
real_function_stack_depth: usize,
/// Constant through existence of the object.
max_function_stack_trace_depth: usize,
/// This stack is guaranteed to have maximum of `max_function_stack_trace_depth` elements.
stack: VecWithLimitedCapacity<FunctionCall>,
/// The last element of this vector is always a number of elements of the stack before the last
/// function call.
previous_stack_lengths: Vec<usize>,
}

impl CallStack {
pub fn new(max_function_stack_trace_depth: usize) -> Self {
Self {
stack: vec![],
real_function_stack_depth: 0,
max_function_stack_trace_depth,
stack: VecWithLimitedCapacity::new(max_function_stack_trace_depth),
previous_stack_lengths: vec![],
}
}

pub fn enter_function_call(&mut self, function_name: FunctionName) {
if let Some(stack_element) = self.stack.last_mut() {
if function_name == stack_element.function_name {
stack_element.recursive_calls_count += 1;
return;
}
pub fn enter_function_call(&mut self, new_call_stack: VecWithLimitedCapacity<FunctionCall>) {
self.previous_stack_lengths.push(self.stack.len());

self.stack = new_call_stack;
}

pub fn exit_function_call(&mut self) -> Option<()> {
let previous_stack_len = self.previous_stack_lengths.pop()?;
self.stack.truncate(previous_stack_len);
Some(())
}

/// Returns current call stack truncated to `max_function_stack_trace_depth`.
pub fn current_call_stack(
&self,
current_function_name: FunctionName,
) -> VecWithLimitedCapacity<FunctionCall> {
let mut current_call_stack = self.stack.clone();

current_call_stack.push(FunctionCall::InternalFunctionCall(
InternalFunctionCall::NonInlined(current_function_name),
));

current_call_stack
}
}

#[derive(Clone)]
pub struct VecWithLimitedCapacity<T> {
vector: Vec<T>,
max_capacity: usize,
}

impl<T> VecWithLimitedCapacity<T> {
pub fn new(max_capacity: usize) -> Self {
Self {
vector: vec![],
max_capacity,
}
}

if self.real_function_stack_depth < self.max_function_stack_trace_depth {
self.stack.push(CallStackElement {
function_name,
recursive_calls_count: 0,
});
pub fn from(mut vector: Vec<T>, max_capacity: usize) -> Self {
vector.truncate(max_capacity);
Self {
vector,
max_capacity,
}
self.real_function_stack_depth += 1;
}

pub fn exit_function_call(&mut self) -> Option<()> {
if self.real_function_stack_depth <= self.max_function_stack_trace_depth {
let mut stack_element = self.stack.pop()?;

if stack_element.recursive_calls_count > 0 {
// Recursive function exited.
stack_element.recursive_calls_count -= 1;
self.stack.push(stack_element);

Some(())
} else {
// Regular function exited.
self.real_function_stack_depth -= 1;
Some(())
}
} else {
// Hidden function exited.
self.real_function_stack_depth -= 1;
Some(())
pub fn push(&mut self, el: T) {
if self.vector.len() < self.max_capacity {
self.vector.push(el);
}
}

pub fn truncate(&mut self, len: usize) {
self.vector.truncate(len);
}

pub fn len(&self) -> usize {
self.vector.len()
}

pub fn deconstruct(self) -> (Vec<T>, usize) {
(self.vector, self.max_capacity)
}
}

impl<T> From<VecWithLimitedCapacity<T>> for Vec<T> {
fn from(value: VecWithLimitedCapacity<T>) -> Self {
value.vector
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,63 +7,58 @@ use crate::sierra_loader::StatementsFunctionsMap;
use crate::trace_reader::function_name::FunctionName;
use crate::trace_reader::sample::{FunctionCall, InternalFunctionCall};

// TODO: add comments + better names
pub(super) fn build_original_function_stack(
// TODO: add comments
pub(super) fn build_original_call_stack_with_inlined_calls(
sierra_statement_idx: StatementIdx,
maybe_statements_functions_map: Option<&StatementsFunctionsMap>,
original_call_stack_of_last_non_inlined_function_call: Vec<FunctionCall>,
current_call_stack: Vec<FunctionCall>,
) -> Vec<FunctionCall> {
let maybe_original_function_names_stack = maybe_statements_functions_map
let maybe_original_call_stack = maybe_statements_functions_map
.as_ref()
.and_then(|statements_functions_map| statements_functions_map.get(sierra_statement_idx));

if let Some(mappings) = maybe_original_function_names_stack {
if let Some(original_call_stack) = maybe_original_call_stack {
// Statements functions map represents callstack from the least meaningful elements.
let mappings = mappings.iter().rev().collect_vec();
construct_original_function_stack(
original_call_stack_of_last_non_inlined_function_call,
&mappings,
)
let original_call_stack = original_call_stack.iter().rev().collect_vec();
construct_original_call_stack_with_inlined_calls(current_call_stack, &original_call_stack)
} else {
original_call_stack_of_last_non_inlined_function_call
current_call_stack
}
}

// TODO: test mutually recursive functions (+ inline always with 3 funcs) - add tests!!!
fn construct_original_function_stack(
original_call_stack_of_last_non_inlined_function_call: Vec<FunctionCall>,
mappings: &[&FunctionName],
fn construct_original_call_stack_with_inlined_calls(
current_call_stack: Vec<FunctionCall>,
original_call_stack: &[&FunctionName],
) -> Vec<FunctionCall> {
let start_index = max(
original_call_stack_of_last_non_inlined_function_call.len() as i128
- mappings.len() as i128,
current_call_stack.len() as i128 - original_call_stack.len() as i128,
0,
)
.try_into()
.expect("Non-negative i128 to usize cast should never fail");
let mut num_of_overlapping_functions = 0;
for i in start_index..original_call_stack_of_last_non_inlined_function_call.len() {
for i in start_index..current_call_stack.len() {
let mut overlapped = true;

for j in 0..original_call_stack_of_last_non_inlined_function_call.len() - i {
if mappings[j]
!= original_call_stack_of_last_non_inlined_function_call[i + j].function_name()
{
for j in 0..current_call_stack.len() - i {
if original_call_stack[j] != current_call_stack[i + j].function_name() {
overlapped = false;
break;
}
}

if overlapped {
num_of_overlapping_functions =
original_call_stack_of_last_non_inlined_function_call.len() - i;
num_of_overlapping_functions = current_call_stack.len() - i;
break;
}
}

let mut result = original_call_stack_of_last_non_inlined_function_call;
let mut result = current_call_stack;

for &function_name in &mappings[num_of_overlapping_functions..mappings.len()] {
for &function_name in
&original_call_stack[num_of_overlapping_functions..original_call_stack.len()]
{
result.push(FunctionCall::InternalFunctionCall(
InternalFunctionCall::Inlined(function_name.clone()),
));
Expand Down

0 comments on commit 005fd0b

Please sign in to comment.