print line by line cycle counts in execution history

Author: morganthomas

crates/lean_vm/src/execution/runner.rs

@@ -1,6 +1,6 @@
 //! VM execution runner
 
-use crate::HintExecutionContext;
+use crate::{SourceLineNumber, HintExecutionContext};
 use crate::core::{
     DIMENSION, F, ONE_VEC_PTR, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR,
     PUBLIC_INPUT_START, VECTOR_LEN, ZERO_VEC_PTR,
@@ -72,36 +72,55 @@ pub fn execute_bytecode(
 ) -> ExecutionResult {
     let mut std_out = String::new();
     let mut instruction_history = ExecutionHistory::new();
-    execute_bytecode_helper(
-        bytecode,
-        public_input,
-        private_input,
-        no_vec_runtime_memory,
-        &mut std_out,
-        &mut instruction_history,
-        profiler,
-        function_locations,
-    )
-    .unwrap_or_else(|err| {
-        let lines_history = &instruction_history.lines;
-        let latest_instructions =
-            &lines_history[lines_history.len().saturating_sub(STACK_TRACE_INSTRUCTIONS)..];
-        println!(
-            "\n{}",
-            crate::diagnostics::pretty_stack_trace(
-                source_code,
-                latest_instructions,
-                function_locations
-            )
-        );
-        if !std_out.is_empty() {
-            println!("╔══════════════════════════════════════════════════════════════╗");
-            println!("║                         STD-OUT                              ║");
-            println!("╚══════════════════════════════════════════════════════════════╝\n");
-            print!("{std_out}");
-        }
-        panic!("Error during bytecode execution: {err}");
-    })
+    let result =
+        execute_bytecode_helper(
+            bytecode,
+            public_input,
+            private_input,
+            no_vec_runtime_memory,
+            &mut std_out,
+            &mut instruction_history,
+            profiler,
+            function_locations,
+        )
+        .unwrap_or_else(|err| {
+            let lines_history = &instruction_history.lines;
+            let latest_instructions =
+                &lines_history[lines_history.len().saturating_sub(STACK_TRACE_INSTRUCTIONS)..];
+            println!(
+                "\n{}",
+                crate::diagnostics::pretty_stack_trace(
+                    source_code,
+                    latest_instructions,
+                    function_locations
+                )
+            );
+            if !std_out.is_empty() {
+                println!("╔══════════════════════════════════════════════════════════════╗");
+                println!("║                         STD-OUT                              ║");
+                println!("╚══════════════════════════════════════════════════════════════╝\n");
+                print!("{std_out}");
+            }
+            panic!("Error during bytecode execution: {err}");
+        });
+    if profiler {
+        print_line_cycle_counts(bytecode, instruction_history);
+    }
+    result
+}
+
+fn print_line_cycle_counts(bytecode: &Bytecode, history: ExecutionHistory) {
+    println!("Line by line cycle counts");
+    println!("=========================\n");
+
+    let mut gross_cycle_counts: BTreeMap<SourceLineNumber, usize> = BTreeMap::new();
+    for (line, cycle_count) in history.lines.iter().zip(history.cycles.iter()) {
+        let prev_count = gross_cycle_counts.get(line).unwrap_or(&0);
+        gross_cycle_counts.insert(*line, *prev_count + cycle_count);
+    }
+    for (line, cycle_count) in gross_cycle_counts.iter() {
+        println!("line {line}: {cycle_count} cycles");
+    }
 }
 
 /// Helper function that performs the actual bytecode execution