analyzer+ir: automatically overlay state-local variables

Before this change, state-local variables (`state Foo { var x: u8 = 0 }`)
were silently no-ops: the analyzer allocated a ZP slot for them, but
the codegen's `var_addrs` map only covered IR globals and scope-qualified
function locals — so every `LoadVar` / `StoreVar` whose `VarId` pointed
at a state-local resolved to no address and emitted nothing. Existing
examples compiled and matched their goldens because none of them observed
the dropped writes within the 180-frame harness window.

The overlay changes the analyzer's state-local pass to snapshot both the
ZP and RAM cursors after the globals have been laid out, then rewind to
that snapshot before each state's locals and track the running max.
`ZP_CURRENT_STATE` keeps exactly one state active at runtime, so every
state's locals are mutually exclusive with every other state's and can
share the same bytes. The IR lowerer now pushes each state's locals into
the IR globals table (with `init_value=None`) so the codegen resolves
their addresses the same way it does program globals, and prepends the
declared initializers to each state's `on_enter` handler (synthesizing
an empty one where needed) so a freshly-entered state re-establishes its
bytes before user code runs.

`--memory-map` now tags each allocation with its owning state
(`[@Title]`, `[@Playing]`, ...) and counts distinct bytes rather than
summed allocation sizes so overlaid slots don't double-count. The
`AnalysisResult.state_local_owners` map exposes the ownership to any
tool that wants to group allocations the same way.

Only `state_machine.ne` and `platformer.ne` declare state-level vars,
so they're the only example ROMs whose bytes change. `platformer.ne`'s
audio golden shifts slightly (the now-working `blink` counter in Title
adds a few cycles per frame before the auto-transition to Playing, which
offsets APU register writes within each frame); its video golden and
every other example ROM stay byte-for-byte identical.

Fixes #22.

https://claude.ai/code/session_015kvJu3iEFLSRJoShPBfm3X

src/analyzer/mod.rs

@@ -31,6 +31,10 @@ pub struct AnalysisResult {
     pub diagnostics: Vec<Diagnostic>,
     pub call_graph: HashMap<String, Vec<String>>,
     pub max_depths: HashMap<String, u32>,
+    /// For each state-local variable name, the state it belongs to.
+    /// Consumed by the memory-map printer to group overlaid slots by
+    /// their owning state. Empty for programs without state-locals.
+    pub state_local_owners: HashMap<String, String>,
 }
 
 /// Default call stack depth limit for the NES runtime.
@@ -124,12 +128,20 @@ pub fn analyze(program: &Program) -> AnalysisResult {
     };
     analyzer.analyze_program(program);
 
+    let mut state_local_owners = HashMap::new();
+    for state in &program.states {
+        for var in &state.locals {
+            state_local_owners.insert(var.name.clone(), state.name.clone());
+        }
+    }
+
     AnalysisResult {
         symbols: analyzer.symbols,
         var_allocations: analyzer.var_allocations,
         diagnostics: analyzer.diagnostics,
         call_graph: analyzer.call_graph,
         max_depths: analyzer.max_depths,
+        state_local_owners,
     }
 }
 
@@ -524,12 +536,46 @@ impl Analyzer {
             self.register_fun(fun);
         }
 
-        // Register state-local variables
+        // Register state-local variables with automatic memory
+        // overlaying. At runtime only one state is active at a time
+        // (a single `ZP_CURRENT_STATE` byte picks the handler), so
+        // every state's locals are mutually exclusive with every
+        // other state's — their RAM footprints can share the same
+        // addresses. The allocator snapshots both cursors after the
+        // globals have been laid out, then rewinds to that snapshot
+        // before each state's locals and tracks the running max.
+        // The overall cursor advances to the max at the end, so
+        // anything allocated after the state-locals (function
+        // parameters, function bodies' locals) picks up past every
+        // state's overlay window.
+        //
+        // Each state's on_enter handler re-initializes the locals
+        // from their declared initializers — the IR lowering moves
+        // those stores into the handler's prologue so a freshly
+        // entered state doesn't read another state's leftover
+        // bytes. State-locals whose name collides with a global or
+        // another state's local are still rejected via E0501 at
+        // `register_var` because the symbol table is keyed by the
+        // bare name.
+        let overlay_zp_base = self.next_zp_addr;
+        let overlay_ram_base = self.next_ram_addr;
+        let mut max_zp = overlay_zp_base;
+        let mut max_ram = overlay_ram_base;
         for state in &program.states {
+            self.next_zp_addr = overlay_zp_base;
+            self.next_ram_addr = overlay_ram_base;
             for var in &state.locals {
                 self.register_var(var);
             }
+            if self.next_zp_addr > max_zp {
+                max_zp = self.next_zp_addr;
+            }
+            if self.next_ram_addr > max_ram {
+                max_ram = self.next_ram_addr;
+            }
         }
