feat(profiler): allow to preserve original stack trace order

crates/cairo-lang-runner/src/lib.rs

@@ -316,7 +316,7 @@ impl SierraCasmRunner {
         // function).
         // The value is the weight of the stack trace so far, not including the pending weight being
         // tracked at the time.
-        let mut stack_trace_weights = UnorderedHashMap::default();
+        let mut stack_trace_weights = OrderedHashMap::default();
         let mut end_of_program_reached = false;
         // The total weight of each Sierra statement.
         // Note the header and footer (CASM instructions added for running the program by the

crates/cairo-lang-runner/src/profiling.rs

@@ -7,6 +7,7 @@ use cairo_lang_sierra_generator::db::SierraGenGroup;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
 use cairo_lang_utils::{require, LookupIntern};
+use itertools::Itertools;
 use smol_str::SmolStr;
 
 #[cfg(test)]
@@ -25,7 +26,8 @@ pub struct ProfilingInfo {
     /// as a vector of indices of the functions in the stack (indices of the functions according to
     /// the list in the sierra program).
     /// The value is the weight of the stack trace.
-    pub stack_trace_weights: UnorderedHashMap<Vec<usize>, usize>,
+    /// The stack trace entries are sorted in the order they occur.
+    pub stack_trace_weights: OrderedHashMap<Vec<usize>, usize>,
 }
 
 /// Weights per libfunc.
@@ -205,6 +207,9 @@ pub struct ProfilingInfoProcessorParams {
     /// Whether to process the profiling info by Cairo stack trace (that is, no generated
     /// functions in the traces).
     pub process_by_cairo_stack_trace: bool,
+    /// Whether to sort Sierra stack trace items by weight (desc) and stack entry (asc),
+    /// alternatively the original order is preserved.
+    pub sort_sierra_stack_trace: bool,
 }
 impl Default for ProfilingInfoProcessorParams {
     fn default() -> Self {
@@ -218,6 +223,7 @@ impl Default for ProfilingInfoProcessorParams {
             process_by_cairo_function: true,
             process_by_stack_trace: true,
             process_by_cairo_stack_trace: true,
+            sort_sierra_stack_trace: true,
         }
     }
 }
@@ -310,36 +316,38 @@ impl<'a> ProfilingInfoProcessor<'a> {
         params: &ProfilingInfoProcessorParams,
     ) -> StackTraceWeights {
         let sierra_stack_trace_weights = params.process_by_stack_trace.then(|| {
-            raw_profiling_info
-                .stack_trace_weights
-                .iter_sorted_by_key(|(idx_stack_trace, weight)| {
-                    (usize::MAX - **weight, (*idx_stack_trace).clone())
-                })
-                .map(|(idx_stack_trace, weight)| {
-                    (
-                        index_stack_trace_to_name_stack_trace(
-                            &self.sierra_program,
-                            idx_stack_trace,
-                        ),
-                        *weight,
-                    )
-                })
-                .collect()
+            let resolve_names = |(idx_stack_trace, weight): (&Vec<usize>, &usize)| {
+                (
+                    index_stack_trace_to_name_stack_trace(&self.sierra_program, &idx_stack_trace),
+                    *weight,
+                )
+            };
+            if params.sort_sierra_stack_trace {
+                raw_profiling_info
+                    .stack_trace_weights
+                    .iter()
+                    .sorted_by_key(|&(trace, weight)| (usize::MAX - *weight, trace.clone()))
+                    .map(resolve_names)
+                    .collect()
+            } else {
+                raw_profiling_info.stack_trace_weights.iter().map(resolve_names).collect()
+            }
         });
 
         let cairo_stack_trace_weights = params.process_by_cairo_stack_trace.then(|| {
             let db = self.db.expect("DB must be set with `process_by_cairo_stack_trace=true`.");
             raw_profiling_info
                 .stack_trace_weights
-                .filter_cloned(|trace, _| is_cairo_trace(db, &self.sierra_program, trace))
-                .into_iter_sorted_by_key(|(trace, weight)| (usize::MAX - *weight, trace.clone()))
+                .iter()
+                .filter(|(trace, _)| is_cairo_trace(db, &self.sierra_program, trace))
+                .sorted_by_key(|&(trace, weight)| (usize::MAX - *weight, trace.clone()))
                 .map(|(idx_stack_trace, weight)| {
                     (
                         index_stack_trace_to_name_stack_trace(
                             &self.sierra_program,
                             &idx_stack_trace,
                         ),
-                        weight,
+                        *weight,
                     )
                 })
                 .collect()