Fix m68k frame teardown for intra-subroutine dispatch

codeGenM68k/src/prog8/codegen/m68k/AsmGen.kt

@@ -293,6 +293,7 @@ internal class AsmGen(val program: IRProgram, internal val target: ICompilationT
     // the code chunks of the subroutine currently being emitted, to tell a jump that stays
     // inside it apart from one that leaves it
     private var currentSubChunks: Set<IRCodeChunkBase> = emptySet()
+    private var previousInsn: IRInstruction? = null
 
     // Virtual registers of the subroutine being emitted, by register number -> frame offset.
     // A re-entrant subroutine (one that can have two live activations) keeps them per-activation in
@@ -317,6 +318,21 @@ internal class AsmGen(val program: IRProgram, internal val target: ICompilationT
         return target !in currentSubChunks
     }
 
+    internal fun jumpTableStaysInCurrentSub(pointer: VirtualRegister): Boolean {
+        val load = previousInsn ?: return false
+        if (load.opcode != Opcode.LOADX) return false
+        if (load.dest?.register != pointer) return false
+        val indexed = load.memory as? MemoryReference.Indexed ?: return false
+        val symbol = indexed.base as? AddressBase.Symbol ?: return false
+        val variable = program.st.lookup(symbol.name) as? IRStStaticVariable ?: return false
+        val initializer = variable.initializationValue as? IRVariableInitializer.Array ?: return false
+        if (initializer.elements.isEmpty()) return false
+        return initializer.elements.all { element ->
+            val label = (element as? IRStSymbolicReference.Symbol)?.name ?: return@all false
+            !labelIsOutsideCurrentSub(label)
+        }
+    }
+
     // === virtual register file layout (1, 2 or 4 bytes per slot depending on type, word-aligned) ===
     private data class RegFileLayout(val offsets: Map<Int, Int>, val totalSize: Int)
 
@@ -915,6 +931,7 @@ internal class AsmGen(val program: IRProgram, internal val target: ICompilationT
         deadStoreSuppressionAllowed: Boolean = true,
         preserveLoopCounter: Boolean = false
     ) {
+        previousInsn = null
         invalidateD0Cache()                 // new basic block / branch target
         emitSourceComment(chunk.sourceLinesPositions)
         val callOptimizations = mutableMapOf<Int, ImmediateCallOptimization>()
@@ -959,6 +976,7 @@ internal class AsmGen(val program: IRProgram, internal val target: ICompilationT
                 val mem = insn.memory as MemoryReference.Direct
                 val slot = mem.base as AddressBase.FrameSlot
                 emitFrameZeroLoop(slot.offset + mem.displacement, zeroRun)
+                previousInsn = null
                 index += zeroRun
                 continue
             }
@@ -970,6 +988,7 @@ internal class AsmGen(val program: IRProgram, internal val target: ICompilationT
                 callOptimizations[index],
                 index in deadLoadIndices
             )
+            previousInsn = insn
             if (clobbersLoopCounter)
                 emitLine("move.w  (sp)+,d7", "restore loop counter after call")
             index++

codeGenM68k/src/prog8/codegen/m68k/AsmOptimizer.kt

@@ -145,6 +145,7 @@ private fun optimizeRedundantLea(allPretrimmed: List<String>): List<Modification
 
 /** Matches a parenthesised group, for stripping addressing modes. */
 private val parenthesised = Regex("""\([^)]*\)""")
+private val d0IndexOperand = Regex("""\(\s*a0\s*,\s*d0\.w(?:\*\d+)?\s*\)""")
 
 /** A single edit to apply: either remove a line (keeping its label if present) or replace it. */
 private class Modification(val lineIndex: Int, val remove: Boolean, val replacement: String?)
@@ -671,11 +672,19 @@ private fun isSlotReferencedInLine(slot: String, line: String): Boolean {
  * or for 68020+ JUMPI/CALLI:
  *   jmp/jsr ([p8_regfile+M])     ; memory-indirect
  *
- * Fused (preserves current unscaled semantics; no *4 scaling):
+ * The same sequence with the scale folded into the addressing mode instead of
+ * into d0 (68020+; this is what a pointer-array dispatch such as on..goto emits):
+ *   lea     table_label, a0
+ *   move.l  (a0,d0.w*4), d0
+ *   move.l  d0, p8_regfile+M
+ *   movea.l p8_regfile+M, a0
+ *   jmp/jsr (a0)
+ *
+ * Fused, keeping whichever index operand was matched so the scale is preserved:
  *   moveq   #0, d0
  *   move.b  p8_regfile+N, d0
  *   lea     table_label, a0
- *   move.l  (a0,d0.w), a0
+ *   move.l  (a0,d0.w), a0        ; or (a0,d0.w*4), a0 on 68020+
  *   jmp/jsr (a0)
  */
 private fun optimizeFuseLoadxJumpi(allPretrimmed: List<String>, lines: MutableList<String>): List<Modification> {
@@ -686,11 +695,13 @@ private fun optimizeFuseLoadxJumpi(allPretrimmed: List<String>, lines: MutableLi
     val mods = mutableListOf<Modification>()
     val alreadyModified = mutableSetOf<Int>()
 
-    fun isLoadIndexed(instr: String): Boolean {
-        // matches "move.l  (a0,d0.w), d0" with flexible spacing
-        if (!instr.startsWith("move.l")) return false
-        if ("(a0,d0.w)" !in instr) return false
-        return instr.trim().endsWith(", d0") || instr.trim().endsWith(",d0")
+    fun loadIndexedOperand(instr: String): String? {
+        // matches "move.l  (a0,d0.w), d0" (68000, index pre-scaled into d0) and the scaled
+        // "move.l  (a0,d0.w*4), d0" forms (68020+, scale folded into the addressing mode).
+        // The matched operand is returned so the fused instruction keeps the same scale.
+        if (!instr.startsWith("move.l")) return null
+        if (!(instr.trim().endsWith(", d0") || instr.trim().endsWith(",d0"))) return null
+        return d0IndexOperand.find(instr)?.value
     }
 
     fun isSpill(instr: String): String? {
@@ -719,7 +730,7 @@ private fun optimizeFuseLoadxJumpi(allPretrimmed: List<String>, lines: MutableLi
         val loadIdx = leaIdx + 1
         val spillIdx = leaIdx + 2
         if (filtered[loadIdx].index in alreadyModified || filtered[spillIdx].index in alreadyModified) continue
-        if (!isLoadIndexed(filtered[loadIdx].instruction)) continue
+        val loadOperand = loadIndexedOperand(filtered[loadIdx].instruction) ?: continue
         val slotM = isSpill(filtered[spillIdx].instruction) ?: continue
 
         // Check jump/call variant: 68000 uses movea reload + jmp/jsr (a0),
@@ -746,7 +757,7 @@ private fun optimizeFuseLoadxJumpi(allPretrimmed: List<String>, lines: MutableLi
             if (hasLabel(filtered[spillIdx].trimmed) || hasLabel(filtered[nextIdx].trimmed) || hasLabel(filtered[jmpIdx].trimmed)) continue
 
             val loadIndent = lines[filtered[loadIdx].index].takeWhile { it.isWhitespace() }
-            mods.add(Modification(filtered[loadIdx].index, false, "${loadIndent}move.l  (a0,d0.w), a0"))
+            mods.add(Modification(filtered[loadIdx].index, false, "${loadIndent}move.l  $loadOperand, a0"))
             mods.add(Modification(filtered[spillIdx].index, true, null))
             mods.add(Modification(filtered[nextIdx].index, true, null))
             alreadyModified.add(filtered[loadIdx].index)
@@ -759,7 +770,7 @@ private fun optimizeFuseLoadxJumpi(allPretrimmed: List<String>, lines: MutableLi
             val loadIndent = lines[filtered[loadIdx].index].takeWhile { it.isWhitespace() }
             val jmpIndent = lines[filtered[nextIdx].index].takeWhile { it.isWhitespace() }
             val mnemonic = if (is68020MemIndirectJsr) "jsr" else "jmp"
-            mods.add(Modification(filtered[loadIdx].index, false, "${loadIndent}move.l  (a0,d0.w), a0"))
+            mods.add(Modification(filtered[loadIdx].index, false, "${loadIndent}move.l  $loadOperand, a0"))
             mods.add(Modification(filtered[spillIdx].index, true, null))
             mods.add(Modification(filtered[nextIdx].index, false, "${jmpIndent}$mnemonic  (a0)"))
             alreadyModified.add(filtered[loadIdx].index)

codeGenM68k/src/prog8/codegen/m68k/InstrControl.kt

@@ -20,7 +20,6 @@ import prog8.code.core.CpuType
 import prog8.code.core.toHex
 import prog8.code.target.Amiga500Target
 import prog8.intermediate.*
-import prog8.intermediate.incomingSlotSize
 
 internal fun AsmGen.translateControl(insn: IRInstruction, forwardedImmediateCall: ImmediateCallOptimization? = null) {
     when (insn.opcode) {
@@ -48,8 +47,9 @@ internal fun AsmGen.translateControl(insn: IRInstruction, forwardedImmediateCall
             val ref = insn.requireTarget() as? CodeReference.Indirect ?: error("JUMPI needs an indirect target register")
             val reg = ref.pointer.intNumber
             invalidateD0Cache()
-            // a dispatch jump leaves this subroutine for good: the target's rts must find our caller's return address, so our frame has to be gone first
-            emitFrameUnlk()
+            // Keep the frame for a table whose targets all belong to this subroutine.
+            if (!jumpTableStaysInCurrentSub(ref.pointer.register))
+                emitFrameUnlk()
             if(program.options.compTarget.cpu >= CpuType.M68020) {
                 // 68020+ supports memory-indirect addressing: fetch the target address from memory directly.
                 emitLine("jmp  ([${regAddr(reg)}])")

codeGenM68k/test/TestInstructionSelectionOptimizations.kt

@@ -1208,6 +1208,41 @@ class TestInstructionSelectionOptimizations : FunSpec({
         lines.count { it.contains("(a0,d0.w*") } shouldBe 0
     }
 
+    test("loadx+jumpi fusion removes the regfile roundtrip and keeps the scaled index operand") {
+        // a pointer-array dispatch (on..goto / rewritten when) emits the scale in the addressing mode
+        val lines = generateAsm(
+            tempRoot.resolve("test-m68k-fuse-loadx-jumpi-scaled"),
+            listOf(
+                IRInstructions.loadMemory(Opcode.LOADX, IRDataType.POINTER, 6, IRMemory.indexed("table", 3, IRDataType.WORD, 4)),
+                IRInstructions.jumpIndirect(6)
+            )
+        )
+
+        lines.any { it == "move.l  (a0,d0.w*4), a0" } shouldBe true
+        lines.any { it == "jmp  (a0)" } shouldBe true
+        // the spill and the reload it feeds are gone
+        lines.count { it.startsWith("move.l  d0,p8_regfile") } shouldBe 0
+        lines.count { it.startsWith("movea.l") } shouldBe 0
+        lines.any { it.startsWith("jmp  ([p8_regfile") } shouldBe false
+    }
+
+    test("loadx+jumpi fusion keeps the unscaled index operand on 68000") {
+        val lines = generateAsm(
+            tempRoot.resolve("test-m68k-fuse-loadx-jumpi-unscaled"),
+            listOf(
+                IRInstructions.loadMemory(Opcode.LOADX, IRDataType.POINTER, 6, IRMemory.indexed("table", 3, IRDataType.WORD, 4)),
+                IRInstructions.jumpIndirect(6)
+            ),
+            target = Amiga500Target()
+        )
+
+        // on 68000 the scale is applied to d0 by hand, so the addressing mode stays unscaled
+        lines.any { it == "lsl.w  #2,d0" } shouldBe true
+        lines.any { it == "move.l  (a0,d0.w), a0" } shouldBe true
+        lines.any { it == "jmp  (a0)" } shouldBe true
+        lines.count { it.startsWith("move.l  d0,p8_regfile") } shouldBe 0
+    }
+
     test("indirect memory lowering picks displacement addressing or explicit pointer adjustment") {
         val lines = generateAsm(
             tempRoot.resolve("test-m68k-indirect-displacement"),

docs/source/todo.rst

@@ -66,7 +66,7 @@ Optimizations
 - M68K stack/frame follow-ups: the stack-based memory model for locals and parameters is implemented, including frame slots for locals, parameters in the caller-pushed incoming area (word-packed argument slots are now the default convention), per-activation virtual registers for re-entrant subroutines, prologue zero-store coalescing, and diagnostics/polish items. Remaining low-pri optimization work: reserved outgoing-argument area (allocate the max call-site arg size once in the caller's ``link`` frame; benchmark-driven, likely marginal) and loop-aware liveness for frame slot reuse (needs real backward live-in/live-out dataflow over the CFG; frame-size-driven). Register argument passing lives in ``ideas/m68k-register-allocation.md`` §9; recursion-sound defers is an IR/VM item above. See the user-facing description in ``technical.rst``.
 - ``when`` → ``on..goto``: extend beyond goto-only value choice bodies. ``When.onGotoIndexOffset`` only accepts a ``when`` whose value choices consist of a single ``goto`` (the ``else`` may be any body, it is taken over verbatim as the ``on..goto``'s else part), which is rare in hand-written code. Generalizing it means synthesizing a label per choice body and appending the bodies after the ``when`` with ``AstInsert.after``, so choices containing function calls or assignments become eligible too. This would greatly widen the set of programs that benefit.
 - ``when`` → ``on..goto``: gap filling requires an ``else`` whose body is a single ``goto``, because a gap slot in the jump table has no choice to point at and can only reuse the else target. A gap value could equally well fall through to the statement after the ``when``. That is just a synthesized empty ``Label`` inserted with ``AstInsert.after(when, label, parent)``, with the gap slots pointing at it; labels are valid inside nested scopes (verified: a label and a ``goto`` to it inside an ``if`` body compile fine), and in the compiler AST a statement's parent is always an ``IStatementContainer``. The only missing piece is a unique-name helper, since nothing in the compiler AST constructs a label yet.
-- ``on..goto``: the m68k loadx+jumpi fusion never fires for dispatch sites. ``optimizeFuseLoadxJumpi`` in the m68k ``AsmOptimizer`` already exists and would remove the spill/reload, but its matcher requires the unscaled ``move.l (a0,d0.w), d0``, while an ``on..goto``/``when`` table is a pointer array and is emitted as the scaled ``move.l (a0,d0.w*4), d0``. So the ``p8_regfile`` round-trip (``move.l d0,p8_regfile+M`` then ``jmp ([p8_regfile+M])``) survives, as does the three-instruction byte-to-word index extension (``moveq #0,d0`` / ``move.b ...,d0`` / ``move.w d0,...``). Teaching the fusion to recognise the scaled form would fix both dispatch paths at once. Keeping the index in a hardware register would help further; related: the deferred m68k register allocator above.
+- ``on..goto``: the index is loaded twice per dispatch. The index is read once for the ``index < numlabels`` bounds check and again for the table lookup, on every backend (cx16 emits ``lda sel`` twice with ``cmp #6``/``bcs`` in between; m68k emits ``move.b sel,p8_regfile+0`` then ``move.b sel,p8_regfile+2``). Making ``CodeDesugarer.after(OnGoto)`` always use a temp var is *not* the fix: measured on cx16 that grows a 6-way dispatch from 13 to 16 bytes, since it replaces the first load with ``lda``/``sta`` and still reloads the temp. The value is genuinely needed twice at the AST level, so this needs register reuse in ``IRPeepholeOptimizer``: ``removeDoubleLoadsAndStores`` only works within one chunk, whereas these two loads sit either side of a ``bge``. Note the file's own comment there already lists "detect multiple loads to the same target registers, only keep first (if source is not I/O memory)" as a known un-implemented optimization. A cross-chunk version must first prove the fall-through is the only way into the second chunk, since another branch or jump there would leave the register undefined; there is no liveness data in the IR to lean on (``RegisterPool`` just hands out the next free register).
 - Calibrate the ``when`` → ``on..goto`` ``sizeLimit`` per target instead of the current empirical 4 (65C02) / 5 (m68k) / 6 (other). Measured ``.prg`` sizes: on cx16 (6502, 16-bit pointers) the table already wins from 4 choices up (348 vs 353 bytes at 4 choices, 362 vs 375 at 6), growing ~7 bytes per choice against ~14 for the compare chain, so the cx16 limit is about right. On qemu68k the 32-bit slots invert it: +103 bytes at 6 choices, +63 at 8, break-even at 12, and a win from ~20 up. So the current limit of 5 makes small m68k dispatchers bigger for nothing. Note the speed is never the problem: VM instruction counts for the dispatch are ~30% lower with the table at every size measured, so the limit should be chosen for when the extra ROM is worth it, not for raw speed.