From ac5f45d2d4c5aa2dd3600e30ece17d0f6a1f118a Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Wed, 27 Oct 2021 02:30:58 +0200
Subject: [PATCH] fix nested label lookups in anon scopes (partly)

---
 .../astprocessing/AstIdentifiersChecker.kt    |   2 +-
 compiler/test/TestScoping.kt                  | 112 ++++++++++++++++++
 compiler/test/TestSubroutines.kt              |  30 -----
 compilerAst/src/prog8/ast/AstToplevel.kt      |  84 ++++++++++---
 compilerAst/test/TestProg8Parser.kt           |  31 +++++
 examples/test.p8                              |  30 +++--
 6 files changed, 232 insertions(+), 57 deletions(-)
 create mode 100644 compiler/test/TestScoping.kt

diff --git a/compiler/src/prog8/compiler/astprocessing/AstIdentifiersChecker.kt b/compiler/src/prog8/compiler/astprocessing/AstIdentifiersChecker.kt
index 4e4657ec5..caa12e98b 100644
--- a/compiler/src/prog8/compiler/astprocessing/AstIdentifiersChecker.kt
+++ b/compiler/src/prog8/compiler/astprocessing/AstIdentifiersChecker.kt
@@ -75,7 +75,7 @@ internal class AstIdentifiersChecker(private val program: Program, private val e
             val paramNames = subroutine.parameters.map { it.name }.toSet()
             val paramsToCheck = paramNames.intersect(namesInSub)
             for(name in paramsToCheck) {
-                val labelOrVar = subroutine.getLabelOrVariable(name)
+                val labelOrVar = subroutine.searchLabelOrVariableNotSubscoped(name)
                 if(labelOrVar!=null && labelOrVar.position != subroutine.position)
                     nameError(name, labelOrVar.position, subroutine)
                 val sub = subroutine.statements.firstOrNull { it is Subroutine && it.name==name}
diff --git a/compiler/test/TestScoping.kt b/compiler/test/TestScoping.kt
new file mode 100644
index 000000000..08901e109
--- /dev/null
+++ b/compiler/test/TestScoping.kt
@@ -0,0 +1,112 @@
+package prog8tests
+
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.TestInstance
+import prog8.ast.expressions.NumericLiteralValue
+import prog8.ast.statements.*
+import prog8.compiler.target.C64Target
+import prog8tests.helpers.assertSuccess
+import prog8tests.helpers.compileText
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class TestScoping {
+
+    @Test
+    fun testAnonScopeVarsMovedIntoSubroutineScope() {
+        val src = """
+            main {
+                sub start() {
+                    repeat {
+                        ubyte xx = 99
+                        xx++
+                    }
+                }
+            }
+        """
+
+        val result = compileText(C64Target, false, src, writeAssembly = false).assertSuccess()
+        val module = result.programAst.toplevelModule
+        val mainBlock = module.statements.single() as Block
+        val start = mainBlock.statements.single() as Subroutine
+        val repeatbody = start.statements.filterIsInstance<RepeatLoop>().single().body
+        assertFalse(mainBlock.statements.any { it is VarDecl }, "no vars moved to main block")
+        val subroutineVars = start.statements.filterIsInstance<VarDecl>()
+        assertEquals(1, subroutineVars.size, "var from repeat anonscope must be moved up to subroutine")
+        assertEquals("xx", subroutineVars[0].name)
+        assertFalse(repeatbody.statements.any { it is VarDecl }, "var should have been removed from repeat anonscope")
+        val initassign = repeatbody.statements[0] as? Assignment
+        assertEquals(listOf("xx"), initassign?.target?.identifier?.nameInSource, "vardecl in repeat should be replaced by init assignment")
+        assertEquals(99, (initassign?.value as? NumericLiteralValue)?.number?.toInt(), "vardecl in repeat should be replaced by init assignment")
+        assertTrue(repeatbody.statements[1] is PostIncrDecr)
+    }
+
+    @Test
+    fun testLabelsWithAnonScopes() {
+        val src = """
+            main {
+                sub start() {
+                    uword addr
+                    goto labeloutside
+        
+                    if true {
+                        if true {
+                            addr = &iflabel
+                            addr = &labelinside
+                            addr = &labeloutside
+                            addr = &main.start.nested.nestedlabel
+                            ; addr = &nested.nestedlabel      ; TODO should also work!!
+                            goto labeloutside
+                            goto iflabel
+                            goto main.start.nested.nestedlabel
+                            ; goto nested.nestedlabel         ; TODO should also work!!
+                        }
+            iflabel:
+                    }
+        
+                    repeat {
+                        addr = &iflabel
+                        addr = &labelinside
+                        addr = &labeloutside
+                        addr = &main.start.nested.nestedlabel
+                        ; addr = &nested.nestedlabel      ; TODO should also work!!
+                        goto iflabel
+                        goto labelinside
+                        goto main.start.nested.nestedlabel
+                        ; goto nested.nestedlabel         ; TODO should also work!!
+            labelinside:
+                    }
+
+                    sub nested () {
+            nestedlabel:
+                        addr = &nestedlabel
+                        goto nestedlabel
+                        goto main.start.nested.nestedlabel
+                    }
+
+            labeloutside:
+                    addr = &iflabel
+                    addr = &labelinside
+                    addr = &labeloutside
+                    addr = &main.start.nested.nestedlabel
+                    ; addr = &nested.nestedlabel      ; TODO should also work!!
+                    goto main.start.nested.nestedlabel
+                    ; goto nested.nestedlabel     ; TODO should also work!!
+                }
+            }
+        """
+
+        val result = compileText(C64Target, false, src, writeAssembly = true).assertSuccess()
+        val module = result.programAst.toplevelModule
+        val mainBlock = module.statements.single() as Block
+        val start = mainBlock.statements.single() as Subroutine
+        val labels = start.statements.filterIsInstance<Label>()
+        assertEquals(1, labels.size, "only one label in subroutine scope")
+
+        TODO("fix the partial scoped lookups above as well")
+    }
+
+}
diff --git a/compiler/test/TestSubroutines.kt b/compiler/test/TestSubroutines.kt
index ff262d975..87a384e21 100644
--- a/compiler/test/TestSubroutines.kt
+++ b/compiler/test/TestSubroutines.kt
@@ -5,7 +5,6 @@ import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.TestInstance
 import prog8.ast.base.DataType
 import prog8.ast.expressions.IdentifierReference
-import prog8.ast.expressions.NumericLiteralValue
 import prog8.ast.expressions.TypecastExpression
 import prog8.ast.statements.*
 import prog8.compiler.target.C64Target
@@ -182,33 +181,4 @@ class TestSubroutines {
         assertEquals(DataType.ARRAY_UB, func.parameters.single().type)
         assertTrue(func.statements.isEmpty())
     }
-
-    @Test
-    fun testAnonScopeVarsMovedIntoSubroutineScope() {
-        val src = """
-            main {
-                sub start() {
-                    repeat {
-                        ubyte xx = 99
-                        xx++
-                    }
-                }
-            }
-        """
-
-        val result = compileText(C64Target, false, src, writeAssembly = false).assertSuccess()
-        val module = result.programAst.toplevelModule
-        val mainBlock = module.statements.single() as Block
-        val start = mainBlock.statements.single() as Subroutine
-        val repeatbody = start.statements.filterIsInstance<RepeatLoop>().single().body
-        assertFalse(mainBlock.statements.any { it is VarDecl }, "no vars moved to main block")
-        val subroutineVars = start.statements.filterIsInstance<VarDecl>()
-        assertEquals(1, subroutineVars.size, "var from repeat anonscope must be moved up to subroutine")
-        assertEquals("xx", subroutineVars[0].name)
-        assertFalse(repeatbody.statements.any { it is VarDecl }, "var should have been removed from repeat anonscope")
-        val initassign = repeatbody.statements[0] as? Assignment
-        assertEquals(listOf("xx"), initassign?.target?.identifier?.nameInSource, "vardecl in repeat should be replaced by init assignment")
-        assertEquals(99, (initassign?.value as? NumericLiteralValue)?.number?.toInt(), "vardecl in repeat should be replaced by init assignment")
-        assertTrue(repeatbody.statements[1] is PostIncrDecr)
-    }
 }
diff --git a/compilerAst/src/prog8/ast/AstToplevel.kt b/compilerAst/src/prog8/ast/AstToplevel.kt
index 015272ac0..bc5970d65 100644
--- a/compilerAst/src/prog8/ast/AstToplevel.kt
+++ b/compilerAst/src/prog8/ast/AstToplevel.kt
@@ -84,6 +84,71 @@ interface IStatementContainer {
     fun isEmpty(): Boolean = statements.isEmpty()
     fun isNotEmpty(): Boolean = statements.isNotEmpty()
 
+    fun searchLabelOrVariableNotSubscoped(name: String): Statement? {         // TODO return INamedStatement instead?  and rename to searchSymbol ?
+        // this is called quite a lot and could perhaps be optimized a bit more,
+        // but adding a memoization cache didn't make much of a practical runtime difference...
+        for (stmt in statements) {
+            when(stmt) {
+//                is INamedStatement -> {
+//                    if(stmt.name==name) return stmt
+//                }
+                is VarDecl -> {
+                    if(stmt.name==name) return stmt
+                }
+                is Label -> {
+                    if(stmt.name==name) return stmt
+                }
+                is AnonymousScope -> {
+                    val found = stmt.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is IfStatement -> {
+                    val found = stmt.truepart.searchLabelOrVariableNotSubscoped(name) ?: stmt.elsepart.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is BranchStatement -> {
+                    val found = stmt.truepart.searchLabelOrVariableNotSubscoped(name) ?: stmt.elsepart.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is ForLoop -> {
+                    val found = stmt.body.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is WhileLoop -> {
+                    val found = stmt.body.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is RepeatLoop -> {
+                    val found = stmt.body.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is UntilLoop -> {
+                    val found = stmt.body.searchLabelOrVariableNotSubscoped(name)
+                    if(found!=null)
+                        return found
+                }
+                is WhenStatement -> {
+                    stmt.choices.forEach {
+                        val found = it.statements.searchLabelOrVariableNotSubscoped(name)
+                        if(found!=null)
+                            return found
+                    }
+                }
+                else -> {
+                    // NOTE: when other nodes containing AnonymousScope are introduced,
+                    //       these should be added here as well to look into!
+                }
+            }
+        }
+        return null
+    }
+
     val allDefinedSymbols: List<Pair<String, Statement>>
         get() {
             return statements.mapNotNull {
@@ -101,17 +166,7 @@ interface IStatementContainer {
 interface INameScope: IStatementContainer, INamedStatement {
     fun subScope(name: String): INameScope?  = statements.firstOrNull { it is INameScope && it.name==name } as? INameScope
 
-    fun getLabelOrVariable(name: String): Statement? {
-        // this is called A LOT and could perhaps be optimized a bit more,
-        // but adding a memoization cache didn't make much of a practical runtime difference
-        for (stmt in statements) {
-            if (stmt is VarDecl && stmt.name==name) return stmt
-            if (stmt is Label && stmt.name==name) return stmt
-        }
-        return null
-    }
-
-    fun lookup(scopedName: List<String>, localContext: Node) : Statement? {
+    fun lookup(scopedName: List<String>, localContext: Node) : Statement? {             // TODO return INamedStatement instead?
         if(scopedName.size>1) {
             // a scoped name refers to a name in another module.
             // it's a qualified name, look it up from the root of the module's namespace (consider all modules in the program)
@@ -123,7 +178,7 @@ interface INameScope: IStatementContainer, INamedStatement {
                         break
                 }
                 if(scope!=null) {
-                    val result = scope.getLabelOrVariable(scopedName.last())
+                    val result = scope.searchLabelOrVariableNotSubscoped(scopedName.last())
                     if(result!=null)
                         return result
                     return scope.subScope(scopedName.last()) as Statement?
@@ -136,7 +191,7 @@ interface INameScope: IStatementContainer, INamedStatement {
             var statementScope = localContext
             while(statementScope !is ParentSentinel) {
                 val localScope = statementScope.definingScope
-                val result = localScope.getLabelOrVariable(scopedName[0])
+                val result = localScope.searchLabelOrVariableNotSubscoped(scopedName[0])
                 if (result != null)
                     return result
                 val subscope = localScope.subScope(scopedName[0]) as Statement?
@@ -168,7 +223,6 @@ interface Node {
 
     val definingSubroutine: Subroutine? get() = findParentNode<Subroutine>(this)
 
-    // TODO CHECK IF THIS IS ACTUALLY WHAT'S DESIRED: (namescope v.s. statementcontainer)
     val definingScope: INameScope
         get() {
             val scope = findParentNode<INameScope>(this)
@@ -374,7 +428,7 @@ class GlobalNamespace(val modules: Iterable<Module>, private val builtinFunction
         throw FatalAstException("cannot replace anything in the namespace")
     }
 
-    override fun lookup(scopedName: List<String>, localContext: Node): Statement? {
+    override fun lookup(scopedName: List<String>, localContext: Node): Statement? {         // TODO return INamedStatement instead?
         if (scopedName.size == 1 && scopedName[0] in builtinFunctionNames) {
             // builtin functions always exist, return a dummy localContext for them
             val builtinPlaceholder = Label("builtin::${scopedName.last()}", localContext.position)
diff --git a/compilerAst/test/TestProg8Parser.kt b/compilerAst/test/TestProg8Parser.kt
index 8fb68c911..5d9f656d7 100644
--- a/compilerAst/test/TestProg8Parser.kt
+++ b/compilerAst/test/TestProg8Parser.kt
@@ -632,4 +632,35 @@ class TestProg8Parser {
         assertTrue(repeatbody.statements[1] is PostIncrDecr)
         // the ast processing steps used in the compiler, will eventually move the var up to the containing scope (subroutine).
     }
+
+    @Test
+    fun testLabelsWithAnonScopesParsesFine() {
+        val src = SourceCode.Text("""
+            main {
+                sub start() {
+                    goto mylabeloutside
+        
+                    if true {
+                        if true {
+                            goto labeloutside
+                            goto iflabel
+                        }
+            iflabel:
+                    }
+        
+                    repeat {
+                        goto labelinside
+            labelinside:
+                    }
+        
+            labeloutside:
+                }
+            }
+        """)
+        val module = parseModule(src)
+        val mainBlock = module.statements.single() as Block
+        val start = mainBlock.statements.single() as Subroutine
+        val labels = start.statements.filterIsInstance<Label>()
+        assertEquals(1, labels.size, "only one label in subroutine scope")
+    }
 }
diff --git a/examples/test.p8 b/examples/test.p8
index b412a7b3e..0bb89f262 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -1,16 +1,24 @@
-
 main {
     sub start() {
-        ubyte col
-        ubyte row
-        repeat {
-            col = rnd() % 33
-            row = rnd() % 33
-            ;cx16logo.logo_at(col, row)
-            ;txt.plot(col-3, row+7 )
-            ;txt.print("commander x16")
-            col++
-            row++
+        goto labeloutside
+
+        if true {
+            if true {
+                goto labeloutside
+                goto iflabel
+            }
+iflabel:
         }
+
+        repeat {
+            goto labelinside
+labelinside:
+        }
+
+labeloutside:
     }
+
+    start:
+    start:
+    start:
 }