Codegen: on_scanline per-state dispatch + NMI reload

Extends the \`on_scanline\` codegen to support multiple scanline
handlers across states:

- \`__irq_user\` now dispatches by \`current_state\`: each state with a
  scanline handler gets a CMP/BNE/JSR entry in the dispatch table.
  States without a handler fall through to just acknowledge the IRQ.
- New \`__ir_mmc3_reload\` helper that (re)loads the MMC3 counter
  latch based on \`current_state\`. States without a scanline handler
  fall through to disable the IRQ (\$E000 write).
- Linker detects the \`__ir_mmc3_reload\` label in user code and
  splices a JSR into it at the top of the NMI handler, so the
  counter is reloaded once per frame with the current state's
  target scanline.
- IRQ handler no longer re-enables IRQ on ACK (the NMI reload now
  handles that) so it won't fire multiple times per frame.
- Program init chooses the start state's scanline count (if any) or
  the first scanline handler found as a fallback.

Also fixes \`dump_asm\`: a \`NOP\` with a \`Label\` operand is a label
definition, but any other opcode with a \`Label\` operand is a real
instruction like \`JSR foo\`. The old dump was hiding JSR/JMP targets.

https://claude.ai/code/session_01W6eQFStA66EuMKHUFo2rx3

src/codegen/ir_codegen.rs

@@ -186,11 +186,22 @@ impl<'a> IrCodeGen<'a> {
         }
 
         // 2b. If the program has any `on scanline` handlers, set up
-        // the MMC3 IRQ counter. We pick the first scanline handler's
+        // the MMC3 IRQ counter using the scanline count of the
-        // line count as the latch value. Only supports a single-
+        // *start* state (if it has one). Per-state scanline switching
-        // scanline-per-program setup for now.
+        // is then handled by writing to `$C000`/`$C001`/`$E001` at
-        let scanline_info = find_first_scanline_handler(ir);
+        // the top of each on_enter handler via inline asm (future).
-        if let Some((_state_name, line)) = &scanline_info {
+        // For now we always set up with the first scanline found in
+        // the start state, or any scanline in the program as a
+        // fallback.
+        let scanline_handlers = collect_scanline_handlers(ir);
+        if !scanline_handlers.is_empty() {
+            // Prefer a scanline in the start state; otherwise use the
+            // first one found.
+            let (_state_name, line) = scanline_handlers
+                .iter()
+                .find(|(s, _)| *s == ir.start_state)
+                .or_else(|| scanline_handlers.first())
+                .unwrap();
             // Write (line-1) to $C000 (scanline latch), any value to
             // $C001 (reload counter), any value to $E001 (enable IRQ).
             self.emit(LDA, AM::Immediate(line.saturating_sub(1)));
@@ -235,11 +246,11 @@ impl<'a> IrCodeGen<'a> {
         }
 
         // 5. If we have scanline handlers, emit an IRQ handler that
-        // saves registers, ACKs the MMC3 IRQ, JSRs into the first
+        // saves registers, ACKs the MMC3 IRQ, dispatches to the
-        // scanline handler, restores registers, and RTIs. The linker
+        // current state's scanline handler (if any), restores
-        // picks up `__irq_user` and uses it as the IRQ vector instead
+        // registers, and RTIs. The linker picks up `__irq_user` and
-        // of the default stub.
+        // uses it as the IRQ vector instead of the default stub.
-        if let Some((state_name, line)) = scanline_info {
+        if !scanline_handlers.is_empty() {
             self.emit_label("__irq_user");
             // Save registers onto the stack.
             self.emit(PHA, AM::Implied);
@@ -247,14 +258,25 @@ impl<'a> IrCodeGen<'a> {
             self.emit(PHA, AM::Implied);
             self.emit(TYA, AM::Implied);
             self.emit(PHA, AM::Implied);
-            // Acknowledge and reload the MMC3 IRQ counter.
+            // Acknowledge the MMC3 IRQ ($E000 = disable/ack). NMI
-            self.emit(STA, AM::Absolute(0xE000)); // disable / ack
+            // will re-enable each frame by reloading the counter.
-            self.emit(STA, AM::Absolute(0xE001)); // re-enable
+            self.emit(LDA, AM::Immediate(0));
-                                                  // JSR into the scanline handler. We don't dispatch by
+            self.emit(STA, AM::Absolute(0xE000));
-                                                  // state yet — the first scanline handler found is used
+            // Dispatch by current_state: for each state with a
-                                                  // unconditionally.
+            // scanline handler, CMP and JSR.
-            let handler = format!("{state_name}_scanline_{line}");
+            self.emit(LDA, AM::ZeroPage(ZP_CURRENT_STATE));
-            self.emit(JSR, AM::Label(format!("__ir_fn_{handler}")));
+            let done_label = "__irq_user_done".to_string();
+            for (i, (state_name, line)) in scanline_handlers.iter().enumerate() {
+                let state_idx = self.state_indices.get(state_name).copied().unwrap_or(0);
+                let skip_label = format!("__irq_disp_skip_{i}");
+                self.emit(CMP, AM::Immediate(state_idx));
+                self.emit(BNE, AM::LabelRelative(skip_label.clone()));
+                let handler = format!("{state_name}_scanline_{line}");
+                self.emit(JSR, AM::Label(format!("__ir_fn_{handler}")));
+                self.emit(JMP, AM::Label(done_label.clone()));
+                self.emit_label(&skip_label);
+            }
+            self.emit_label(&done_label);
             // Restore registers and return from interrupt.
             self.emit(PLA, AM::Implied);
             self.emit(TAY, AM::Implied);
@@ -262,6 +284,34 @@ impl<'a> IrCodeGen<'a> {
             self.emit(TAX, AM::Implied);
             self.emit(PLA, AM::Implied);
             self.emit(RTI, AM::Implied);
+
+            // Also emit a helper that the NMI handler can use to
+            // reload the MMC3 counter each frame. Linker extends the
+            // NMI with a JSR into this when `__ir_mmc3_reload` exists.
+            self.emit_label("__ir_mmc3_reload");
+            // Dispatch on current_state to pick the right scanline
+            // count, write (count-1) to $C000, reload $C001, enable
+            // $E001. States without a scanline handler fall through
+            // to disable ($E000).
+            self.emit(LDA, AM::ZeroPage(ZP_CURRENT_STATE));
+            let reload_done = "__ir_mmc3_reload_done".to_string();
+            for (i, (state_name, line)) in scanline_handlers.iter().enumerate() {
+                let state_idx = self.state_indices.get(state_name).copied().unwrap_or(0);
+                let skip_label = format!("__ir_reload_skip_{i}");
+                self.emit(CMP, AM::Immediate(state_idx));
+                self.emit(BNE, AM::LabelRelative(skip_label.clone()));
+                self.emit(LDA, AM::Immediate(line.saturating_sub(1)));
+                self.emit(STA, AM::Absolute(0xC000));
+                self.emit(STA, AM::Absolute(0xC001));
+                self.emit(STA, AM::Absolute(0xE001));
+                self.emit(JMP, AM::Label(reload_done.clone()));
+                self.emit_label(&skip_label);
+            }
+            // No matching state — disable IRQ for this frame.
+            self.emit(LDA, AM::Immediate(0));
+            self.emit(STA, AM::Absolute(0xE000));
+            self.emit_label(&reload_done);
+            self.emit(RTS, AM::Implied);
         }
 
         self.instructions
@@ -629,20 +679,21 @@ enum CmpKind {
     GtEq,
 }
 
-/// Scan the IR functions for the first scanline handler (named
+/// Scan the IR functions for all scanline handlers (named
-/// `{state}_scanline_{line}`) and return the state name and line.
+/// `{state}_scanline_{line}`) and return them in IR function order.
-fn find_first_scanline_handler(ir: &IrProgram) -> Option<(String, u8)> {
+fn collect_scanline_handlers(ir: &IrProgram) -> Vec<(String, u8)> {
+    let mut out = Vec::new();
     for func in &ir.functions {
         // Match the pattern `<state>_scanline_<N>` by splitting on
         // the last `_scanline_`. This keeps it simple and avoids
         // re-threading scanline info through lowering.
         if let Some((state_name, tail)) = func.name.rsplit_once("_scanline_") {
             if let Ok(line) = tail.parse::<u8>() {
-                return Some((state_name.to_string(), line));
+                out.push((state_name.to_string(), line));
             }
         }
     }
-    None
+    out
 }
 
 #[cfg(test)]

src/linker/mod.rs

@@ -22,6 +22,15 @@ pub struct SpriteData {
     pub chr_bytes: Vec<u8>,
 }
 
+/// True if `instructions` contains a label definition with the given
+/// name. Labels are emitted as `NOP` pseudo-instructions whose mode
+/// is `AddressingMode::Label(name)`.
+fn has_label(instructions: &[Instruction], name: &str) -> bool {
+    instructions
+        .iter()
+        .any(|i| matches!(&i.mode, AM::Label(n) if n == name))
+}
+
 /// A smiley face CHR tile for the default sprite (M1).
 const DEFAULT_SPRITE_CHR: [u8; 16] = [
     // Plane 0 (low bits)
@@ -109,6 +118,17 @@ impl Linker {
 
         // NMI handler
         all_instructions.push(Instruction::new(NOP, AM::Label("__nmi".into())));
+        // If user code emits an MMC3 reload hook, splice in a JSR
+        // before the regular NMI runs. This reloads the scanline IRQ
+        // counter each frame so the handler fires at the right line.
+        // The presence of the `__ir_mmc3_reload` label is detected
+        // during assembly via the labels map; we unconditionally
+        // emit a conditional JSR whose target is resolved at link
+        // time. The helper emits an RTS so it's safe to call even
+        // when there's no work to do.
+        if has_label(user_code, "__ir_mmc3_reload") {
+            all_instructions.push(Instruction::new(JSR, AM::Label("__ir_mmc3_reload".into())));
+        }
         all_instructions.extend(runtime::gen_nmi());
 
         // IRQ handler

src/main.rs

@@ -90,10 +90,18 @@ fn main() {
 }
 
 fn dump_asm(instructions: &[nescript::asm::Instruction]) {
+    use nescript::asm::{AddressingMode, Opcode};
     for inst in instructions {
-        if let nescript::asm::AddressingMode::Label(name) = &inst.mode {
+        // A bare `NOP` with a `Label` operand is a label *definition*
-            println!("{name}:");
+        // (the pseudo-instruction the codegen emits when marking a
-            continue;
+        // position). Any other opcode with `Label` mode is an actual
+        // instruction like `JSR foo` or `JMP bar`, so we show the
+        // opcode + target.
+        if inst.opcode == Opcode::NOP {
+            if let AddressingMode::Label(name) = &inst.mode {
+                println!("{name}:");
+                continue;
+            }
         }
         println!("    {:?} {:?}", inst.opcode, inst.mode);
     }

tests/integration_test.rs

@@ -156,6 +156,28 @@ fn program_with_on_scanline_mmc3() {
     rom::validate_ines(&rom_data).expect("should be valid iNES");
 }
 
+#[test]
+fn program_with_on_scanline_per_state() {
+    // Two states, each with its own scanline handler at a different
+    // position. The IR codegen should emit per-state dispatch in
+    // both `__irq_user` and `__ir_mmc3_reload`.
+    let source = r#"
+        game "MultiSL" { mapper: MMC3 }
+        var s: u8 = 0
+        state A {
+            on frame { wait_frame }
+            on scanline(64) { scroll(0, 0) }
+        }
+        state B {
+            on frame { wait_frame }
+            on scanline(192) { scroll(0, 0) }
+        }
+        start A
+    "#;
+    let rom_data = compile(source);
+    rom::validate_ines(&rom_data).expect("should be valid iNES");
+}
+
 #[test]
 fn program_with_structs() {
     let source = r#"