+        self.next_zp_addr = max_zp;
+        self.next_ram_addr = max_ram;
 
         // Validate state references
         let state_names: Vec<&str> = program.states.iter().map(|s| s.name.as_str()).collect();

src/ir/lowering.rs

@@ -655,6 +655,35 @@ impl LoweringContext {
         // enforced.
         self.capture_inline_bodies(program);
 
+        // Register state-local variables as IR globals so the codegen
+        // resolves their addresses through the same `ir.globals`
+        // pathway it uses for program globals — the analyzer records
+        // them under their bare names in `var_allocations`, which
+        // `IrCodeGen::new` then matches against each global's
+        // `name` field. Without this, a `LoadVar`/`StoreVar` on a
+        // state-local variable resolved its `VarId` to no address
+        // and the codegen silently emitted nothing — the root
+        // cause of the "state-local variables don't actually work"
+        // bug that this change ships with the overlay feature.
+        //
+        // `init_value` / `init_array` are intentionally left blank:
+        // state-locals are re-initialized in each state's on_enter
+        // handler below, not at program reset. The analyzer's
+        // overlay allocation means one state's initial bytes would
+        // stomp on another state's if we emitted them at reset.
+        for state in &program.states {
+            for var in &state.locals {
+                let var_id = self.get_or_create_var(&var.name);
+                self.globals.push(IrGlobal {
+                    var_id,
+                    name: var.name.clone(),
+                    size: type_size(&var.var_type),
+                    init_value: None,
+                    init_array: Vec::new(),
+                });
+            }
+        }
+
         // Lower user functions
         for fun in &program.functions {
             self.lower_function(fun);
@@ -737,7 +766,26 @@ impl LoweringContext {
         // `Title::on frame` and one in `Playing::on frame` get
         // different VarIds.
 
-        if let Some(on_enter) = &state.on_enter {
+        // State-local variables with initializers need their values
+        // re-established every time the state is entered, because
+        // the analyzer overlays state-locals across mutually
+        // exclusive states and another state's writes can clobber
+        // the bytes in between. If the state already has an
+        // on_enter handler, `lower_handler` prepends the
+        // initializer stores; if not, synthesize an empty one here
+        // so the dispatch path still calls into the prelude.
+        let needs_synthetic_enter =
+            state.on_enter.is_none() && state.locals.iter().any(|v| v.init.is_some());
+        let synthetic_enter = Block {
+            statements: Vec::new(),
+            span: state.span,
+        };
+        let on_enter_block: Option<&Block> = state.on_enter.as_ref().or(if needs_synthetic_enter {
+            Some(&synthetic_enter)
+        } else {
+            None
+        });
+        if let Some(on_enter) = on_enter_block {
             self.lower_handler(
                 &format!("{}_enter", state.name),
                 &format!("{}__enter", state.name),
@@ -813,6 +861,53 @@ impl LoweringContext {
 
         let entry = self.fresh_label(&format!("{name}_entry"));
         self.start_block(&entry);
+
+        // on_enter handlers carry the state-local initializer
+        // prologue: every `var x: u8 = expr` declared at
+        // `state Foo { ... }` level gets a store emitted at the
+        // top of on_enter so the state's locals are reset every
+        // time the state is entered. This is what makes the
+        // analyzer's overlay allocation safe — another state
+        // having written into these bytes no longer matters,
+        // because we unconditionally re-initialize them here.
+        // User code inside the on_enter body then runs on top.
+        // Locals without an initializer are left at whatever
+        // bytes the previous state wrote; the programmer can
+        // explicitly assign them if they want a fresh value.
+        if name.ends_with("_enter") {
+            for var in &state.locals {
+                let Some(init) = &var.init else { continue };
+                let var_id = self.get_or_create_var(&var.name);
+                if let Expr::ArrayLiteral(_, _) = init {
+                    // Array initializers for state-locals aren't
+                    // supported yet — a runtime memcpy loop from a
+                    // ROM blob would be the natural lowering.
+                    // Programs that try this should get a diagnostic
+                    // from the analyzer; for now, silently skip.
+                    continue;
+                }
+                if let Expr::StructLiteral(_, fields, _) = init {
+                    for (fname, fexpr) in fields {
+                        let full = format!("{}.{fname}", var.name);
+                        let fvid = self.get_or_create_var(&full);
+                        let val = self.lower_expr(fexpr);
+                        self.emit(IrOp::StoreVar(fvid, val));
+                    }
+                    continue;
+                }
+                let val = self.lower_expr(init);
+                self.emit(IrOp::StoreVar(var_id, val));
+                // u16-typed state-locals also need the high byte
+                // of the initializer stored at base+1. Mirror the
+                // `VarDecl` lowering in `lower_statement` so wide
+                // inits round-trip cleanly.
+                if matches!(var.var_type, NesType::U16) {
+                    let (_, hi) = self.widen(val);
+                    self.emit(IrOp::StoreVarHi(var_id, hi));
+                }
+            }
+        }
+
         self.lower_block(block);
         self.end_block(IrTerminator::Return(None));
 

src/main.rs

@@ -155,7 +155,24 @@ fn write_memory_map(
     backgrounds: &[BackgroundData],
 ) -> std::io::Result<()> {
     let mut allocs: Vec<_> = analysis.var_allocations.iter().collect();
-    allocs.sort_by_key(|a| a.address);
+    // Sort by address, then by state-local owner (None before Some),
+    // so the memory map groups overlaid state-locals together under
+    // their shared base address.
+    allocs.sort_by(|a, b| {
+        a.address.cmp(&b.address).then_with(|| {
+            analysis
+                .state_local_owners
+                .get(&a.name)
+                .cmp(&analysis.state_local_owners.get(&b.name))
+        })
+    });
+
+    let fmt_tag = |name: &str| -> String {
+        match analysis.state_local_owners.get(name) {
+            Some(state) => format!("[@{state}]"),
+            None => "[USER]   ".to_string(),
+        }
+    };
 
     writeln!(w, "=== NEScript Memory Map ===")?;
     writeln!(w, "Zero Page ($00-$FF):")?;
@@ -164,14 +181,16 @@ fn write_memory_map(
         "  $00-$0F  [SYSTEM]  reserved (frame flag, input, state, params, scratch)"
     )?;
     for a in allocs.iter().filter(|a| a.address < 0x100) {
+        let tag = fmt_tag(&a.name);
         if a.size == 1 {
-            writeln!(w, "  ${:04X}    [USER]    {} (u8)", a.address, a.name)?;
+            writeln!(w, "  ${:04X}    {} {} (u8)", a.address, tag, a.name)?;
         } else {
             writeln!(
                 w,
-                "  ${:04X}-${:04X}  [USER]  {} ({} bytes)",
+                "  ${:04X}-${:04X}  {} {} ({} bytes)",
                 a.address,
                 a.address + a.size - 1,
+                tag,
                 a.name,
                 a.size
             )?;
@@ -183,14 +202,16 @@ fn write_memory_map(
         writeln!(w, "\nRAM ($0200-$07FF):")?;
         writeln!(w, "  $0200-$02FF  [SYSTEM]  OAM shadow buffer")?;
         for a in &ram_allocs {
+            let tag = fmt_tag(&a.name);
             if a.size == 1 {
-                writeln!(w, "  ${:04X}        [USER]    {} (u8)", a.address, a.name)?;
+                writeln!(w, "  ${:04X}        {} {} (u8)", a.address, tag, a.name)?;
             } else {
                 writeln!(
                     w,
-                    "  ${:04X}-${:04X}  [USER]  {} ({} bytes)",
+                    "  ${:04X}-${:04X}  {} {} ({} bytes)",
                     a.address,
                     a.address + a.size - 1,
+                    tag,
                     a.name,
                     a.size
                 )?;
@@ -198,17 +219,24 @@ fn write_memory_map(
         }
     }
 
-    // Summary line.
-    let zp_used: u16 = allocs
-        .iter()
-        .filter(|a| a.address < 0x80)
-        .map(|a| a.size)
-        .sum();
-    let ram_used: u16 = allocs
-        .iter()
-        .filter(|a| a.address >= 0x300)
-        .map(|a| a.size)
-        .sum();
+    // Summary counts distinct byte addresses in use, not the sum of
+    // allocation sizes, so overlaid state-locals are only counted
+    // once per shared byte. Non-state-local allocations and the
+    // per-state allocations each contribute their own bytes.
+    let mut zp_bytes_used: std::collections::HashSet<u16> = std::collections::HashSet::new();
+    let mut ram_bytes_used: std::collections::HashSet<u16> = std::collections::HashSet::new();
+    for a in &allocs {
+        for offset in 0..a.size {
+            let byte = a.address + offset;
+            if byte < 0x80 {
+                zp_bytes_used.insert(byte);
+            } else if byte >= 0x300 {
+                ram_bytes_used.insert(byte);
+            }
+        }
+    }
+    let zp_used = zp_bytes_used.len();
+    let ram_used = ram_bytes_used.len();
     writeln!(w)?;
     writeln!(w, "Zero Page: {zp_used}/128 bytes used")?;
     writeln!(w, "Main RAM:  {ram_used}/1280 bytes used")?;
@@ -513,6 +541,7 @@ mod tests {
             diagnostics: Vec::new(),
             call_graph: HashMap::new(),
             max_depths: HashMap::new(),
+            state_local_owners: HashMap::new(),
         }
     }