optimize: rewrite suitable when into on..goto

codeOptimizers/src/prog8/optimizer/StatementOptimizer.kt

@@ -501,6 +501,27 @@ class StatementOptimizer(private val program: Program,
             }
         }
 
+        if(whenStmt.betterAsOnGoto(program, options)) {
+            // rewrite when into a on..goto , which is faster and also smaller for ~5+ cases
+            var elseJump: Jump? = null
+            val jumps = mutableListOf<Pair<Int, Jump>>()
+            whenStmt.choices.forEach { choice ->
+                if(choice.values==null) {
+                    elseJump = choice.statements.statements.single() as Jump
+                } else {
+                    choice.values!!.forEach { value ->
+                        jumps.add(value.constValue(program)!!.number.toInt() to choice.statements.statements.single() as Jump)
+                    }
+                }
+            }
+
+            val jumpLabels = jumps.sortedBy { it.first }.map { it.second.target as IdentifierReference }
+
+            val elsePart = if(elseJump==null) null else AnonymousScope(mutableListOf(elseJump), elseJump.position)
+            val onGoto = OnGoto(false, whenStmt.condition, jumpLabels, elsePart, whenStmt.position)
+            return listOf(IAstModification.ReplaceNode(whenStmt, onGoto, parent))
+        }
+
         return noModifications
     }
 

compiler/src/prog8/compiler/astprocessing/AstChecker.kt

@@ -1669,26 +1669,8 @@ internal class AstChecker(private val program: Program,
         if(whenStmt.condition.constValue(program)!=null)
             errors.warn("when-value is a constant and will always result in the same choice", whenStmt.condition.position)
 
-
-        // a when that has only goto's and the values 0,1,2,3,4... is better written as a on..goto
-        if(whenStmt.choices.size>=3) {
-            if (whenStmt.condition.inferType(program).isBytes) {
-                if (whenStmt.choices.all { it.statements.statements.singleOrNull() is Jump }) {
-                    val values = whenStmt.choices.flatMap {
-                        it.values ?: mutableListOf()
-                    }.map {
-                        it.constValue(program)?.number?.toInt()
-                    }
-                    if(null !in values) {
-                        val sortedValues = values.filterNotNull().sorted()
-                        val range = IntRange(sortedValues.first(), sortedValues.last())
-                        if(range.toList() == sortedValues) {
-                            errors.info("when statement can be replaced with on..goto", whenStmt.position)
-                        }
-                    }
-                }
-            }
-        }
+        if(whenStmt.betterAsOnGoto(program, compilerOptions))
+            errors.info("when statement can be replaced with on..goto", whenStmt.position)
 
         super.visit(whenStmt)
     }

compilerAst/src/prog8/ast/statements/AstStatements.kt

@@ -1195,6 +1195,30 @@ class When(var condition: Expression,
     override fun accept(visitor: AstWalker, parent: Node) = visitor.visit(this, parent)
     override fun referencesIdentifier(nameInSource: List<String>): Boolean =
         condition.referencesIdentifier(nameInSource) || choices.any { it.referencesIdentifier(nameInSource) }
+
+    fun betterAsOnGoto(program: Program, compilerOptions: CompilationOptions): Boolean {
+        // a when that has only goto's and the values 0,1,2,3,4... is better written as a on..goto
+        val sizeLimit = if(compilerOptions.compTarget.cpu == CpuType.CPU65C02) 4 else 6
+        if(choices.size >= sizeLimit) {
+            if (condition.inferType(program).isBytes) {
+                if (choices.all { (it.statements.statements.singleOrNull() as? Jump)?.target is IdentifierReference }) {
+                    val values = choices.flatMap {
+                        it.values ?: mutableListOf()
+                    }.map {
+                        it.constValue(program)?.number?.toInt()
+                    }
+                    if(null !in values) {
+                        val sortedValues = values.filterNotNull().sorted()
+                        val range = IntRange(sortedValues.first(), sortedValues.last())
+                        if(range.toList() == sortedValues) {
+                            return true
+                        }
+                    }
+                }
+            }
+        }
+        return false
+    }
 }
 
 class WhenChoice(var values: MutableList<Expression>?,           // if null,  this is the 'else' part

examples/test.p8

@@ -16,12 +16,14 @@ main {
             1 -> goto second
             2 -> goto third
             3 -> goto fourth
+            4 -> goto third
             else -> goto other
         }
 
-        ; TODO 3+ options better as on:
+        ; 65c02: 3+ options better as on:
+        ; 6502: 5+ options better as on:
 
-        ;on value goto (first, second, third, fourth) else goto other
+        ;on value goto (first, second, third, fourth, third, fourth)
 
         sub first() {
             cx16.r0++