M2: Add call graph analysis, recursion detection, and optimizer

Analyzer extensions:
- Call graph construction from function bodies and state handlers
- DFS-based recursion detection (direct and mutual) with E0402 errors
- Max call depth computation per entry point with E0401 enforcement
- Function declarations registered as symbols (E0503 for undefined calls)
- Collects calls from all statement/expression types recursively

Optimizer (new module):
- Constant folding: evaluate known-constant arithmetic at compile time
- Dead code elimination: remove ops with unused destination temps
- Both operate per-basic-block in a single pass

171 tests total (22 new: 6 analyzer + 11 IR lowering + 5 optimizer)

https://claude.ai/code/session_01W6eQFStA66EuMKHUFo2rx3

src/analyzer/mod.rs

@@ -1,7 +1,7 @@
 #[cfg(test)]
 mod tests;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use crate::errors::{Diagnostic, ErrorCode};
 use crate::lexer::Span;
@@ -29,8 +29,13 @@ pub struct AnalysisResult {
     pub symbols: HashMap<String, Symbol>,
     pub var_allocations: Vec<VarAllocation>,
     pub diagnostics: Vec<Diagnostic>,
+    pub call_graph: HashMap<String, Vec<String>>,
+    pub max_depths: HashMap<String, u32>,
 }
 
