improvements for inlined subroutines: fix identifier scoping

compiler/src/prog8/compiler/Compiler.kt

@@ -253,7 +253,8 @@ private fun processAst(programAst: Program, errors: IErrorReporter, compilerOpti
     errors.report()
     programAst.addTypecasts(errors)
     errors.report()
-    programAst.variousCleanups()
+    programAst.variousCleanups(programAst, errors, compilerOptions)
+    errors.report()
     programAst.checkValid(compilerOptions, errors, compilerOptions.compTarget)
     errors.report()
     programAst.checkIdentifiers(errors, compilerOptions.compTarget)
@@ -283,7 +284,7 @@ private fun optimizeAst(programAst: Program, errors: IErrorReporter, functions:
 private fun postprocessAst(programAst: Program, errors: IErrorReporter, compilerOptions: CompilationOptions) {
     programAst.addTypecasts(errors)
     errors.report()
-    programAst.variousCleanups()
+    programAst.variousCleanups(programAst, errors, compilerOptions)
     programAst.checkValid(compilerOptions, errors, compilerOptions.compTarget)          // check if final tree is still valid
     errors.report()
     val callGraph = CallGraph(programAst, ::loadAsmIncludeFile)

compiler/src/prog8/compiler/astprocessing/AstExtensions.kt

@@ -56,12 +56,14 @@ internal fun Program.checkIdentifiers(errors: IErrorReporter, compTarget: ICompi
     }
 }
 
-internal fun Program.variousCleanups() {
-    val process = VariousCleanups()
+internal fun Program.variousCleanups(program: Program, errors: IErrorReporter, options: CompilationOptions) {
+    val process = VariousCleanups(program, errors, options)
     process.visit(this)
-    process.applyModifications()
+    if(errors.isEmpty())
+        process.applyModifications()
 }
 
+
 internal fun Program.moveMainAndStartToFirst() {
     // the module containing the program entrypoint is moved to the first in the sequence.
     // the "main" block containing the entrypoint is moved to the top in there,

compiler/src/prog8/compiler/astprocessing/VariousCleanups.kt

@@ -3,17 +3,17 @@ package prog8.compiler.astprocessing
 import prog8.ast.IFunctionCall
 import prog8.ast.INameScope
 import prog8.ast.Node
+import prog8.ast.Program
 import prog8.ast.base.Position
-import prog8.ast.expressions.DirectMemoryRead
-import prog8.ast.expressions.FunctionCall
-import prog8.ast.expressions.NumericLiteralValue
-import prog8.ast.expressions.TypecastExpression
+import prog8.ast.expressions.*
 import prog8.ast.statements.*
 import prog8.ast.walk.AstWalker
 import prog8.ast.walk.IAstModification
+import prog8.compiler.CompilationOptions
+import prog8.compiler.IErrorReporter
 
 
-internal class VariousCleanups: AstWalker() {
+internal class VariousCleanups(private val program: Program, val errors: IErrorReporter, private val compilerOptions: CompilationOptions): AstWalker() {
     private val noModifications = emptyList<IAstModification>()
 
     override fun before(nopStatement: NopStatement, parent: Node): Iterable<IAstModification> {
@@ -56,6 +56,14 @@ internal class VariousCleanups: AstWalker() {
     }
 
     private fun before(functionCall: IFunctionCall, parent: Node, position: Position): Iterable<IAstModification> {
+
+        if(compilerOptions.optimize) {
+            val sub = functionCall.target.targetSubroutine(program)
+            if(sub!=null && sub.inline && !sub.isAsmSubroutine)
+                annotateInlinedSubroutineIdentifiers(sub)
+        }
+
+
         if(functionCall.target.nameInSource==listOf("peek")) {
             // peek(a) is synonymous with @(a)
             val memread = DirectMemoryRead(functionCall.args.single(), position)
@@ -70,4 +78,35 @@ internal class VariousCleanups: AstWalker() {
         return noModifications
     }
 
+    private fun annotateInlinedSubroutineIdentifiers(sub: Subroutine) {
+        // this adds full name prefixes to all identifiers used in the subroutine,
+        // so that the statements can be inlined (=copied) in the call site and still reference
+        // the correct symbols as seen from the scope of the subroutine.
+
+        class Annotator: AstWalker() {
+            var numReturns=0
+
+            override fun before(identifier: IdentifierReference, parent: Node): Iterable<IAstModification> {
+                val stmt = identifier.targetStatement(program)!!
+                val prefixed = stmt.makeScopedName(identifier.nameInSource.last()).split('.')
+                val withPrefix = IdentifierReference(prefixed, identifier.position)
+                return listOf(IAstModification.ReplaceNode(identifier, withPrefix, identifier.parent))
+            }
+
+            override fun before(returnStmt: Return, parent: Node): Iterable<IAstModification> {
+                numReturns++
+                if(parent !== sub || sub.indexOfChild(returnStmt)<sub.statements.size-1)
+                    errors.err("return statement must be the very last statement in the inlined subroutine", sub.position)
+                return emptyList()
+            }
+        }
+
+        val annotator = Annotator()
+        sub.accept(annotator, sub.parent)
+        if(annotator.numReturns>1)
+            errors.err("inlined subroutine can only have one return statement", sub.position)
+        else
+            annotator.applyModifications()
+    }
+
 }

compiler/src/prog8/compiler/target/cpu6502/codegen/FunctionCallAsmGen.kt

@@ -111,7 +111,10 @@ internal class FunctionCallAsmGen(private val program: Program, private val asmg
             }
         }
 
-        if(sub.inline && asmgen.options.optimize) {
+        if(!sub.inline || !asmgen.options.optimize) {
+            asmgen.out("  jsr  $subName")
+        } else {
+            // inline the subroutine
             if(sub.containsDefinedVariables())
                 throw AssemblyError("can't inline sub with vars")
             if(!sub.isAsmSubroutine && sub.parameters.isNotEmpty())
@@ -120,9 +123,6 @@ internal class FunctionCallAsmGen(private val program: Program, private val asmg
             val statements = sub.statements.filter { it !is ParameterVarDecl && it !is Directive }
             statements.forEach { asmgen.translate(it) }
         }
-        else {
-            asmgen.out("  jsr  $subName")
-        }
 
         // remember: dealing with the X register and/or dealing with return values is the responsibility of the caller
     }

compilerAst/src/prog8/ast/AstToSourceCode.kt

@@ -137,8 +137,11 @@ class AstToSourceCode(val output: (text: String) -> Unit, val program: Program):
 
     override fun visit(subroutine: Subroutine) {
         output("\n")
+        outputi("")
+        if(subroutine.inline)
+            output("inline ")
         if(subroutine.isAsmSubroutine) {
-            outputi("asmsub ${subroutine.name} (")
+            output("asmsub ${subroutine.name} (")
             for(param in subroutine.parameters.zip(subroutine.asmParameterRegisters)) {
                 val reg =
                         when {
@@ -152,7 +155,7 @@ class AstToSourceCode(val output: (text: String) -> Unit, val program: Program):
             }
         }
         else {
-            outputi("sub ${subroutine.name} (")
+            output("sub ${subroutine.name} (")
             for(param in subroutine.parameters) {
                 output("${datatypeString(param.type)} ${param.name}")
                 if(param!==subroutine.parameters.last())

compilerAst/src/prog8/ast/expressions/AstExpressions.kt

@@ -711,8 +711,7 @@ internal fun makeRange(fromVal: Int, toVal: Int, stepVal: Int): IntProgression {
     }
 }
 
-data class IdentifierReference(val nameInSource: List<String>, override val position: Position) : Expression(),
-    IAssignable {
+data class IdentifierReference(val nameInSource: List<String>, override val position: Position) : Expression(), IAssignable {
     override lateinit var parent: Node
 
     fun targetStatement(program: Program) =

docs/source/todo.rst

@@ -2,6 +2,8 @@
 TODO
 ====
 
+- fix inlining issue with return statement inlined literally at call site
+
 - allow inlining of subroutines with vardecls
 - optimize several inner loops in gfx2
 - hoist all variable declarations up to the subroutine scope *before* even the constant folding takes place (to avoid undefined symbol errors when referring to a variable from another nested scope in the subroutine)

examples/test.p8

@@ -1,11 +1,12 @@
 %import textio
 %zeropage basicsafe
+%option no_sysinit
 
 main {
 
     sub start() {
         ubyte thing = otherblock.othersub()
-        txt.print_ub(thing)     ; should print 99!
+        txt.print_ub(thing)     ; should print 21!
 
 ;        str filename = "titlescreen.bin"
 ;        ubyte success = cx16.vload(filename, 8, 0, $0000)
@@ -22,13 +23,21 @@ main {
     }
 }
 
+block3 {
+    ubyte returnvalue=10
+
+    sub thing()->ubyte {
+        return returnvalue
+    }
+}
+
 otherblock {
 
     ubyte othervar=20
     ubyte calcparam=10
 
     sub calc(ubyte zz) -> ubyte {
-        return zz+1
+        return zz+1+block3.thing()
     }
 
     inline sub othersub() -> ubyte {