more robust 'return' statement checks in subroutines

compiler/src/prog8/compiler/astprocessing/BeforeAsmAstChanger.kt

@@ -135,15 +135,20 @@ internal class BeforeAsmAstChanger(val program: Program,
         val mods = mutableListOf<IAstModification>()
 
         // add the implicit return statement at the end (if it's not there yet), but only if it's not a kernal routine.
-        // and if an assembly block doesn't contain a rts/rti, and some other situations.
+        // and if an assembly block doesn't contain a rts/rti.
         if (!subroutine.isAsmSubroutine) {
-            if(subroutine.statements.isEmpty() ||
+            if(subroutine.isEmpty()) {
-                (!subroutine.hasRtsInAsm(options.compTarget)
-                        && subroutine.statements.lastOrNull { it !is VarDecl } !is Return
-                        && subroutine.statements.last() !is Subroutine
-                        && subroutine.statements.last() !is Return)) {
                 val returnStmt = Return(null, subroutine.position)
                 mods += IAstModification.InsertLast(returnStmt, subroutine)
+            } else {
+                val last = subroutine.statements.last()
+                if((last !is InlineAssembly || !last.hasReturnOrRts(options.compTarget)) && last !is Return) {
+                    val lastStatement = subroutine.statements.reversed().firstOrNull { it !is Subroutine }
+                    if(lastStatement !is Return) {
+                        val returnStmt = Return(null, subroutine.position)
+                        mods += IAstModification.InsertLast(returnStmt, subroutine)
+                    }
+                }
             }
         }
 
@@ -174,6 +179,14 @@ internal class BeforeAsmAstChanger(val program: Program,
             }
         }
 
+        if(subroutine.isNotEmpty() && subroutine.statements.last() is Return) {
+            // maybe the last return can be removed because there is a fall-through prevention above it
+            val lastStatementBefore = subroutine.statements.reversed().drop(1).firstOrNull { it !is Subroutine }
+            if(lastStatementBefore is Return) {
+                mods += IAstModification.Remove(subroutine.statements.last(), subroutine)
+            }
+        }
+
         return mods
     }
 

compiler/test/vm/TestCompilerVirtual.kt

@@ -5,6 +5,7 @@ import io.kotest.matchers.shouldBe
 import io.kotest.matchers.shouldNotBe
 import prog8.ast.expressions.BuiltinFunctionCall
 import prog8.ast.statements.Assignment
+import prog8.code.target.C64Target
 import prog8.code.target.Cx16Target
 import prog8.code.target.VMTarget
 import prog8.vm.VmRunner
@@ -110,6 +111,10 @@ mylabel_inside:
     }
 }"""
 
+        val target1 = C64Target()
+        val result1 = compileText(target1, false, src, writeAssembly = false)!!
+        result1 shouldNotBe null
+
         val target = VMTarget()
         val result = compileText(target, false, src, writeAssembly = true)!!
         val virtfile = result.compilationOptions.outputDir.resolve(result.program.name + ".p8ir")
@@ -164,4 +169,24 @@ main {
             vm.memory.getUB(3) shouldBe 43u
         }
     }
+
+    test("memory mapped var as for loop counter") {
+        val src = """
+main {
+    sub start() {
+        for cx16.r0 in 0 to 10 {
+            cx16.r1++
+        }
+    }
+}"""
+        val othertarget = Cx16Target()
+        compileText(othertarget, true, src, writeAssembly = true, keepIR=true) shouldNotBe null
+
+        val target = VMTarget()
+        val result = compileText(target, false, src, writeAssembly = true)!!
+        val virtfile = result.compilationOptions.outputDir.resolve(result.program.name + ".p8ir")
+        VmRunner().runAndTestProgram(virtfile.readText()) { vm ->
+            vm.stepCount shouldBe 9999  // TODO FIX
+        }
+    }
 })

docs/source/todo.rst

@@ -4,7 +4,7 @@ TODO
 For next release
 ^^^^^^^^^^^^^^^^
 - ir: fix unit tests
-- ir: join more codechunks
+- ir: join more codechunks, before ir is written to file
 - ir: fix removeWeirdBranches in IR optimizer
 - update diagram in technical.rst?