+/// Default call stack depth limit for the NES runtime.
+const DEFAULT_STACK_DEPTH: u32 = 8;
+
 /// Analyze a parsed program for semantic errors.
 pub fn analyze(program: &Program) -> AnalysisResult {
     let mut analyzer = Analyzer {
@@ -39,6 +44,9 @@ pub fn analyze(program: &Program) -> AnalysisResult {
         diagnostics: Vec::new(),
         next_ram_addr: 0x0300, // $0300 is first usable RAM after OAM buffer
         next_zp_addr: 0x10,    // $10 is first usable zero-page after reserved area
+        call_graph: HashMap::new(),
+        max_depths: HashMap::new(),
+        stack_depth_limit: DEFAULT_STACK_DEPTH,
     };
     analyzer.analyze_program(program);
 
@@ -46,6 +54,8 @@ pub fn analyze(program: &Program) -> AnalysisResult {
         symbols: analyzer.symbols,
         var_allocations: analyzer.var_allocations,
         diagnostics: analyzer.diagnostics,
+        call_graph: analyzer.call_graph,
+        max_depths: analyzer.max_depths,
     }
 }
 
@@ -55,6 +65,9 @@ struct Analyzer {
     diagnostics: Vec<Diagnostic>,
     next_ram_addr: u16,
     next_zp_addr: u8,
+    call_graph: HashMap<String, Vec<String>>,
+    max_depths: HashMap<String, u32>,
+    stack_depth_limit: u32,
 }
 
 impl Analyzer {
@@ -69,6 +82,11 @@ impl Analyzer {
             self.register_var(var);
         }
 
+        // Register functions as symbols
+        for fun in &program.functions {
+            self.register_fun(fun);
+        }
+
         // Register state-local variables
         for state in &program.states {
             for var in &state.locals {
@@ -100,6 +118,27 @@ impl Analyzer {
                 self.check_block(block, &state_names);
             }
         }
+
+        // Type-check function bodies
+        for fun in &program.functions {
+            self.check_block(&fun.body, &state_names);
+        }
+
+        // Build call graph
+        self.build_call_graph(program);
+
+        // Detect recursion
+        let recursive_fns = detect_recursion(&self.call_graph);
+        for name in &recursive_fns {
+            self.diagnostics.push(Diagnostic::error(
+                ErrorCode::E0402,
+                format!("recursion detected in function '{name}'"),
+                program.span,
+            ));
+        }
+
+        // Compute max call depths from entry points (state handlers)
+        self.compute_max_depths(program);
     }
 
     fn register_const(&mut self, c: &ConstDecl) {
@@ -152,6 +191,27 @@ impl Analyzer {
         });
     }
 
+    fn register_fun(&mut self, fun: &FunDecl) {
+        if self.symbols.contains_key(&fun.name) {
+            self.diagnostics.push(Diagnostic::error(
+                ErrorCode::E0501,
+                format!("duplicate declaration of '{}'", fun.name),
+                fun.span,
+            ));
+            return;
+        }
+        let sym_type = fun.return_type.clone().unwrap_or(NesType::U8);
+        self.symbols.insert(
+            fun.name.clone(),
+            Symbol {
+                name: fun.name.clone(),
+                sym_type,
+                is_const: false,
+                span: fun.span,
+            },
+        );
+    }
+
     fn allocate_ram(&mut self, size: u16) -> u16 {
         // For M1: simple linear allocator using zero-page for u8 vars
         if size == 1 && self.next_zp_addr < 0xFF {
@@ -165,6 +225,75 @@ impl Analyzer {
         }
     }
 
+    fn build_call_graph(&mut self, program: &Program) {
+        // Record calls from each function body
+        for fun in &program.functions {
+            let callees = collect_calls(&fun.body);
+            self.call_graph.insert(fun.name.clone(), callees);
+        }
+
+        // Record calls from each state handler
+        for state in &program.states {
+            if let Some(block) = &state.on_enter {
+                let key = format!("{}::enter", state.name);
+                let callees = collect_calls(block);
+                self.call_graph.insert(key, callees);
+            }
+            if let Some(block) = &state.on_exit {
+                let key = format!("{}::exit", state.name);
+                let callees = collect_calls(block);
+                self.call_graph.insert(key, callees);
+            }
+            if let Some(block) = &state.on_frame {
+                let key = format!("{}::frame", state.name);
+                let callees = collect_calls(block);
+                self.call_graph.insert(key, callees);
+            }
+        }
+    }
+
+    fn compute_max_depths(&mut self, program: &Program) {
+        let mut cache = HashMap::new();
+
+        // Entry points are state handlers
+        for state in &program.states {
+            let handler_keys: Vec<String> = [
+                state
+                    .on_enter
+                    .as_ref()
+                    .map(|_| format!("{}::enter", state.name)),
+                state
+                    .on_exit
+                    .as_ref()
+                    .map(|_| format!("{}::exit", state.name)),
+                state
+                    .on_frame
+                    .as_ref()
+                    .map(|_| format!("{}::frame", state.name)),
+            ]
+            .into_iter()
+            .flatten()
+            .collect();
+
+            for key in handler_keys {
+                let mut visited = HashSet::new();
+                let depth = compute_depth(&key, &self.call_graph, &mut visited, &mut cache);
+                self.max_depths.insert(key.clone(), depth);
+
+                if depth > self.stack_depth_limit {
+                    self.diagnostics.push(Diagnostic::error(
+                        ErrorCode::E0401,
+                        format!(
+                            "call depth {depth} in handler '{key}' exceeds stack limit {}",
+                            self.stack_depth_limit
+                        ),
+                        program.span,
+                    ));
+                }
+            }
+        }
+    }
+
     fn check_block(&mut self, block: &Block, state_names: &[&str]) {
         for stmt in &block.statements {
             self.check_statement(stmt, state_names);
@@ -225,9 +354,8 @@ impl Analyzer {
             }
             Statement::Call(name, _args, span) => {
                 if !self.symbols.contains_key(name) {
-                    // Not a known function yet — for M1, just warn
                     self.diagnostics.push(Diagnostic::error(
-                        ErrorCode::E0502,
+                        ErrorCode::E0503,
                         format!("undefined function '{name}'"),
                         *span,
                     ));
@@ -334,6 +462,179 @@ impl Analyzer {
     }
 }
 
+/// Collect all function/call names from a block.
+fn collect_calls(block: &Block) -> Vec<String> {
+    let mut calls = Vec::new();
+    for stmt in &block.statements {
+        collect_calls_stmt(stmt, &mut calls);
+    }
+    calls
+}
+
+fn collect_calls_stmt(stmt: &Statement, calls: &mut Vec<String>) {
+    match stmt {
+        Statement::Call(name, args, _) => {
+            calls.push(name.clone());
+            for arg in args {
+                collect_calls_expr(arg, calls);
+            }
+        }
+        Statement::If(cond, then_b, elifs, else_b, _) => {
+            collect_calls_expr(cond, calls);
+            collect_calls_block(then_b, calls);
+            for (c, b) in elifs {
+                collect_calls_expr(c, calls);
+                collect_calls_block(b, calls);
+            }
+            if let Some(b) = else_b {
+                collect_calls_block(b, calls);
+            }
+        }
+        Statement::While(cond, body, _) => {
+            collect_calls_expr(cond, calls);
+            collect_calls_block(body, calls);
+        }
+        Statement::Loop(body, _) => {
+            collect_calls_block(body, calls);
+        }
+        Statement::Assign(_, _, expr, _) => {
+            collect_calls_expr(expr, calls);
+        }
+        Statement::VarDecl(var) => {
+            if let Some(init) = &var.init {
+                collect_calls_expr(init, calls);
+            }
+        }
+        Statement::Return(Some(expr), _) => {
+            collect_calls_expr(expr, calls);
+        }
+        Statement::Draw(draw) => {
+            collect_calls_expr(&draw.x, calls);
+            collect_calls_expr(&draw.y, calls);
+            if let Some(f) = &draw.frame {
+                collect_calls_expr(f, calls);
+            }
+        }
+        Statement::Return(None, _)
+        | Statement::Transition(_, _)
+        | Statement::WaitFrame(_)
+        | Statement::Break(_)
+        | Statement::Continue(_) => {}
+    }
+}
+
+fn collect_calls_block(block: &Block, calls: &mut Vec<String>) {
+    for stmt in &block.statements {
+        collect_calls_stmt(stmt, calls);
+    }
+}
+
+fn collect_calls_expr(expr: &Expr, calls: &mut Vec<String>) {
+    match expr {
+        Expr::Call(name, args, _) => {
+            calls.push(name.clone());
+            for arg in args {
+                collect_calls_expr(arg, calls);
+            }
+        }
+        Expr::BinaryOp(lhs, _, rhs, _) => {
+            collect_calls_expr(lhs, calls);
+            collect_calls_expr(rhs, calls);
+        }
+        Expr::UnaryOp(_, inner, _) => {
+            collect_calls_expr(inner, calls);
+        }
+        Expr::ArrayIndex(_, idx, _) => {
+            collect_calls_expr(idx, calls);
+        }
+        Expr::ArrayLiteral(elems, _) => {
+            for e in elems {
+                collect_calls_expr(e, calls);
+            }
+        }
+        Expr::IntLiteral(_, _)
+        | Expr::BoolLiteral(_, _)
+        | Expr::Ident(_, _)
+        | Expr::ButtonRead(_, _, _) => {}
+    }
+}
+
+/// Detect cycles in the call graph using DFS. Returns the names of all
+/// functions that participate in a cycle (direct or mutual recursion).
+fn detect_recursion(graph: &HashMap<String, Vec<String>>) -> Vec<String> {
+    let mut recursive = Vec::new();
+    let mut visited = HashSet::new();
+    let mut on_stack = HashSet::new();
+
+    for node in graph.keys() {
+        if !visited.contains(node) {
+            detect_recursion_dfs(node, graph, &mut visited, &mut on_stack, &mut recursive);
+        }
+    }
+
+    recursive.sort();
+    recursive.dedup();
+    recursive
+}
+
+fn detect_recursion_dfs(
+    node: &str,
+    graph: &HashMap<String, Vec<String>>,
+    visited: &mut HashSet<String>,
+    on_stack: &mut HashSet<String>,
+    recursive: &mut Vec<String>,
+) {
+    visited.insert(node.to_string());
+    on_stack.insert(node.to_string());
+
+    if let Some(callees) = graph.get(node) {
+        for callee in callees {
+            if on_stack.contains(callee) {
+                // Found a cycle — mark the callee (the one we recursed back to)
+                recursive.push(callee.clone());
+            } else if !visited.contains(callee) {
+                detect_recursion_dfs(callee, graph, visited, on_stack, recursive);
+            }
+        }
+    }
+
+    on_stack.remove(node);
+}
+
+/// Compute the maximum call depth starting from a given node in the call graph.
+/// Returns `None` if a cycle is encountered (handled separately by recursion detection).
+fn compute_depth(
+    node: &str,
+    graph: &HashMap<String, Vec<String>>,
+    visited: &mut HashSet<String>,
+    cache: &mut HashMap<String, u32>,
+) -> u32 {
+    if let Some(&depth) = cache.get(node) {
+        return depth;
+    }
+    if visited.contains(node) {
+        // Cycle — return 0 to avoid infinite recursion; the cycle itself
+        // is flagged by detect_recursion.
+        return 0;
+    }
+    visited.insert(node.to_string());
+    let mut max_child: u32 = 0;
+    if let Some(callees) = graph.get(node) {
+        for callee in callees {
+            let child = compute_depth(callee, graph, visited, cache);
+            max_child = max_child.max(child);
+        }
+    }
+    visited.remove(node);
+    let depth = if graph.get(node).is_none_or(Vec::is_empty) {
+        0
+    } else {
+        1 + max_child
+    };
+    cache.insert(node.to_string(), depth);
+    depth
+}
+
 fn type_size(t: &NesType) -> u16 {
     match t {
         NesType::U8 | NesType::I8 | NesType::Bool => 1,

src/analyzer/tests.rs

@@ -126,3 +126,100 @@ fn analyze_const_symbol() {
     let sym = result.symbols.get("SPEED").unwrap();
     assert!(sym.is_const);
 }
+
+#[test]
+fn analyze_function_registered() {
+    let result = analyze_ok(
+        r#"
+        game "Test" { mapper: NROM }
+        fun add(a: u8, b: u8) -> u8 { return a }
+        on frame { wait_frame }
+        start Main
+    "#,
+    );
+    assert!(result.symbols.contains_key("add"));
+}
+
+#[test]
+fn analyze_recursion_detected() {
+    let errors = analyze_errors(
+        r#"
+        game "Test" { mapper: NROM }
+        fun a() { a() }
+        on frame { wait_frame }
+        start Main
+    "#,
+    );
+    assert!(errors.contains(&ErrorCode::E0402));
+}
+
+#[test]
+fn analyze_mutual_recursion() {
+    let errors = analyze_errors(
+        r#"
+        game "Test" { mapper: NROM }
+        fun a() { b() }
+        fun b() { a() }
+        on frame { wait_frame }
+        start Main
+    "#,
+    );
+    assert!(errors.contains(&ErrorCode::E0402));
+}
+
+#[test]
+fn analyze_call_depth_ok() {
+    // 3 levels of nesting — well within the default limit of 8
+    let result = analyze_ok(
+        r#"
+        game "Test" { mapper: NROM }
+        fun c() { wait_frame }
+        fun b() { c() }
+        fun a() { b() }
+        on frame { a() }
+        start Main
+    "#,
+    );
+    // The frame handler's depth should be <= 8
+    for &depth in result.max_depths.values() {
+        assert!(depth <= 8, "depth {depth} should be within limit");
+    }
+}
+
+#[test]
+fn analyze_call_depth_exceeded() {
+    // Build a call chain deeper than 8: f1 -> f2 -> ... -> f10
+    let result = analyze_errors(
+        r#"
+        game "Test" { mapper: NROM }
+        fun f10() { wait_frame }
+        fun f9() { f10() }
+        fun f8() { f9() }
+        fun f7() { f8() }
+        fun f6() { f7() }
+        fun f5() { f6() }
+        fun f4() { f5() }
+        fun f3() { f4() }
+        fun f2() { f3() }
+        fun f1() { f2() }
+        on frame { f1() }
+        start Main
+    "#,
+    );
+    assert!(
+        result.contains(&ErrorCode::E0401),
+        "expected E0401 for exceeded call depth, got: {result:?}"
+    );
+}
+
+#[test]
+fn analyze_undefined_function() {
+    let errors = analyze_errors(
+        r#"
+        game "Test" { mapper: NROM }
+        on frame { no_such_fn() }
+        start Main
+    "#,
+    );
+    assert!(errors.contains(&ErrorCode::E0503));
+}

src/errors/diagnostic.rs

@@ -28,9 +28,7 @@ pub enum ErrorCode {
     E0301, // zero-page overflow
 
     // E04xx: Control flow errors
-    #[allow(dead_code)]
     E0401, // call depth exceeded
-    #[allow(dead_code)]
     E0402, // recursion detected
     #[allow(dead_code)]
     E0403, // unreachable state
@@ -39,7 +37,6 @@ pub enum ErrorCode {
     // E05xx: Declaration errors
     E0501, // duplicate declaration
     E0502, // undefined variable
-    #[allow(dead_code)]
     E0503, // undefined function
     E0504, // missing start declaration
     #[allow(dead_code)]

src/ir/lowering.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use super::*;
-use crate::analyzer::{AnalysisResult, VarAllocation};
+use crate::analyzer::AnalysisResult;
 use crate::parser::ast::*;
 
 /// Lower a parsed & analyzed program into IR.
@@ -334,7 +334,9 @@ impl LoweringContext {
                 let idx = self.lower_expr(index);
                 let val = self.lower_expr(expr);
                 // For compound assignment on arrays, load first
-                if op != AssignOp::Assign {
+                if op == AssignOp::Assign {
+                    self.emit(IrOp::ArrayStore(var_id, idx, val));
+                } else {
                     let current = self.fresh_temp();
                     self.emit(IrOp::ArrayLoad(current, var_id, idx));
                     let result = self.fresh_temp();
@@ -348,8 +350,6 @@ impl LoweringContext {
                     };
                     self.emit(ir_op);
                     self.emit(IrOp::ArrayStore(var_id, idx, result));
-                } else {
-                    self.emit(IrOp::ArrayStore(var_id, idx, val));
                 }
             }
         }
@@ -496,9 +496,7 @@ impl LoweringContext {
                 self.emit(IrOp::ArrayLoad(t, var_id, idx));
                 t
             }
-            Expr::BinaryOp(left, op, right, _) => {
+            Expr::BinaryOp(left, op, right, _) => self.lower_binop(left, *op, right),
-                self.lower_binop(left, *op, right)
-            }
             Expr::UnaryOp(op, inner, _) => {
                 let val = self.lower_expr(inner);
                 let t = self.fresh_temp();

src/ir/mod.rs

@@ -133,7 +133,7 @@ pub enum IrOp {
 pub enum IrTerminator {
     /// Unconditional jump to a label.
     Jump(String),
-    /// Conditional branch: if temp != 0 goto true_label else goto false_label.
+    /// Conditional branch: if temp != 0 goto `true_label` else goto `false_label`.
     Branch(IrTemp, String, String),
     /// Return from function, optionally with a value.
     Return(Option<IrTemp>),

src/ir/tests.rs

@@ -42,7 +42,11 @@ fn lower_var_assignment() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     // Should have a StoreVar op
     let has_store = frame_fn
         .blocks
@@ -62,7 +66,11 @@ fn lower_plus_assign() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     let has_add = frame_fn
         .blocks
         .iter()
@@ -83,12 +91,19 @@ fn lower_if_creates_branch() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     let has_branch = frame_fn
         .blocks
         .iter()
         .any(|b| matches!(&b.terminator, IrTerminator::Branch(..)));
-    assert!(has_branch, "if statement should produce a Branch terminator");
+    assert!(
+        has_branch,
+        "if statement should produce a Branch terminator"
+    );
 }
 
 #[test]
@@ -103,7 +118,11 @@ fn lower_while_creates_loop() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     // A while loop needs at least 3 blocks: condition check, body, and exit
     assert!(
         frame_fn.blocks.len() >= 3,
@@ -124,7 +143,11 @@ fn lower_button_read() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     let has_input = frame_fn
         .blocks
         .iter()
@@ -144,7 +167,11 @@ fn lower_draw_sprite() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     let has_draw = frame_fn
         .blocks
         .iter()
@@ -164,7 +191,11 @@ fn lower_constants_become_immediates() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     // SPEED should be lowered to LoadImm(_, 3)
     let has_imm3 = frame_fn
         .blocks
@@ -232,7 +263,11 @@ fn lower_wait_frame() {
         start Main
     "#,
     );
-    let frame_fn = ir.functions.iter().find(|f| f.name.contains("frame")).unwrap();
+    let frame_fn = ir
+        .functions
+        .iter()
+        .find(|f| f.name.contains("frame"))
+        .unwrap();
     let has_wait = frame_fn
         .blocks
         .iter()

src/lib.rs

@@ -2,8 +2,10 @@ pub mod analyzer;
 pub mod asm;
 pub mod codegen;
 pub mod errors;
+pub mod ir;
 pub mod lexer;
 pub mod linker;
+pub mod optimizer;
 pub mod parser;
 pub mod rom;
 pub mod runtime;

src/optimizer/mod.rs

@@ -0,0 +1,301 @@
+#[cfg(test)]
+mod tests;
+
+use std::collections::{HashMap, HashSet};
+
+use crate::ir::{IrBasicBlock, IrFunction, IrOp, IrProgram, IrTemp, IrTerminator};
+
+/// Run all optimization passes on the IR program.
+pub fn optimize(program: &mut IrProgram) {
+    const_fold(program);
+    dead_code(program);
+}
+
+// ---------------------------------------------------------------------------
+// Constant folding
+// ---------------------------------------------------------------------------
+
+/// Single-pass constant folding within each basic block.
+///
+/// When we see `LoadImm(t, v)`, we record `t -> v`. When a binary op or
+/// comparison has both operands as known constants we replace the instruction
+/// with a single `LoadImm`. After folding we remove `LoadImm` ops whose
+/// destination temps are no longer referenced anywhere in the block.
+fn const_fold(program: &mut IrProgram) {
+    for func in &mut program.functions {
+        for block in &mut func.blocks {
+            const_fold_block(block);
+        }
+    }
+}
+
+fn const_fold_block(block: &mut IrBasicBlock) {
+    let mut constants: HashMap<IrTemp, u8> = HashMap::new();
+
+    // First pass: fold arithmetic / comparison ops into LoadImm where possible.
+    for op in &mut block.ops {
+        match *op {
+            IrOp::LoadImm(t, v) => {
+                constants.insert(t, v);
+            }
+            IrOp::Add(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = va.wrapping_add(vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::Sub(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = va.wrapping_sub(vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::And(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = va & vb;
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::Or(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = va | vb;
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::Xor(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = va ^ vb;
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::CmpEq(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = u8::from(va == vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::CmpNe(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = u8::from(va != vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::CmpLt(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = u8::from(va < vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::CmpGt(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = u8::from(va > vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::CmpLtEq(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = u8::from(va <= vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            IrOp::CmpGtEq(dest, a, b) => {
+                if let (Some(&va), Some(&vb)) = (constants.get(&a), constants.get(&b)) {
+                    let result = u8::from(va >= vb);
+                    *op = IrOp::LoadImm(dest, result);
+                    constants.insert(dest, result);
+                }
+            }
+            _ => {}
+        }
+    }
+
+    // Second pass: remove LoadImm ops whose dest temps are no longer referenced
+    // as source operands by anything else in the block (ops + terminator).
+    let used = collect_used_temps_in_block(block);
+    block.ops.retain(|op| {
+        if let IrOp::LoadImm(t, _) = op {
+            used.contains(t)
+        } else {
+            true
+        }
+    });
+}
+
+// ---------------------------------------------------------------------------
+// Dead code elimination
+// ---------------------------------------------------------------------------
+
+/// Remove ops whose destination temps are never used, and remove unreachable
+/// basic blocks (those that have no incoming edges, except the entry block).
+fn dead_code(program: &mut IrProgram) {
+    for func in &mut program.functions {
+        dead_code_eliminate(func);
+        remove_unreachable_blocks(func);
+    }
+}
+
+fn dead_code_eliminate(func: &mut IrFunction) {
+    let used = collect_used_temps(func);
+    for block in &mut func.blocks {
+        block.ops.retain(|op| {
+            if let Some(dest) = op_dest(op) {
+                used.contains(&dest)
+            } else {
+                true // ops without a dest (StoreVar, WaitFrame, etc.) are always kept
+            }
+        });
+    }
+}
+
+/// Collect every `IrTemp` that is used as a *source* operand anywhere in the
+/// function (ops + terminators).
+fn collect_used_temps(func: &IrFunction) -> HashSet<IrTemp> {
+    let mut used = HashSet::new();
+    for block in &func.blocks {
+        collect_used_from_block(block, &mut used);
+    }
+    used
+}
+
+/// Collect temps used as source operands within a single block (ops + terminator).
+fn collect_used_temps_in_block(block: &IrBasicBlock) -> HashSet<IrTemp> {
+    let mut used = HashSet::new();
+    collect_used_from_block(block, &mut used);
+    used
+}
+
+fn collect_used_from_block(block: &IrBasicBlock, used: &mut HashSet<IrTemp>) {
+    for op in &block.ops {
+        collect_source_temps(op, used);
+    }
+    match &block.terminator {
+        IrTerminator::Branch(t, _, _) | IrTerminator::Return(Some(t)) => {
+            used.insert(*t);
+        }
+        IrTerminator::Jump(_) | IrTerminator::Return(None) | IrTerminator::Unreachable => {}
+    }
+}
+
+/// Add all source-operand temps of an op to `used`.
+fn collect_source_temps(op: &IrOp, used: &mut HashSet<IrTemp>) {
+    match op {
+        IrOp::LoadImm(_, _) => {} // dest only, no source temps
+        IrOp::LoadVar(_, _) => {} // dest only
+        IrOp::StoreVar(_, src) => {
+            used.insert(*src);
+        }
+        IrOp::Add(_, a, b)
+        | IrOp::Sub(_, a, b)
+        | IrOp::Mul(_, a, b)
+        | IrOp::And(_, a, b)
+        | IrOp::Or(_, a, b)
+        | IrOp::Xor(_, a, b)
+        | IrOp::CmpEq(_, a, b)
+        | IrOp::CmpNe(_, a, b)
+        | IrOp::CmpLt(_, a, b)
+        | IrOp::CmpGt(_, a, b)
+        | IrOp::CmpLtEq(_, a, b)
+        | IrOp::CmpGtEq(_, a, b) => {
+            used.insert(*a);
+            used.insert(*b);
+        }
+        IrOp::ShiftLeft(_, src, _) | IrOp::ShiftRight(_, src, _) => {
+            used.insert(*src);
+        }
+        IrOp::Negate(_, src) | IrOp::Complement(_, src) => {
+            used.insert(*src);
+        }
+        IrOp::ArrayLoad(_, _, idx) => {
+            used.insert(*idx);
+        }
+        IrOp::ArrayStore(_, idx, val) => {
+            used.insert(*idx);
+            used.insert(*val);
+        }
+        IrOp::Call(_, _, args) => {
+            for a in args {
+                used.insert(*a);
+            }
+        }
+        IrOp::DrawSprite { x, y, frame, .. } => {
+            used.insert(*x);
+            used.insert(*y);
+            if let Some(f) = frame {
+                used.insert(*f);
+            }
+        }
+        IrOp::ReadInput | IrOp::WaitFrame | IrOp::Transition(_) | IrOp::SourceLoc(_) => {}
+    }
+}
+
+/// Return the destination temp of an op, if it has one.
+fn op_dest(op: &IrOp) -> Option<IrTemp> {
+    match op {
+        IrOp::LoadImm(d, _)
+        | IrOp::LoadVar(d, _)
+        | IrOp::Add(d, _, _)
+        | IrOp::Sub(d, _, _)
+        | IrOp::Mul(d, _, _)
+        | IrOp::And(d, _, _)
+        | IrOp::Or(d, _, _)
+        | IrOp::Xor(d, _, _)
+        | IrOp::ShiftLeft(d, _, _)
+        | IrOp::ShiftRight(d, _, _)
+        | IrOp::Negate(d, _)
+        | IrOp::Complement(d, _)
+        | IrOp::CmpEq(d, _, _)
+        | IrOp::CmpNe(d, _, _)
+        | IrOp::CmpLt(d, _, _)
+        | IrOp::CmpGt(d, _, _)
+        | IrOp::CmpLtEq(d, _, _)
+        | IrOp::CmpGtEq(d, _, _)
+        | IrOp::ArrayLoad(d, _, _) => Some(*d),
+        IrOp::Call(dest, _, _) => *dest,
+        IrOp::StoreVar(_, _)
+        | IrOp::ArrayStore(_, _, _)
+        | IrOp::DrawSprite { .. }
+        | IrOp::ReadInput
+        | IrOp::WaitFrame
+        | IrOp::Transition(_)
+        | IrOp::SourceLoc(_) => None,
+    }
+}
+
+/// Remove basic blocks that have no incoming edges (except the entry block,
+/// which is always reachable by definition).
+fn remove_unreachable_blocks(func: &mut IrFunction) {
+    if func.blocks.is_empty() {
+        return;
+    }
+
+    let entry_label = func.blocks[0].label.clone();
+
+    // Collect all labels that are jump/branch targets.
+    let mut reachable_labels: HashSet<String> = HashSet::new();
+    reachable_labels.insert(entry_label);
+
+    for block in &func.blocks {
+        match &block.terminator {
+            IrTerminator::Jump(lbl) => {
+                reachable_labels.insert(lbl.clone());
+            }
+            IrTerminator::Branch(_, t, f) => {
+                reachable_labels.insert(t.clone());
+                reachable_labels.insert(f.clone());
+            }
+            IrTerminator::Return(_) | IrTerminator::Unreachable => {}
+        }
+    }
+
+    func.blocks.retain(|b| reachable_labels.contains(&b.label));
+}

src/optimizer/tests.rs

@@ -0,0 +1,160 @@
+use super::*;
+use crate::ir::{IrBasicBlock, IrFunction, IrOp, IrProgram, IrTemp, IrTerminator, VarId};
+use crate::lexer::Span;
+
+/// Helper: build a minimal `IrProgram` wrapping a single function with one block.
+fn make_program(ops: Vec<IrOp>, terminator: IrTerminator) -> IrProgram {
+    IrProgram {
+        functions: vec![IrFunction {
+            name: "test_fn".to_string(),
+            blocks: vec![IrBasicBlock {
+                label: "entry".to_string(),
+                ops,
+                terminator,
+            }],
+            locals: vec![],
+            param_count: 0,
+            has_return: false,
+            source_span: Span::new(0, 0, 0),
+        }],
+        globals: vec![],
+        rom_data: vec![],
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Constant folding tests
+// ---------------------------------------------------------------------------
+
+#[test]
+fn fold_add_constants() {
+    // LoadImm(t0, 3), LoadImm(t1, 5), Add(t2, t0, t1)
+    // should become LoadImm(t2, 8) with unused t0/t1 removed.
+    let t0 = IrTemp(0);
+    let t1 = IrTemp(1);
+    let t2 = IrTemp(2);
+
+    let ops = vec![
+        IrOp::LoadImm(t0, 3),
+        IrOp::LoadImm(t1, 5),
+        IrOp::Add(t2, t0, t1),
+    ];
+    let mut prog = make_program(ops, IrTerminator::Return(Some(t2)));
+    optimize(&mut prog);
+
+    let block = &prog.functions[0].blocks[0];
+    // After folding + dead-code removal we should have exactly one LoadImm(t2, 8).
+    assert_eq!(block.ops.len(), 1, "expected 1 op, got {:?}", block.ops);
+    assert!(
+        matches!(block.ops[0], IrOp::LoadImm(t, 8) if t == t2),
+        "expected LoadImm(t2, 8), got {:?}",
+        block.ops[0]
+    );
+}
+
+#[test]
+fn fold_sub_constants() {
+    let t0 = IrTemp(0);
+    let t1 = IrTemp(1);
+    let t2 = IrTemp(2);
+
+    let ops = vec![
+        IrOp::LoadImm(t0, 10),
+        IrOp::LoadImm(t1, 3),
+        IrOp::Sub(t2, t0, t1),
+    ];
+    let mut prog = make_program(ops, IrTerminator::Return(Some(t2)));
+    optimize(&mut prog);
+
+    let block = &prog.functions[0].blocks[0];
+    assert_eq!(block.ops.len(), 1, "expected 1 op, got {:?}", block.ops);
+    assert!(
+        matches!(block.ops[0], IrOp::LoadImm(t, 7) if t == t2),
+        "expected LoadImm(t2, 7), got {:?}",
+        block.ops[0]
+    );
+}
+
+#[test]
+fn fold_comparison() {
+    // CmpEq with two equal constants should fold to LoadImm(dest, 1).
+    let t0 = IrTemp(0);
+    let t1 = IrTemp(1);
+    let t2 = IrTemp(2);
+
+    let ops = vec![
+        IrOp::LoadImm(t0, 42),
+        IrOp::LoadImm(t1, 42),
+        IrOp::CmpEq(t2, t0, t1),
+    ];
+    let mut prog = make_program(ops, IrTerminator::Return(Some(t2)));
+    optimize(&mut prog);
+
+    let block = &prog.functions[0].blocks[0];
+    assert_eq!(block.ops.len(), 1, "expected 1 op, got {:?}", block.ops);
+    assert!(
+        matches!(block.ops[0], IrOp::LoadImm(t, 1) if t == t2),
+        "expected LoadImm(t2, 1), got {:?}",
+        block.ops[0]
+    );
+
+    // CmpEq with two different constants should fold to LoadImm(dest, 0).
+    let ops2 = vec![
+        IrOp::LoadImm(t0, 10),
+        IrOp::LoadImm(t1, 20),
+        IrOp::CmpEq(t2, t0, t1),
+    ];
+    let mut prog2 = make_program(ops2, IrTerminator::Return(Some(t2)));
+    optimize(&mut prog2);
+
+    let block2 = &prog2.functions[0].blocks[0];
+    assert_eq!(block2.ops.len(), 1, "expected 1 op, got {:?}", block2.ops);
+    assert!(
+        matches!(block2.ops[0], IrOp::LoadImm(t, 0) if t == t2),
+        "expected LoadImm(t2, 0), got {:?}",
+        block2.ops[0]
+    );
+}
+
+#[test]
+fn dead_code_removes_unused() {
+    // t0 is loaded but never used anywhere — should be eliminated.
+    let t0 = IrTemp(0);
+    let t1 = IrTemp(1);
+
+    let ops = vec![
+        IrOp::LoadImm(t0, 99), // dead — t0 never referenced
+        IrOp::LoadImm(t1, 42), // alive — returned
+    ];
+    let mut prog = make_program(ops, IrTerminator::Return(Some(t1)));
+    optimize(&mut prog);
+
+    let block = &prog.functions[0].blocks[0];
+    assert_eq!(block.ops.len(), 1, "expected 1 op, got {:?}", block.ops);
+    assert!(
+        matches!(block.ops[0], IrOp::LoadImm(t, 42) if t == t1),
+        "expected only LoadImm(t1, 42), got {:?}",
+        block.ops[0]
+    );
+}
+
+#[test]
+fn optimize_preserves_used_ops() {
+    // Every op here is live — nothing should be removed.
+    let t0 = IrTemp(0);
+    let t1 = IrTemp(1);
+    let v0 = VarId(0);
+
+    let ops = vec![
+        IrOp::LoadImm(t0, 5),
+        IrOp::LoadVar(t1, v0),
+        IrOp::StoreVar(v0, t0),
+    ];
+    let mut prog = make_program(ops.clone(), IrTerminator::Return(Some(t1)));
+    optimize(&mut prog);
+
+    let block = &prog.functions[0].blocks[0];
+    // LoadImm(t0, 5) is used by StoreVar; LoadVar(t1, v0) is used by Return.
+    // StoreVar has no dest so it's always kept.
+    assert_eq!(block.ops.len(), 3, "expected 3 ops, got {:?}", block.ops);
+}