From 1152191f4814556b4a3d42d5d9364228401826b9 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Wed, 15 Mar 2023 21:11:37 +0100
Subject: [PATCH] add optimization: replace simple for loops by repeat loop

---
 .../src/prog8/optimizer/StatementOptimizer.kt | 10 ++++++++++
 .../src/prog8/optimizer/UnusedCodeRemover.kt  |  4 ++--
 .../src/prog8/ast/statements/AstStatements.kt | 20 +++++++++++++++++++
 docs/source/todo.rst                          |  3 ---
 4 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt b/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt
index e4958e063..37ea7e2f6 100644
--- a/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt
+++ b/codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt
@@ -166,6 +166,16 @@ class StatementOptimizer(private val program: Program,
             }
         }
 
+        val iterationCount = forLoop.constIterationCount(program)
+        if(iterationCount!=null) {
+            val loopName = forLoop.loopVar.nameInSource
+            if(!forLoop.iterable.referencesIdentifier(loopName) && !forLoop.body.referencesIdentifier(loopName)) {
+                errors.warn("for loop can be replaced with repeat loop, possibly also remove the loop variable", forLoop.position)
+                val repeat = RepeatLoop(NumericLiteral.optimalNumeric(iterationCount, forLoop.position), forLoop.body, forLoop.position)
+                return listOf(IAstModification.ReplaceNode(forLoop, repeat, parent))
+            }
+        }
+
         return noModifications
     }
 
diff --git a/codeOptimizers/src/prog8/optimizer/UnusedCodeRemover.kt b/codeOptimizers/src/prog8/optimizer/UnusedCodeRemover.kt
index 3bb14cee4..d42dadf84 100644
--- a/codeOptimizers/src/prog8/optimizer/UnusedCodeRemover.kt
+++ b/codeOptimizers/src/prog8/optimizer/UnusedCodeRemover.kt
@@ -125,8 +125,8 @@ class UnusedCodeRemover(private val program: Program,
                     if(usages.size==1) {
                         val singleUse = usages[0].parent
                         if(singleUse is AssignTarget) {
-                            val assignment = singleUse.parent as Assignment
-                            if(assignment.origin==AssignmentOrigin.VARINIT) {
+                            val assignment = singleUse.parent as? Assignment
+                            if(assignment!=null && assignment.origin==AssignmentOrigin.VARINIT) {
                                 if(assignment.value.isSimple) {
                                     // remove the vardecl
                                     if(!decl.definingModule.isLibrary)
diff --git a/compilerAst/src/prog8/ast/statements/AstStatements.kt b/compilerAst/src/prog8/ast/statements/AstStatements.kt
index 6fc26ca84..331d0134c 100644
--- a/compilerAst/src/prog8/ast/statements/AstStatements.kt
+++ b/compilerAst/src/prog8/ast/statements/AstStatements.kt
@@ -863,6 +863,26 @@ class ForLoop(var loopVar: IdentifierReference,
                 body.referencesIdentifier(nameInSource)
 
     fun loopVarDt(program: Program) = loopVar.inferType(program)
+
+    fun constIterationCount(program: Program): Int? {
+        return when (val iter = iterable) {
+            is IdentifierReference -> {
+                val target = iter.targetVarDecl(program)!!
+                if (target.isArray)
+                    target.arraysize!!.constIndex()
+                else if (target.value is StringLiteral)
+                    (target.value as StringLiteral).value.length
+                else if (target.value is ArrayLiteral)
+                    (target.value as ArrayLiteral).value.size
+                else null
+            }
+            is ArrayLiteral -> iter.value.size
+            is RangeExpression -> iter.size()
+            is StringLiteral -> iter.value.length
+            else -> null
+        }
+    }
+
 }
 
 class WhileLoop(var condition: Expression,
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index e96a98b81..74ce5aac5 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -3,9 +3,6 @@ TODO
 
 For next minor release
 ^^^^^^^^^^^^^^^^^^^^^^
-- when a loopvariable of a forloop is *only* referenced in the for loop as loopvariable, and the number of iterations is known, replace the loop by a repeatloop (and remove the vardecl)
-  but we have no efficient way right now to see if the body references a variable. (what about ReferencesIdentifier?)
-
 ...