preparing to optimize 16 bit repeat loop

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

@@ -976,10 +976,12 @@ internal class AsmGen(private val program: Program,
                 val name = asmVariableName(stmt.iterations as IdentifierReference)
                 when(vardecl.datatype) {
                     DataType.UBYTE, DataType.BYTE -> {
-                        repeatByteCountVar(name, repeatLabel, endLabel, stmt.body)
+                        assignVariableToRegister(name, RegisterOrPair.A)
+                        repeatByteCountInA(null, repeatLabel, endLabel, stmt.body)
                     }
                     DataType.UWORD, DataType.WORD -> {
-                        repeatWordCountVar(name, repeatLabel, endLabel, stmt.body)
+                        assignVariableToRegister(name, RegisterOrPair.AY)
+                        repeatWordCountInAY(null, repeatLabel, endLabel, stmt.body)
                     }
                     else -> throw AssemblyError("invalid loop variable datatype $vardecl")
                 }
@@ -1009,22 +1011,26 @@ internal class AsmGen(private val program: Program,
         if(constIterations==0)
             return
         // note: A/Y must have been loaded with the number of iterations already!
-        // TODO can be even more optimized by iterating over pages
         val counterVar = makeLabel("repeatcounter")
         out("""
-                sta  $counterVar
+            cmp  #0
-                sty  $counterVar+1
+            bne  +
-$repeatLabel    lda  $counterVar
+            cpy  #0
-                bne  +
+            beq  $endLabel      ; skip if 0 iterations
-                lda  $counterVar+1
++           sta  $counterVar
-                beq  $endLabel
+            sty  $counterVar+1
-+               lda  $counterVar
+$repeatLabel
-                bne  +
+            lda  $counterVar
-                dec  $counterVar+1
+            bne  +
-+               dec  $counterVar
+            lda  $counterVar+1
-""")
+            beq  $endLabel
++           lda  $counterVar
+            bne  +
+            dec  $counterVar+1
++           dec  $counterVar""")
         translate(body)
         jmp(repeatLabel)
+
         if(constIterations!=null && constIterations>=16 && zeropage.available() > 1) {
             // allocate count var on ZP
             val zpAddr = zeropage.allocate(counterVar, DataType.UWORD, body.position, errors)
@@ -1033,8 +1039,8 @@ $repeatLabel    lda  $counterVar
             out("""
 $counterVar    .word  0""")
         }
-        out(endLabel)
 
+        out(endLabel)
     }
 
     private fun repeatByteCountInA(constIterations: Int?, repeatLabel: String, endLabel: String, body: AnonymousScope) {
@@ -1055,46 +1061,6 @@ $counterVar    .byte  0""")
         out(endLabel)
     }
 
-    private fun repeatByteCountVar(repeatCountVar: String, repeatLabel: String, endLabel: String, body: AnonymousScope) {
-        // note: cannot use original counter variable because it should retain its original value
-        val counterVar = makeLabel("repeatcounter")
-        out("  lda  $repeatCountVar |  beq  $endLabel |  sta  $counterVar")
-        out(repeatLabel)
-        translate(body)
-        out("  dec  $counterVar |  bne  $repeatLabel")
-        // inline countervar:
-        out("""
-                beq  $endLabel  
-$counterVar    .byte  0""")
-        out(endLabel)
-    }
-
-    private fun repeatWordCountVar(repeatCountVar: String, repeatLabel: String, endLabel: String, body: AnonymousScope) {
-        // TODO can be even more optimized by iterating over pages
-        // note: cannot use original counter variable because it should retain its original value
-        val counterVar = makeLabel("repeatcounter")
-        out("""
-            lda  $repeatCountVar
-            sta  $counterVar
-            ora  $repeatCountVar+1
-            beq  $endLabel
-            lda  $repeatCountVar+1
-            sta  $counterVar+1""")
-        out(repeatLabel)
-        translate(body)
-        out("""
-            lda  $counterVar
-            bne  +
-            dec  $counterVar+1
-+           dec  $counterVar
-            lda  $counterVar
-            ora  $counterVar+1
-            bne  $repeatLabel
-            beq  $endLabel  
-$counterVar    .word  0""")
-        out(endLabel)
-    }
-
     private fun translate(stmt: WhileLoop) {
         checkBooleanExpression(stmt.condition)  // we require the condition to be of the form  'x <comparison> <value>'
         val booleanCondition = stmt.condition as BinaryExpression

examples/test.p8

@@ -7,17 +7,35 @@ main {
         uword xx1
         uword xx2
         uword xx3
-        uword total
+        uword iterations
 
         c64.SETTIM(0,0,0)
-        repeat 600 {
+
-            repeat 600 {
+        iterations = 0
+        repeat iterations {
+            repeat iterations {
+                xx1++
+                xx2++
+                xx3++
+            }
+        }
+
+        iterations = 600
+        repeat iterations {
+            repeat iterations {
                 xx1++
                 xx2++
                 xx3++
             }
         }
         uword time = c64.RDTIM16()
+        txt.print("time: ")
         txt.print_uw(time)
+        txt.print("\n$7e40? :\n")
+        txt.print_uwhex(xx1,true)
+        txt.nl()
+        txt.print_uwhex(xx2,true)
+        txt.nl()
+        txt.print_uwhex(xx3,true)
     }
 }