From 51e6bf0d45572ab8504caa55acf20e3f90f09bf1 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Sun, 28 Feb 2021 17:34:18 +0100
Subject: [PATCH] optimized repeat loop for word counts

---
 .../compiler/target/cpu6502/codegen/AsmGen.kt | 34 +++++++-------
 examples/test.p8                              | 46 +++++++++++++++----
 2 files changed, 54 insertions(+), 26 deletions(-)

diff --git a/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt b/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
index 715accc23..1fd39ad45 100644
--- a/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
+++ b/compiler/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
@@ -1008,28 +1008,30 @@ 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 if 0 iterations
++""")
+        }
         val counterVar = makeLabel("repeatcounter")
         out("""
-            cmp  #0
-            bne  +
-            cpy  #0
-            beq  $endLabel      ; skip if 0 iterations
-+           sta  $counterVar
+            sta  $counterVar
+            iny
             sty  $counterVar+1
-$repeatLabel
-            lda  $counterVar
-            bne  +
-            lda  $counterVar+1
-            beq  $endLabel
-+           lda  $counterVar
-            bne  +
-            dec  $counterVar+1
-+           dec  $counterVar""")
+$repeatLabel""")
         translate(body)
-        jmp(repeatLabel)
+        out("""
+            dec  $counterVar
+            bne  $repeatLabel
+            dec  $counterVar+1
+            bne  $repeatLabel
+            beq  $endLabel""")
 
         if(constIterations!=null && constIterations>=16 && zeropage.available() > 1) {
             // allocate count var on ZP
diff --git a/examples/test.p8 b/examples/test.p8
index ecab6a0b3..63bc95b70 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -3,25 +3,51 @@
 
 main {
 
+    sub iter(uword iterations) -> uword {
+        uword total = 0
+        repeat iterations {
+            repeat iterations {
+                total++
+            }
+        }
+
+        return total
+
+    }
+
     sub start() {
         uword xx1
         uword xx2
         uword xx3
         uword iterations
 
+        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()
+
+
         c64.SETTIM(0,0,0)
 
-        iterations = 0
-        repeat iterations {
-            repeat iterations {
-                xx1++
-                xx2++
-                xx3++
-            }
-        }
-
+        xx1=0
+        xx2=0
+        xx3=0
         iterations = 600
-        repeat iterations {
+        repeat 600 {
             repeat iterations {
                 xx1++
                 xx2++