flamegraph: fix tail-call misattribution, trie-based fold, addr2line enrichment

bench_vs/lambda/recursion/Cargo.toml

@@ -11,3 +11,6 @@ lambda-vm-prover = { path = "../../../prover", default-features = false, feature
 ] }
 lambda-vm-syscalls = { path = "../../../syscalls" }
 postcard = { version = "1.0", features = ["alloc"] }
+
+[profile.release]
+debug = 2

bin/cli/Cargo.toml

@@ -10,6 +10,7 @@ prover = { path = "../../prover", package = "lambda-vm-prover" }
 stark = { path = "../../crypto/stark" }
 clap = { version = "4.3.10", features = ["derive"] }
 bincode = "1"
+tempfile = "3"
 tikv-jemallocator = "0.6"
 tikv-jemalloc-ctl = { version = "0.6", features = ["stats"], optional = true }
 env_logger = "0.11"

bin/cli/src/main.rs

@@ -2,19 +2,15 @@
 
 use std::fs::File;
 use std::io::{BufWriter, Write};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::process::ExitCode;
 use std::time::Instant;
 
 use clap::{Parser, Subcommand, ValueHint};
 
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
-use executor::{
-    elf::{Elf, SymbolTable},
-    flamegraph::FlamegraphGenerator,
-    vm::execution::Executor,
-};
+use executor::{elf::Elf, flamegraph::FlamegraphGenerator, vm::execution::Executor};
 use prover::VmProof;
 use stark::proof::options::GoldilocksCubicProofOptions;
 
