From 2764d235a9a3bb33f2c05a98cd13f3fbab5db8d1 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Sat, 25 Nov 2023 21:09:19 +0100
Subject: [PATCH] optimizing for x in 0 to something

---
 .../src/prog8/optimizer/StatementOptimizer.kt | 46 ++++++++++++++++
 docs/source/todo.rst                          |  2 -
 examples/test.p8                              | 54 ++++++++++++++++++-
 3 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt b/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt
index 508f670e3..1949a09aa 100644
--- a/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt
+++ b/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt
@@ -168,6 +168,52 @@ class StatementOptimizer(private val program: Program,
             }
         }
 
+        val loopvarDt = forLoop.loopVarDt(program)
+        if(loopvarDt.istype(DataType.UWORD) || loopvarDt.istype(DataType.UBYTE)) {
+            if (range != null && range.from.constValue(program)?.number == 0.0 && range.step.constValue(program)?.number==1.0) {
+                val toBinExpr = range.to as? BinaryExpression
+                if(toBinExpr!=null && toBinExpr.operator=="-" && toBinExpr.right.constValue(program)?.number==1.0) {
+                    // FOR var IN 0 TO X-1 .... ---> var=0, DO {... , var++} UNTIL var==X
+                    val pos = forLoop.position
+                    val condition = BinaryExpression(forLoop.loopVar.copy(), "==", toBinExpr.left, pos)
+                    val incOne = PostIncrDecr(AssignTarget(forLoop.loopVar.copy(), null, null, pos), "++", pos)
+                    forLoop.body.statements.add(incOne)
+                    val replacement = AnonymousScope(mutableListOf(
+                        Assignment(AssignTarget(forLoop.loopVar.copy(), null, null, pos),
+                            NumericLiteral.optimalNumeric(0.0, pos),
+                            AssignmentOrigin.OPTIMIZER, pos),
+                        UntilLoop(forLoop.body, condition, pos)
+                    ), pos)
+                    return listOf(IAstModification.ReplaceNode(forLoop, replacement, parent))
+                }
+
+                if(options.compTarget.name!=VMTarget.NAME) {
+                    // this optimization is not effective for the VM target.
+                    val toConst = range.to.constValue(program)
+                    if (toConst == null) {
+                        // FOR var in 0 TO X ... --->  var=0, REPEAT { ... , IF var==X break , var++ }
+                        val pos = forLoop.position
+                        val incOne = PostIncrDecr(AssignTarget(forLoop.loopVar.copy(), null, null, pos), "++", pos)
+                        val breakCondition = IfElse(
+                            BinaryExpression(forLoop.loopVar, "==", range.to, pos),
+                            AnonymousScope(mutableListOf(Break(pos)), pos),
+                            AnonymousScope(mutableListOf(), pos),
+                            pos
+                        )
+                        forLoop.body.statements.add(breakCondition)
+                        forLoop.body.statements.add(incOne)
+                        val replacement = AnonymousScope(mutableListOf(
+                            Assignment(AssignTarget(forLoop.loopVar.copy(), null, null, pos),
+                                NumericLiteral.optimalNumeric(0.0, pos),
+                                AssignmentOrigin.OPTIMIZER, pos),
+                            RepeatLoop(null, forLoop.body, pos)
+                        ), pos)
+                        return listOf(IAstModification.ReplaceNode(forLoop, replacement, parent))
+                    }
+                }
+            }
+        }
+
         return noModifications
     }
 
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index 04bdb6a2c..b79ffce3f 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -2,8 +2,6 @@
 TODO
 ====
 
-- optimize 6502 codegen: "for 0 to end" into repeat loops. (only for ubyte and uword, 6502 cpu)
-
 - [on branch: shortcircuit] investigate McCarthy evaluation again? this may also reduce code size perhaps for things like if a>4 or a<2 ....
 
 ...
diff --git a/examples/test.p8 b/examples/test.p8
index 4806dbb92..a300eb315 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -1,13 +1,31 @@
 %import textio
 %zeropage basicsafe
 
+; $3d5
+
 main {
     ubyte counter
     uword wcounter
-    ubyte end=10
-    uword wend=10
+    ubyte end
+    uword wend
 
     sub start() {
+        end=10
+        for cx16.r2L in 0 to end-1 {
+            txt.print_ub(cx16.r2L)
+            txt.spc()
+        }
+        txt.nl()
+
+        cx16.r2L=0
+    labeltje:
+            txt.print_ub(cx16.r2L)
+            txt.spc()
+            cx16.r2L++
+            if cx16.r2L!=end
+                goto labeltje
+        txt.nl()
+
         cx16.r0=0
         forloops()
         txt.print_uw(cx16.r0)
@@ -24,19 +42,31 @@ main {
         for counter in 0 to end {
             cx16.r0++
         }
+        for counter in 0 to end-1 {
+            cx16.r0++
+        }
         end=0
         for counter in 0 to end {
             cx16.r0++
         }
+        for counter in 0 to end-1 {
+            cx16.r0++
+        }
         end=255
         for counter in 0 to end {
             cx16.r0++
         }
+        for counter in 0 to end-1 {
+            cx16.r0++
+        }
 
         wend=1000
         for wcounter in 0 to wend {
             cx16.r0++
         }
+        for wcounter in 0 to wend-1 {
+            cx16.r0++
+        }
     }
 
     sub untilloops() {
@@ -49,6 +79,11 @@ main {
                 break
             counter++
         }
+        counter = 0
+        do {
+            cx16.r0++
+            counter++
+        } until counter==end
 
         end=0
         counter = 0
@@ -58,6 +93,11 @@ main {
                 break
             counter++
         }
+        counter = 0
+        do {
+            cx16.r0++
+            counter++
+        } until counter==end
 
         counter = 0
         end=255
@@ -67,6 +107,11 @@ main {
                 break
             counter++
         }
+        counter = 0
+        do {
+            cx16.r0++
+            counter++
+        } until counter==end
 
         wcounter = 0
         wend=1000
@@ -76,5 +121,10 @@ main {
                 break
             wcounter++
         }
+        wcounter = 0
+        do {
+            cx16.r0++
+            wcounter++
+        } until wcounter==wend
     }
 }