From d4b69ac79c6bdd46ead746ff6b1c43640d1d3409 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Sun, 30 May 2021 15:30:34 +0200
Subject: [PATCH] improved repeat counter vars allocation (re-use var if
 possible)

---
 .../compiler/target/cpu6502/codegen/AsmGen.kt | 42 +++++++++++--------
 .../src/prog8/ast/statements/AstStatements.kt |  3 +-
 docs/source/todo.rst                          |  1 -
 examples/test.p8                              |  3 +-
 4 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt b/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
index 9f8dc839a..d3552a200 100644
--- a/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
+++ b/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
@@ -880,11 +880,10 @@ internal class AsmGen(private val program: Program,
             }
 
             out("; variables")
-            for((name, addr) in sub.asmGenInfo.extraVarsZP) {
-                out("$name = $addr")
-            }
-            for((dt, name) in sub.asmGenInfo.extraVars) {
-                when(dt) {
+            for((dt, name, addr) in sub.asmGenInfo.extraVars) {
+                if(addr!=null)
+                    out("$name = $addr")
+                else when(dt) {
                     DataType.UBYTE -> out("$name    .byte  0")
                     DataType.UWORD -> out("$name    .word  0")
                     else -> throw AssemblyError("weird dt")
@@ -1069,33 +1068,40 @@ $repeatLabel    lda  $counterVar
     }
 
     private fun createRepeatCounterVar(dt: DataType, constIterations: Int?, stmt: RepeatLoop): String {
-        // TODO share counter variables between subroutines or even between repeat loops as long as they're not nested
-//        var parent = stmt.parent
-//        while(parent !is ParentSentinel) {
-//            if(parent is RepeatLoop)
-//                break
-//            parent = parent.parent
-//        }
-//        val isNested = parent is RepeatLoop
-        val counterVar = makeLabel("repeatcounter")
         val asmInfo = stmt.definingSubroutine()!!.asmGenInfo
+        var parent = stmt.parent
+        while(parent !is ParentSentinel) {
+            if(parent is RepeatLoop)
+                break
+            parent = parent.parent
+        }
+        val isNested = parent is RepeatLoop
+
+        if(!isNested) {
+            // we can re-use a counter var from the subroutine if it already has one for that datatype
+            val existingVar = asmInfo.extraVars.firstOrNull { it.first==dt }
+            if(existingVar!=null)
+                return existingVar.second
+        }
+
+        val counterVar = makeLabel("repeatcounter")
         when(dt) {
             DataType.UBYTE -> {
                 if(constIterations!=null && constIterations>=16 && zeropage.hasByteAvailable()) {
                     // allocate count var on ZP
                     val zpAddr = zeropage.allocate(counterVar, DataType.UBYTE, stmt.position, errors)
-                    asmInfo.extraVarsZP.add(Pair(counterVar, zpAddr))
+                    asmInfo.extraVars.add(Triple(DataType.UBYTE, counterVar, zpAddr))
                 } else {
-                    asmInfo.extraVars.add(Pair(DataType.UBYTE, counterVar))
+                    asmInfo.extraVars.add(Triple(DataType.UBYTE, counterVar, null))
                 }
             }
             DataType.UWORD -> {
                 if(constIterations!=null && constIterations>=16 && zeropage.hasWordAvailable()) {
                     // allocate count var on ZP
                     val zpAddr = zeropage.allocate(counterVar, DataType.UWORD, stmt.position, errors)
-                    asmInfo.extraVarsZP.add(Pair(counterVar, zpAddr))
+                    asmInfo.extraVars.add(Triple(DataType.UWORD, counterVar, zpAddr))
                 } else {
-                    asmInfo.extraVars.add(Pair(DataType.UWORD, counterVar))
+                    asmInfo.extraVars.add(Triple(DataType.UWORD, counterVar, null))
                 }
             }
             else -> throw AssemblyError("invalidt dt")
diff --git a/compilerAst/src/prog8/ast/statements/AstStatements.kt b/compilerAst/src/prog8/ast/statements/AstStatements.kt
index e5d3aa0df..2b77496a3 100644
--- a/compilerAst/src/prog8/ast/statements/AstStatements.kt
+++ b/compilerAst/src/prog8/ast/statements/AstStatements.kt
@@ -589,8 +589,7 @@ class AsmGenInfo {
     var usedFloatEvalResultVar1 = false
     var usedFloatEvalResultVar2 = false
 
-    val extraVars = mutableListOf<Pair<DataType, String>>()
-    val extraVarsZP = mutableListOf<Pair<String, Int>>()
+    val extraVars = mutableListOf<Triple<DataType, String, Int?>>()
 }
 
 // the subroutine class covers both the normal user-defined subroutines,
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index 3edde3869..c8d812cb4 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -5,7 +5,6 @@ TODO
 - test all examples (including imgviewer, assembler and petaxian) before release of the new version
 
 - simplify cx16.joystick_get2() once this cx16 rom issue is resolved: https://github.com/commanderx16/x16-rom/issues/203
-- improve repeat counter vars allocation in asmgen.createRepeatCounterVar()
 - c64: make the graphics.BITMAP_ADDRESS configurable (VIC banking)
 - get rid of all other TODO's in the code ;-)
 
diff --git a/examples/test.p8 b/examples/test.p8
index 96bc979bc..989603c95 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -1,11 +1,10 @@
-%import textio ; txt.*
+%import textio
 %zeropage basicsafe
 
 main {
 
     ; test program for the optimization of repeat var allocation (asmgen.createRepeatCounterVar)
     ; output must be:  60  6164  6224  12328
-    ; original program size: $046b
 
     uword xx