@@ -112,6 +108,22 @@ enum Commands {
         #[arg(long, value_hint = ValueHint::FilePath)]
         flamegraph: Option<PathBuf>,
 
+        /// Key the folded stacks by raw hex address instead of resolving
+        /// through the ELF symtab (pairs with scripts/enrich_flamegraph.py).
+        /// Only meaningful with --flamegraph.
+        #[arg(long, requires = "flamegraph")]
+        flamegraph_raw: bool,
+
+        /// Checkpoint the flamegraph's folded output to --flamegraph every N
+        /// cycles, so a killed run still leaves usable (partial) output on
+        /// disk. Only meaningful with --flamegraph.
+        #[arg(long, requires = "flamegraph")]
+        flamegraph_checkpoint_cycles: Option<u64>,
+
+        /// Stop execution early once at least this many cycles have run.
+        #[arg(long)]
+        cycle_budget: Option<u64>,
+
         /// Print the dynamic instruction (cycle) count
         #[arg(long)]
         cycles: bool,
@@ -207,8 +219,21 @@ fn main() -> ExitCode {
             elf,
             private_input,
             flamegraph,
+            flamegraph_raw,
+            flamegraph_checkpoint_cycles,
+            cycle_budget,
+            cycles,
+        } => cmd_execute(
+            elf,
+            private_input,
+            FlamegraphCliOptions {
+                path: flamegraph,
+                raw: flamegraph_raw,
+                checkpoint_cycles: flamegraph_checkpoint_cycles,
+            },
+            cycle_budget,
             cycles,
-        } => cmd_execute(elf, private_input, flamegraph, cycles),
+        ),
         Commands::Prove {
             elf,
             output,
@@ -272,10 +297,53 @@ fn count_cycles(elf_data: &[u8], private_inputs: &[u8]) -> Result<u64, String> {
         .map_err(|e| format!("Execution failed during cycle count: {e:?}"))
 }
 
+/// Write the flamegraph's current (possibly partial) folded output to
+/// `output_path`, replacing any previous contents. Used both for the final
+/// write and for periodic checkpoints during a long run.
+///
+/// Writes to a `tempfile` in the same directory, flushes it, then persists
+/// (renames) it over `output_path` — the whole file is replaced atomically,
+/// so a kill mid-write can never leave `output_path` empty or torn (the
+/// previous good checkpoint stays put until the new one is fully on disk).
+fn write_flamegraph_checkpoint(
+    output_path: &PathBuf,
+    generator: &FlamegraphGenerator,
+    raw: bool,
+) -> Result<(), String> {
+    let dir = output_path.parent().unwrap_or_else(|| Path::new("."));
+    let tmp = tempfile::NamedTempFile::new_in(dir)
+        .map_err(|e| format!("Failed to create temp output file: {e}"))?;
+
+    let mut writer = BufWriter::new(tmp.as_file());
+    let result = if raw {
+        generator.write_folded_raw(&mut writer)
+    } else {
+        generator.write_folded(&mut writer)
+    };
+    result.map_err(|e| format!("Failed to write flamegraph output: {e:?}"))?;
+    writer
+        .flush()
+        .map_err(|e| format!("Failed to flush flamegraph output: {e}"))?;
+    drop(writer);
+
+    tmp.persist(output_path)
+        .map_err(|e| format!("Failed to replace {output_path:?} with temp output: {e}"))?;
+    Ok(())
+}
+
+/// Flamegraph-related flags grouped so `cmd_execute` doesn't need a flat
+/// 8-argument signature.
+struct FlamegraphCliOptions {
+    path: Option<PathBuf>,
+    raw: bool,
+    checkpoint_cycles: Option<u64>,
+}
+
 fn cmd_execute(
     elf_path: PathBuf,
     private_input_path: Option<PathBuf>,
-    flamegraph_path: Option<PathBuf>,
+    flamegraph: FlamegraphCliOptions,
+    cycle_budget: Option<u64>,
     cycles: bool,
 ) -> ExitCode {
     let elf_data = match std::fs::read(&elf_path) {
@@ -302,71 +370,94 @@ fn cmd_execute(
         }
     };
 
-    let mut executor = match Executor::new(&program, private_inputs) {
-        Ok(e) => e,
-        Err(e) => {
-            eprintln!("Failed to create executor: {:?}", e);
-            return ExitCode::FAILURE;
-        }
-    };
-
-    // Set up flamegraph generator if requested
-    let mut generator = flamegraph_path.as_ref().map(|_| {
-        let symbols = SymbolTable::parse(&elf_data);
-        FlamegraphGenerator::new(symbols, program.entry_point)
-    });
+    let cycle_count = if let Some(ref output_path) = flamegraph.path {
+        // Shared execute+flamegraph path (executor::flamegraph) instead of
+        // hand-rolling the SymbolTable/Executor/drive-loop wiring here.
+        let mut next_checkpoint = flamegraph.checkpoint_cycles;
+        let result = executor::flamegraph::run_with_flamegraph(
+            &elf_data,
+            &program,
+            private_inputs,
+            cycle_budget,
+            |total_cycles, generator| {
+                let Some(threshold) = next_checkpoint else {
+                    return;
+                };
+                if total_cycles < threshold {
+                    return;
+                }
+                if let Err(e) = write_flamegraph_checkpoint(output_path, generator, flamegraph.raw)
+                {
+                    eprintln!("Warning: flamegraph checkpoint failed: {e}");
+                }
+                next_checkpoint = flamegraph.checkpoint_cycles.map(|step| threshold + step);
+            },
+        );
 
-    // Execute in chunks, counting cycles and (if requested) feeding the flamegraph.
-    let mut cycle_count: u64 = 0;
-    loop {
-        let logs = match executor.resume() {
-            Ok(logs) => logs,
+        let (generator, result) = result;
+        let total_cycles = match result {
+            Ok(total_cycles) => total_cycles,
             Err(e) => {
                 eprintln!("Execution failed: {:?}", e);
+                // Best-effort: persist whatever the generator accumulated
+                // before the fault instead of discarding it outright.
+                match write_flamegraph_checkpoint(output_path, &generator, flamegraph.raw) {
+                    Ok(()) => eprintln!(
+                        "Partial flamegraph written to {:?} ({} instructions)",
+                        output_path,
+                        generator.total_instructions()
+                    ),
+                    Err(e) => eprintln!("Warning: failed to write partial flamegraph: {e}"),
+                }
                 return ExitCode::FAILURE;
             }
         };
-        match logs {
-            Some(logs) => {
-                cycle_count += logs.len() as u64;
-                if let Some(ref mut fg) = generator {
-                    let logs: Vec<_> = logs.to_vec();
-                    if let Err(e) = fg.process_logs(&logs, &executor.instructions) {
-                        eprintln!("Failed to process logs for flamegraph: {:?}", e);
-                        return ExitCode::FAILURE;
-                    }
-                }
-            }
-            None => break,
-        }
-    }
 
-    if let Err(e) = executor.finish() {
-        eprintln!("Failed to finish execution: {:?}", e);
-        return ExitCode::FAILURE;
-    }
+        if let Err(e) = write_flamegraph_checkpoint(output_path, &generator, flamegraph.raw) {
+            eprintln!("{e}");
+            return ExitCode::FAILURE;
+        }
+        eprintln!(
+            "Flamegraph written to {:?} ({} instructions)",
+            output_path,
+            generator.total_instructions()
+        );
 
-    // Write flamegraph output if requested
-    if let (Some(output_path), Some(generator)) = (flamegraph_path, generator) {
-        let file = match File::create(&output_path) {
-            Ok(f) => f,
+        total_cycles
+    } else {
+        let mut executor = match Executor::new(&program, private_inputs) {
+            Ok(e) => e,
             Err(e) => {
-                eprintln!("Failed to create flamegraph output file: {}", e);
+                eprintln!("Failed to create executor: {:?}", e);
                 return ExitCode::FAILURE;
             }
         };
-        let mut writer = BufWriter::new(file);
-        if let Err(e) = generator.write_folded(&mut writer) {
-            eprintln!("Failed to write flamegraph output: {:?}", e);
+
+        let mut cycle_count: u64 = 0;
+        loop {
+            let logs = match executor.resume_budgeted(cycle_count, cycle_budget) {
+                Ok(logs) => logs,
+                Err(e) => {
+                    eprintln!("Execution failed: {:?}", e);
+                    return ExitCode::FAILURE;
+                }
+            };
+            match logs {
+                Some(logs) => cycle_count += logs.len() as u64,
+                None => break,
+            }
+            if cycle_budget.is_some_and(|budget| cycle_count >= budget) {
+                break;
+            }
+        }
+
+        if let Err(e) = executor.finish() {
+            eprintln!("Failed to finish execution: {:?}", e);
             return ExitCode::FAILURE;
         }
 
-        eprintln!(
-            "Flamegraph written to {:?} ({} instructions)",
-            output_path,
-            generator.total_instructions()
-        );
-    }
+        cycle_count
+    };
 
     if cycles {
         println!("Cycles: {}", cycle_count);

executor/src/elf.rs

@@ -587,6 +587,30 @@ impl SymbolTable {
         }
     }
 
+    /// Like [`Self::lookup`], but also returns the exclusive upper bound of the
+    /// addresses that resolve to the returned function — its size-end, capped
+    /// at the next symbol's start so overlapping/nested symbols are respected.
+    /// Every address in `[func.address, end)` resolves to `func` via `lookup`,
+    /// so callers can cache the range and skip re-running `lookup` inside it.
+    pub fn lookup_range(&self, address: u64) -> Option<(&FunctionSymbol, u64)> {
+        let idx = match self.functions.binary_search_by_key(&address, |f| f.address) {
+            Ok(i) => i,
+            Err(0) => return None,
+            Err(i) => i - 1,
+        };
+        let func = &self.functions[idx];
+        let size_end = if func.size == 0 {
+            u64::MAX
+        } else {
+            func.address + func.size
+        };
+        if address >= size_end {
+            return None;
+        }
+        let next_start = self.functions.get(idx + 1).map_or(u64::MAX, |f| f.address);
+        Some((func, size_end.min(next_start)))
+    }
+
     /// Check if the symbol table is empty
     pub fn is_empty(&self) -> bool {
         self.functions.is_empty()