lean_vm: add source-level failure stack traces

crates/lean_compiler/src/b_compile_intermediate.rs

@@ -502,6 +502,12 @@ fn compile_lines(
                 let new_fp_pos = compiler.stack_pos;
                 compiler.stack_pos += 1;
 
+                instructions.push(IntermediateInstruction::CallSite {
+                    caller: compiler.func_name.clone(),
+                    callee: callee_function_name.clone(),
+                    location: *location,
+                    return_label: return_label.clone(),
+                });
                 instructions.extend(emit_call_frame(
                     callee_function_name,
                     args,

crates/lean_compiler/src/c_compile_final.rs

@@ -13,6 +13,7 @@ impl IntermediateInstruction {
             | Self::HintWitness { .. }
             | Self::Inverse { .. }
             | Self::LocationReport { .. }
+            | Self::CallSite { .. }
             | Self::DebugAssert { .. }
             | Self::DerefHint { .. }
             | Self::PanicHint { .. }
@@ -54,6 +55,7 @@ pub fn compile_to_low_level_bytecode(
         .ok_or("Missing main function")?;
 
     let mut hints = BTreeMap::new();
+    let mut call_sites_by_return_pc = BTreeMap::new();
     let mut label_to_pc = BTreeMap::new();
 
     let exit_point = intermediate_bytecode
@@ -123,7 +125,14 @@ pub fn compile_to_low_level_bytecode(
     let mut instructions = Vec::new();
 
     for (pc_start, block) in code_blocks {
-        compile_block(&compiler, &block, pc_start, &mut instructions, &mut hints);
+        compile_block(
+            &compiler,
+            &block,
+            pc_start,
+            &mut instructions,
+            &mut hints,
+            &mut call_sites_by_return_pc,
+        );
     }
 
     debug_assert_eq!(instructions.len(), bytecode_size);
@@ -186,6 +195,7 @@ pub fn compile_to_low_level_bytecode(
         source_code,
         filepaths,
         pc_to_location,
+        call_sites_by_return_pc,
     })
 }
 
@@ -195,6 +205,7 @@ fn compile_block(
     pc_start: CodeAddress,
     low_level_bytecode: &mut Vec<Instruction>,
     hints: &mut BTreeMap<CodeAddress, Vec<Hint>>,
+    call_sites_by_return_pc: &mut BTreeMap<CodeAddress, CallSite>,
 ) {
     let try_as_mem_or_constant = |value: &IntermediateValue| {
         if let Some(cst) = try_as_constant(value, compiler) {
@@ -357,6 +368,28 @@ fn compile_block(
                 let hint = Hint::LocationReport { location };
                 hints.entry(pc).or_default().push(hint);
             }
+            IntermediateInstruction::CallSite {
+                caller,
+                callee,
+                location,
+                return_label,
+            } => {
+                let return_pc = compiler
+                    .label_to_pc
+                    .get(&return_label)
+                    .copied()
+                    .expect("Fatal: unresolved call return label");
+                call_sites_by_return_pc.insert(
+                    return_pc,
+                    CallSite {
+                        caller,
+                        callee,
+                        location,
+                        call_pc: return_pc.saturating_sub(1),
+                        return_pc,
+                    },
+                );
+            }
             IntermediateInstruction::DebugAssert {
                 expr,
                 location,

crates/lean_compiler/src/ir/instruction.rs

@@ -1,6 +1,8 @@
 use super::value::IntermediateValue;
 use crate::lang::{ConstExpression, MathOperation};
-use lean_vm::{BooleanExpr, CustomHint, HintWitnessDestination, Operation, PrecompileArgs, SourceLocation};
+use lean_vm::{
+    BooleanExpr, CustomHint, FunctionName, HintWitnessDestination, Label, Operation, PrecompileArgs, SourceLocation,
+};
 use std::fmt::{Display, Formatter};
 
 /// Core instruction type for the intermediate representation.
@@ -58,6 +60,13 @@ pub enum IntermediateInstruction {
     LocationReport {
         location: SourceLocation,
     },
+    // No-op metadata used to reconstruct source-level call stacks.
+    CallSite {
+        caller: FunctionName,
+        callee: FunctionName,
+        location: SourceLocation,
+        return_label: Label,
+    },
     DebugAssert {
         expr: BooleanExpr<IntermediateValue>,
         location: SourceLocation,
@@ -192,6 +201,15 @@ impl Display for IntermediateInstruction {
                 Ok(())
             }
             Self::LocationReport { .. } => Ok(()),
+            Self::CallSite {
+                caller,
+                callee,
+                location,
+                return_label,
+            } => write!(
+                f,
+                "call_site {caller} -> {callee} at {location}, returns to {return_label}"
+            ),
             Self::DebugAssert { expr, .. } => {
                 write!(f, "debug_assert {expr}")
             }

crates/lean_compiler/src/lib.rs

@@ -146,6 +146,23 @@ pub struct CompilationFlags {
     pub replacements: BTreeMap<String, String>,
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct CompileAndRunOptions {
+    /// Include VM profiling metadata in the execution summary.
+    pub profiler: bool,
+    /// Print a source-level stack trace to stderr when VM execution fails.
+    pub stack_trace: bool,
+}
+
+impl Default for CompileAndRunOptions {
+    fn default() -> Self {
+        Self {
+            profiler: false,
+            stack_trace: true,
+        }
+    }
+}
+
 pub fn try_compile_program_with_flags(
     input: &ProgramSource,
     flags: CompilationFlags,
@@ -184,7 +201,35 @@ pub fn try_compile_and_run(
     Ok(result.metadata.display())
 }
 
+pub fn try_compile_and_run_with_options(
+    input: &ProgramSource,
+    public_input: &[F; PUBLIC_INPUT_LEN],
+    options: CompileAndRunOptions,
+) -> Result<String, Error> {
+    let bytecode = try_compile_program(input)?;
+    let witness = ExecutionWitness::default();
+    let result = try_execute_bytecode_with_options(
+        &bytecode,
+        public_input,
+        &witness,
+        ExecutionOptions {
+            profiling: options.profiler,
+            stack_trace: options.stack_trace,
+        },
+    )?;
+    Ok(result.metadata.display())
+}
+
 pub fn compile_and_run(input: &ProgramSource, public_input: &[F; PUBLIC_INPUT_LEN], profiler: bool) {
     let summary = try_compile_and_run(input, public_input, profiler).unwrap();
     println!("{summary}");
 }
+
+pub fn compile_and_run_with_options(
+    input: &ProgramSource,
+    public_input: &[F; PUBLIC_INPUT_LEN],
+    options: CompileAndRunOptions,
+) {
+    let summary = try_compile_and_run_with_options(input, public_input, options).unwrap();
+    println!("{summary}");
+}

crates/lean_compiler/tests/test_compiler.rs

@@ -1,4 +1,4 @@
-use std::time::Instant;
+use std::{process::Command, time::Instant};
 
 use backend::{BasedVectorSpace, PrimeCharacteristicRing};
 use lean_compiler::*;
@@ -257,6 +257,200 @@ def main():
     compile_and_run(&ProgramSource::Raw(program.to_string()), &[F::ZERO; DIGEST_LEN], false);
 }
 
+#[test]
+fn nested_runtime_failure_prints_full_call_stack() {
+    let current_exe = std::env::current_exe().expect("current test executable path");
+    let output = Command::new(current_exe)
+        .arg("--exact")
+        .arg("child_nested_runtime_failure_prints_stack_trace")
+        .arg("--nocapture")
+        .env("LEANVM_STACK_TRACE_CHILD", "1")
+        .output()
+        .expect("run child stack trace test");
+
+    assert!(
+        output.status.success(),
+        "child test failed\nstdout:\n{}\nstderr:\n{}",
+        String::from_utf8_lossy(&output.stdout),
+        String::from_utf8_lossy(&output.stderr),
+    );
+
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert_stack_frames(&stderr, &["d()", "c()", "b()", "a()", "main()"]);
+    assert!(
+        stderr.contains("pc="),
+        "stack trace should include pc values\nstderr:\n{stderr}",
+    );
+    assert!(
+        stderr.contains(":15"),
+        "stack trace should include the failing source line\nstderr:\n{stderr}",
+    );
+}
+
+#[test]
+fn parallel_runtime_failure_prints_worker_call_stack() {
+    let current_exe = std::env::current_exe().expect("current test executable path");
+    let output = Command::new(current_exe)
+        .arg("--exact")
+        .arg("child_parallel_runtime_failure_prints_stack_trace")
+        .arg("--nocapture")
+        .env("LEANVM_STACK_TRACE_CHILD", "1")
+        .output()
+        .expect("run child parallel stack trace test");
+
+    assert!(
+        output.status.success(),
+        "child test failed\nstdout:\n{}\nstderr:\n{}",
+        String::from_utf8_lossy(&output.stdout),
+        String::from_utf8_lossy(&output.stderr),
+    );
+
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert_stack_frames(&stderr, &["fail()", "main()"]);
+    assert!(
+        stderr.contains("pc="),
+        "parallel stack trace should include pc values\nstderr:\n{stderr}",
+    );
+}
+
+fn assert_stack_frames(stderr: &str, expected_frames: &[&str]) {
+    let stack = stderr
+        .split("CALL STACK")
+        .nth(1)
+        .unwrap_or_else(|| panic!("stderr should contain a CALL STACK section\nstderr:\n{stderr}"));
+    let frame_lines: Vec<&str> = stack.lines().filter(|line| line.contains("pc=")).collect();
+
+    assert_eq!(
+        frame_lines.len(),
+        expected_frames.len(),
+        "stack trace should contain exactly the expected frames\nframes:\n{}\nstderr:\n{stderr}",
+        frame_lines.join("\n"),
+    );
+    for (line, expected) in frame_lines.iter().zip(expected_frames) {
+        assert!(
+            line.contains(expected),
+            "stack frame should contain {expected}\nframe: {line}\nstderr:\n{stderr}",
+        );
+    }
+}
+
+#[test]
+fn child_nested_runtime_failure_prints_stack_trace() {
+    if std::env::var_os("LEANVM_STACK_TRACE_CHILD").is_none() {
+        return;
+    }
+
+    let program = r#"def main():
+    _ = a()
+    return
+
+def a():
+    return b()
+
+def b():
+    return c()
+
+def c():
+    return d()
+
+def d():
+    debug_assert(0 == 1)
+    return 0
+"#;
+
+    let result = try_compile_and_run_with_options(
+        &ProgramSource::Raw(program.to_string()),
+        &[F::ZERO; DIGEST_LEN],
+        CompileAndRunOptions {
+            profiler: false,
+            stack_trace: true,
+        },
+    );
+
+    assert!(
+        matches!(result, Err(Error::Runtime(_))),
+        "program should fail at VM runtime inside d(), got {result:?}",
+    );
+}
+
+#[test]
+fn child_parallel_runtime_failure_prints_stack_trace() {
+    if std::env::var_os("LEANVM_STACK_TRACE_CHILD").is_none() {
+        return;
+    }
+
+    let program = r#"def main():
+    n = 4
+    for i in parallel_range(0, n):
+        fail(i)
+    return
+
+def fail(i):
+    if i == 2:
+        debug_assert(0 == 1)
+    return
+"#;
+
+    let result = try_compile_and_run_with_options(
+        &ProgramSource::Raw(program.to_string()),
+        &[F::ZERO; DIGEST_LEN],
+        CompileAndRunOptions {
+            profiler: false,
+            stack_trace: true,
+        },
+    );
+
+    assert!(
+        matches!(result, Err(Error::Runtime(_))),
+        "parallel program should fail at VM runtime inside fail(), got {result:?}",
+    );
+}
+
+#[test]
+fn legacy_runtime_failure_respects_stack_trace_env_var() {
+    let current_exe = std::env::current_exe().expect("current test executable path");
+    let output = Command::new(current_exe)
+        .arg("--exact")
+        .arg("child_legacy_runtime_failure_respects_stack_trace_env_var")
+        .arg("--nocapture")
+        .env("LEANVM_STACK_TRACE_CHILD", "1")
+        .env("LEANVM_STACK_TRACE", "0")
+        .output()
+        .expect("run child legacy env-var stack trace test");
+
+    assert!(
+        output.status.success(),
+        "child test failed\nstdout:\n{}\nstderr:\n{}",
+        String::from_utf8_lossy(&output.stdout),
+        String::from_utf8_lossy(&output.stderr),
+    );
+
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert!(
+        !stderr.contains("CALL STACK"),
+        "LEANVM_STACK_TRACE=0 should disable legacy stack trace printing\nstderr:\n{stderr}",
+    );
+}
+
+#[test]
+fn child_legacy_runtime_failure_respects_stack_trace_env_var() {
+    if std::env::var_os("LEANVM_STACK_TRACE_CHILD").is_none() {
+        return;
+    }
+
+    let program = r#"def main():
+    debug_assert(0 == 1)
+    return
+"#;
+
+    let result = try_compile_and_run(&ProgramSource::Raw(program.to_string()), &[F::ZERO; DIGEST_LEN], false);
+
+    assert!(
+        matches!(result, Err(Error::Runtime(_))),
+        "program should fail at VM runtime, got {result:?}",
+    );
+}
+
 #[test]
 #[rustfmt::skip]
 fn test_soundness_suite() {