From c1aa5d4e47402cc4a48e8351e42a923d8365bb66 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Sun, 16 Jul 2023 12:10:46 +0200
Subject: [PATCH] IR: optimized when statement translation

---
 .../prog8/codegen/intermediate/IRCodeGen.kt   | 43 ++++++++-----------
 examples/test.p8                              | 26 ++++++++---
 2 files changed, 38 insertions(+), 31 deletions(-)

diff --git a/codeGenIntermediate/src/prog8/codegen/intermediate/IRCodeGen.kt b/codeGenIntermediate/src/prog8/codegen/intermediate/IRCodeGen.kt
index e865a599b..4d1b987f7 100644
--- a/codeGenIntermediate/src/prog8/codegen/intermediate/IRCodeGen.kt
+++ b/codeGenIntermediate/src/prog8/codegen/intermediate/IRCodeGen.kt
@@ -399,42 +399,35 @@ class IRCodeGen(
     }
 
     private fun translate(whenStmt: PtWhen): IRCodeChunks {
-        if(whenStmt.choices.children.isEmpty())
-            return emptyList()
         val result = mutableListOf<IRCodeChunkBase>()
         val valueDt = irType(whenStmt.value.type)
         val valueTr = expressionEval.translateExpression(whenStmt.value)
         addToResult(result, valueTr, valueTr.resultReg, -1)
-        val choices = whenStmt.choices.children.map {it as PtWhenChoice }
+
+        val choices = mutableListOf<Pair<String, PtWhenChoice>>()
         val endLabel = createLabelName()
-        for (choice in choices) {
+        whenStmt.choices.children.forEach {
+            val choice = it as PtWhenChoice
             if(choice.isElse) {
                 result += translateNode(choice.statements)
+                addInstr(result, IRInstruction(Opcode.JUMP, labelSymbol = endLabel), null)
             } else {
-                val skipLabel = createLabelName()
-                val values = choice.values.children.map {it as PtNumber}
-                if(values.size==1) {
-                    val chunk = IRCodeChunk(null, null)
-                    chunk += IRInstruction(Opcode.BNE, valueDt, reg1=valueTr.resultReg, immediate = values[0].number.toInt(), labelSymbol = skipLabel)
-                    result += chunk
-                    result += translateNode(choice.statements)
-                    if(choice.statements.children.last() !is PtReturn)
-                        addInstr(result, IRInstruction(Opcode.JUMP, labelSymbol = endLabel), null)
-                } else {
-                    val matchLabel = createLabelName()
-                    val chunk = IRCodeChunk(null, null)
-                    for (value in values) {
-                        chunk += IRInstruction(Opcode.BEQ, valueDt, reg1=valueTr.resultReg, immediate = value.number.toInt(), labelSymbol = matchLabel)
-                    }
-                    chunk += IRInstruction(Opcode.JUMP, labelSymbol = skipLabel)
-                    result += chunk
-                    result += labelFirstChunk(translateNode(choice.statements), matchLabel)
-                    if(choice.statements.children.last() !is PtReturn)
-                        addInstr(result, IRInstruction(Opcode.JUMP, labelSymbol = endLabel), null)
+                val choiceLabel = createLabelName()
+                choices.add(choiceLabel to choice)
+                choice.values.children.map { it as PtNumber }.sortedBy { it.number }.forEach { value ->
+                    addInstr(result, IRInstruction(Opcode.BEQ, valueDt, reg1=valueTr.resultReg, immediate = value.number.toInt(), labelSymbol = choiceLabel), null)
                 }
-                result += IRCodeChunk(skipLabel, null)
             }
         }
+        addInstr(result, IRInstruction(Opcode.JUMP, labelSymbol = endLabel), null)
+
+        choices.forEach { (label, choice) ->
+            result += labelFirstChunk(translateNode(choice.statements), label)
+            val lastStatement = choice.statements.children.last()
+            if(lastStatement !is PtReturn && lastStatement !is PtJump)
+                addInstr(result, IRInstruction(Opcode.JUMP, labelSymbol = endLabel), null)
+        }
+
         result += IRCodeChunk(endLabel, null)
         return result
     }
diff --git a/examples/test.p8 b/examples/test.p8
index 9d4d528e3..7af43942d 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -5,13 +5,27 @@ main
 {
     sub start()
     {
-        ubyte a=1
-        ubyte b=2
-        ubyte c=3
-        ubyte d=4
+        ubyte res
+        ubyte val=3
+        when val*10 {
+            1 -> res=1
+            10 -> res=10
+            20 -> res=20
+            30 -> res=30
+            40 -> res=40
+            else -> res=0
+        }
 
-        ubyte xx = (a*b)+(c*d)
-        xx++
+        txt.print_ub(res)   ; 30
+        txt.nl()
 
+        when val {
+            5,7,9,1,3 -> res=100
+            2,4,6,8 -> res=200
+            else -> res=0
+        }
+
+        txt.print_ub(res)   ; 100
+        ; 102 instructions, 31 registers, 122 steps
     }
 }