perf(vm): add peephole pass to elide dead Load; Pop pairs ✂️

ndc_interpreter/src/lib.rs

@@ -125,6 +125,17 @@ impl Interpreter {
         Ok(Compiler::compile(expressions.into_iter())?)
     }
 
+    /// Like [`Self::compile_str`] but skips the peephole optimizer.
+    /// Used by the compiler-tests crate to assert against raw bytecode.
+    pub fn compile_str_unoptimized(
+        &mut self,
+        input: &str,
+    ) -> Result<CompiledFunction, InterpreterError> {
+        let source_id = self.source_db.add("<input>", input);
+        let (expressions, _) = self.parse_and_analyse(input, source_id)?;
+        Ok(Compiler::compile_unoptimized(expressions.into_iter())?)
+    }
+
     pub fn disassemble_str(&mut self, input: &str) -> Result<String, InterpreterError> {
         let compiled = self.compile_str(input)?;
         let mut out = String::new();

ndc_vm/Cargo.toml

@@ -7,9 +7,10 @@ version.workspace = true
 trace = []
 
 [dependencies]
-thiserror.workspace = true
-ndc_parser.workspace = true
-ndc_lexer.workspace = true
+ahash.workspace = true
 ndc_core.workspace = true
+ndc_lexer.workspace = true
+ndc_parser.workspace = true
 num.workspace = true
 ordered-float.workspace = true
+thiserror.workspace = true

ndc_vm/src/chunk.rs

@@ -1,4 +1,5 @@
 use crate::Value;
+use ahash::AHashSet;
 use ndc_core::hash_map::HashMap;
 use ndc_lexer::Span;
 use ndc_parser::CaptureSource;
@@ -36,7 +37,6 @@ impl JumpTarget {
     /// Advance `ip` by this offset using wrapping signed arithmetic.
     #[inline]
     pub fn apply(self, ip: usize) -> usize {
-        // TODO: is this bad for perf?
         match self {
             JumpTarget::Offset(o) => ip.wrapping_add_signed(o),
             JumpTarget::Label(_) => panic!("cannot apply instruction pointer to unresolved label"),
@@ -267,6 +267,49 @@ impl OptimizerIr {
         }
     }
 
+    pub(crate) fn peephole(&mut self) {
+        let targets: AHashSet<LabelId> = self
+            .code
+            .iter()
+            .filter_map(|(_, op, _)| match op {
+                OpCode::Jump(JumpTarget::Label(l))
+                | OpCode::JumpIfFalse(JumpTarget::Label(l))
+                | OpCode::JumpIfTrue(JumpTarget::Label(l))
+                | OpCode::IterNext(JumpTarget::Label(l)) => Some(*l),
+                _ => None,
+            })
+            .collect();
+
+        loop {
+            let mut out = Vec::with_capacity(self.code.len());
+            let mut i = 0;
+            let mut removed = 0;
+            while i < self.code.len() {
+                let elide = i + 1 < self.code.len()
+                    && matches!(
+                        self.code[i].1,
+                        OpCode::Constant(_) | OpCode::GetGlobal(_) | OpCode::GetLocal(_)
+                    )
+                    && matches!(self.code[i + 1].1, OpCode::Pop)
+                    && !targets.contains(&self.code[i].0)
+                    && !targets.contains(&self.code[i + 1].0);
+                if elide {
+                    i += 2;
+                    removed += 2;
+                } else {
+                    out.push(self.code[i].clone());
+                    i += 1;
+                }
+            }
+
+            self.code = out;
+
+            if removed == 0 {
+                break;
+            }
+        }
+    }
+
     pub(crate) fn add_constant(&mut self, value: Value) -> usize {
         self.constants.push(value);
         self.constants.len() - 1

ndc_vm/src/compiler.rs

@@ -9,29 +9,69 @@ use ndc_parser::{
 };
 use std::rc::Rc;
 
-#[derive(Default, Clone)]
+#[derive(Clone)]
 pub struct Compiler {
     ir: OptimizerIr,
     num_locals: usize,
     loop_stack: Vec<LoopContext>,
     allow_return: bool,
+    optimize: bool,
+}
+
+impl Default for Compiler {
+    fn default() -> Self {
+        Self {
+            ir: OptimizerIr::default(),
+            num_locals: 0,
+            loop_stack: Vec::new(),
+            allow_return: false,
+            optimize: true,
+        }
+    }
 }
 
 impl Compiler {
     pub fn compile(
         expressions: impl Iterator<Item = ExpressionLocation>,
     ) -> Result<CompiledFunction, CompileError> {
-        Ok(Self::compile_resumable(expressions)?.0)
+        let mut compiler = Self::default();
+        for expr_loc in expressions {
+            compiler.compile_expr(expr_loc)?;
+        }
+        Ok(compiler.finish()?.0)
+    }
+
+    /// Compile expressions without running the peephole optimizer. Useful
+    /// for tests that want to inspect the raw compiler output, and for
+    /// debugging tools (e.g. a future `--no-optimize` disassembler flag).
+    pub fn compile_unoptimized(
+        expressions: impl Iterator<Item = ExpressionLocation>,
+    ) -> Result<CompiledFunction, CompileError> {
+        let mut compiler = Self {
+            optimize: false,
+            ..Default::default()
+        };
+        for expr_loc in expressions {
+            compiler.compile_expr(expr_loc)?;
+        }
+        Ok(compiler.finish()?.0)
     }
 
     /// Compile expressions and return both the finished function and a
     /// checkpoint that can be passed to `resume` to append more code later.
     /// The checkpoint is the compiler state *before* the `Halt` instruction,
     /// so `resume` can extend the bytecode without re-running old instructions.
+    ///
+    /// Optimization is disabled on this path: the REPL's resume-from-
+    /// halt machinery relies on `halt_ip` matching the position of `Halt` in
+    /// the emitted chunk, and shifting instructions invalidates that.
     pub fn compile_resumable(
         expressions: impl Iterator<Item = ExpressionLocation>,
     ) -> Result<(CompiledFunction, Self), CompileError> {
-        let mut compiler = Self::default();
+        let mut compiler = Self {
+            optimize: false,
+            ..Default::default()
+        };
         for expr_loc in expressions {
             compiler.compile_expr(expr_loc)?;
         }
@@ -71,6 +111,10 @@ impl Compiler {
     fn finish(mut self) -> Result<(CompiledFunction, Self), CompileError> {
         let checkpoint = self.clone();
         self.ir.write(OpCode::Halt, Span::synthetic());
+        if self.optimize {
+            self.ir.peephole();
+        }
+
         let function = CompiledFunction {
             name: None,
             static_type: StaticType::Function {
@@ -718,11 +762,16 @@ impl Compiler {
         let mut fn_compiler = Self {
             num_locals: num_params,
             allow_return: true,
+            optimize: self.optimize,
             ..Default::default()
         };
         fn_compiler.compile_expr(body)?;
         fn_compiler.ir.write(OpCode::Return, Span::synthetic());
 
+        if fn_compiler.optimize {
+            fn_compiler.ir.peephole();
+        }
+
         let compiled = CompiledFunction {
             name,
             static_type,

ndc_vm/src/value/function.rs

@@ -64,6 +64,10 @@ impl CompiledFunction {
     pub fn opcodes(&self) -> &[OpCode] {
         self.body.opcodes()
     }
+
+    pub fn body(&self) -> &Chunk {
+        &self.body
+    }
 }
 impl Function {
     pub fn prototype(&self) -> Option<&Rc<CompiledFunction>> {

tests/compiler/tests/compiler.rs

@@ -5,12 +5,16 @@ use ndc_vm::chunk::OpCode;
 use ndc_vm::chunk::OpCode::*;
 use ndc_vm::compiler::Compiler;
 
+// These helpers compile without the peephole optimizer so the tests
+// document the raw compiler output as a specification. Optimizer behaviour
+// is exercised in `tests/optimizer.rs`.
+
 fn compile(input: &str) -> Vec<OpCode> {
     let tokens = Lexer::new(input, SourceId::SYNTHETIC)
         .collect::<Result<Vec<_>, _>>()
         .expect("lex failed");
     let expressions = Parser::from_tokens(tokens).parse().expect("parse failed");
-    Compiler::compile(expressions.into_iter())
+    Compiler::compile_unoptimized(expressions.into_iter())
         .expect("compile failed")
         .opcodes()
         .to_vec()
@@ -19,7 +23,7 @@ fn compile(input: &str) -> Vec<OpCode> {
 fn compile_with_analysis(input: &str) -> Vec<OpCode> {
     let mut interp = ndc_interpreter::Interpreter::capturing();
     interp
-        .compile_str(input)
+        .compile_str_unoptimized(input)
         .expect("compile failed")
         .opcodes()
         .to_vec()

tests/compiler/tests/optimizer.rs

@@ -0,0 +1,130 @@
+use ndc_lexer::{Lexer, SourceId};
+use ndc_parser::Parser;
+use ndc_stdlib as _;
+use ndc_vm::chunk::OpCode;
+use ndc_vm::chunk::OpCode::*;
+use ndc_vm::compiler::Compiler;
+
+fn parse(input: &str) -> Vec<ndc_parser::ExpressionLocation> {
+    let tokens = Lexer::new(input, SourceId::SYNTHETIC)
+        .collect::<Result<Vec<_>, _>>()
+        .expect("lex failed");
+    Parser::from_tokens(tokens).parse().expect("parse failed")
+}
+
+fn unoptimized(input: &str) -> Vec<OpCode> {
+    Compiler::compile_unoptimized(parse(input).into_iter())
+        .expect("compile failed")
+        .opcodes()
+        .to_vec()
+}
+
+fn optimized(input: &str) -> Vec<OpCode> {
+    Compiler::compile(parse(input).into_iter())
+        .expect("compile failed")
+        .opcodes()
+        .to_vec()
+}
+
+// `5;` — a discarded constant. The peephole eliminates the `Constant; Pop`
+// pair entirely, leaving just `Halt`.
+#[test]
+fn elides_constant_pop_statement() {
+    assert_eq!(unoptimized("5;"), [Constant(0), Pop, Halt]);
+    assert_eq!(optimized("5;"), [Halt]);
+}
+
+// A `Pop` that's the target of a conditional jump must not be elided even
+// when the preceding instruction is a `Load` — the jumper relies on it to
+// drop the condition value off the stack. Here, `JumpIfFalse` targets the
+// false-branch `Pop`, and that `Pop` survives optimization.
+#[test]
+fn preserves_pop_as_jump_target() {
+    let raw = unoptimized("if true { 3 } else { 3; }");
+    let opt = optimized("if true { 3 } else { 3; }");
+
+    // Optimization removes the `3; → Constant; Pop` in the else branch.
+    assert!(
+        opt.len() < raw.len(),
+        "optimizer should shrink the chunk; raw={raw:?} opt={opt:?}"
+    );
+
+    // Both branch-condition Pops survive — one per arm. Raw has a third Pop
+    // (the discarded statement in the else branch); optimized does not.
+    let pop_count = |ops: &[OpCode]| ops.iter().filter(|op| matches!(op, Pop)).count();
+    assert_eq!(pop_count(&raw), 3, "raw should have 3 Pops");
+    assert_eq!(pop_count(&opt), 2, "optimized should have 2 Pops");
+
+    // The optimized chunk still has the conditional jump and the
+    // unconditional skip — control flow is intact.
+    assert!(opt.iter().any(|op| matches!(op, JumpIfFalse(_))));
+    assert!(opt.iter().any(|op| matches!(op, Jump(_))));
+}
+
+// When the caller asks for unoptimized output, nested function bodies must
+// also be raw — otherwise debug/inspection callers would silently see
+// optimized bytecode inside the constants table.
+fn first_nested_function_opcodes(compiled: &ndc_vm::value::CompiledFunction) -> Vec<OpCode> {
+    for (_, _, val) in compiled.body().iter() {
+        let Some(ndc_vm::Value::Object(obj)) = val else {
+            continue;
+        };
+        if let ndc_vm::Object::Function(ndc_vm::value::Function::Compiled(f)) = obj.as_ref() {
+            return f.opcodes().to_vec();
+        }
+    }
+    panic!("expected a compiled function constant in the top-level chunk");
+}
+
+#[test]
+fn nested_function_bodies_respect_unoptimized_flag() {
+    let mut interp = ndc_interpreter::Interpreter::capturing();
+    let raw = interp
+        .compile_str_unoptimized("fn f() { 1; }")
+        .expect("compile failed");
+
+    let mut interp = ndc_interpreter::Interpreter::capturing();
+    let opt = interp.compile_str("fn f() { 1; }").expect("compile failed");
+
+    // Raw body for `1;` followed by an implicit unit Return:
+    // `Constant; Pop; Constant; Return`.
+    assert_eq!(
+        first_nested_function_opcodes(&raw),
+        [Constant(0), Pop, Constant(1), Return],
+        "nested function body should be raw under compile_str_unoptimized",
+    );
+
+    // Optimized body: peephole drops the leading `Constant; Pop`.
+    assert_eq!(
+        first_nested_function_opcodes(&opt),
+        [Constant(1), Return],
+        "nested function body should be optimized under compile_str",
+    );
+}
+
+// `{ 1; 2; 3; }` is three discarded statements followed by the block's
+// unit result. Every `Load; Pop` pair is elided; only the trailing block
+// result remains. Exercises the peephole loop catching multiple pairs.
+#[test]
+fn elides_multiple_load_pop_pairs() {
+    let raw = unoptimized("{ 1; 2; 3; }");
+    let opt = optimized("{ 1; 2; 3; }");
+
+    // Raw: three `Constant; Pop` pairs plus the final unit `Constant` plus Halt.
+    assert_eq!(
+        raw,
+        [
+            Constant(0),
+            Pop,
+            Constant(1),
+            Pop,
+            Constant(2),
+            Pop,
+            Constant(3),
+            Halt,
+        ]
+    );
+
+    // Optimized: just the final unit `Constant` plus Halt.
+    assert_eq!(opt, [Constant(3), Halt]);
+}