attempt 2 at optimizing repeats

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

@@ -1008,56 +1008,67 @@ internal class AsmGen(private val program: Program,
     }
 
     private fun repeatWordCountInAY(constIterations: Int?, repeatLabel: String, endLabel: String, body: AnonymousScope) {
+        // note: A/Y must have been loaded with the number of iterations!
         if(constIterations==0)
             return
-        // note: A/Y must have been loaded with the number of iterations already!
+        if(constIterations==null) {
+            out("""
+                cmp  #0
+                bne  +
+                cpy  #0
+                beq  $endLabel   ; skip loop if zero iters
++""")
+        }
         val counterVar = makeLabel("repeatcounter")
         out("""
-            cmp  #0
+                sta  $counterVar
-            bne  +
+                sty  $counterVar+1
-            cpy  #0
+$repeatLabel    lda  $counterVar
-            beq  $endLabel      ; skip if 0 iterations
+                bne  +
-+           sta  $counterVar
+                lda  $counterVar+1
-            sty  $counterVar+1
+                beq  $endLabel
-$repeatLabel
++               lda  $counterVar
-            lda  $counterVar
+                bne  +
-            bne  +
+                dec  $counterVar+1
-            lda  $counterVar+1
++               dec  $counterVar
-            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
+            // allocate count var on ZP  TODO can be shared with countervars from other subroutines
             val zpAddr = zeropage.allocate(counterVar, DataType.UWORD, body.position, errors)
-            out("""$counterVar = $zpAddr  ; auto zp UWORD""")
+            out("$counterVar = $zpAddr  ; auto zp UWORD")
         } else {
-            out("""
+            out("$counterVar    .word  0")
-$counterVar    .word  0""")
         }
 
         out(endLabel)
     }
 
     private fun repeatByteCountInA(constIterations: Int?, repeatLabel: String, endLabel: String, body: AnonymousScope) {
+        // note: A must have been loaded with the number of iterations!
         if(constIterations==0)
             return
-        // note: A must have been loaded with the number of iterations already!
-        val counterVar = makeLabel("repeatcounter")
         if(constIterations==null)
-            out("  beq  $endLabel")
+            out("  beq  $endLabel   ; skip loop if zero iters")
+        val counterVar = makeLabel("repeatcounter")
         out("  sta  $counterVar")
         out(repeatLabel)
         translate(body)
         out("""
              dec  $counterVar
              bne  $repeatLabel
-             beq  $endLabel  
+             beq  $endLabel""")
-$counterVar    .byte  0""")
+
+        if(constIterations!=null && constIterations>=16 && zeropage.available() > 0) {
+            // allocate count var on ZP  TODO can be shared with countervars from other subroutines
+            val zpAddr = zeropage.allocate(counterVar, DataType.UBYTE, body.position, errors)
+            out("$counterVar = $zpAddr  ; auto zp UBYTE")
+        } else {
+            out("$counterVar    .byte  0")
+        }
+
         out(endLabel)
     }
 

examples/test.p8

@@ -3,25 +3,81 @@
 
 main {
 
+    sub iter(uword iterations) -> uword {
+        uword total = 0
+        repeat iterations {
+            repeat iterations {
+                total++
+            }
+        }
+
+        return total
+
+    }
+
+    sub iterb(ubyte iterations) -> uword {
+        uword total = 0
+        repeat iterations {
+            repeat iterations {
+                total++
+            }
+        }
+
+        return total
+    }
+
     sub start() {
         uword xx1
         uword xx2
         uword xx3
         uword iterations
 
+        xx1=0
+        repeat 99 {
+            xx1++
+        }
+        txt.print_uw(xx1)   ;99
+        txt.nl()
+
+        xx1 = iterb(10)   ; 100
+        txt.print_uw(xx1)
+        txt.nl()
+        xx1 = iterb(1)   ; 1
+        txt.print_uw(xx1)
+        txt.nl()
+        xx1 = iterb(0)   ; 0
+        txt.print_uw(xx1)
+        txt.nl()
+        txt.nl()
+
+        xx1 = iter(0)
+        txt.print_uw(xx1)   ; 0
+        txt.nl()
+        xx1 = iter(10)
+        txt.print_uw(xx1)   ; 100
+        txt.nl()
+        xx1 = iter(16)
+        txt.print_uw(xx1)   ; 256
+        txt.nl()
+        xx1 = iter(20)
+        txt.print_uw(xx1)   ; 400
+        txt.nl()
+        xx1 = iter(200)
+        txt.print_uw(xx1)   ; 4000
+        txt.nl()
+        xx1 = iter(600)
+        txt.print_uw(xx1)   ; 32320
+        txt.nl()
+        txt.nl()
+
+
         c64.SETTIM(0,0,0)
 
-        iterations = 0
+        xx1=0
-        repeat iterations {
+        xx2=0
-            repeat iterations {
+        xx3=0
-                xx1++
-                xx2++
-                xx3++
-            }
-        }
-
         iterations = 600
-        repeat iterations {
+        repeat 600 {
             repeat iterations {
                 xx1++
                 xx2++