From ff69da3fa23f95680fbe5360344227284735c55b Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Sun, 5 Dec 2021 21:54:46 +0100
Subject: [PATCH] error when 'else' choice in when statemetn isn't the last
 one, also generate slightly better code for when statements

---
 .../compiler/target/cpu6502/codegen/AsmGen.kt  | 10 +++++-----
 .../astprocessing/StatementReorderer.kt        |  7 +++++++
 docs/source/todo.rst                           |  3 +--
 examples/test.p8                               | 18 +++++++++++-------
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/codeGeneration/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt b/codeGeneration/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
index 63c08605a..ce4fa873a 100644
--- a/codeGeneration/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
+++ b/codeGeneration/src/prog8/compiler/target/cpu6502/codegen/AsmGen.kt
@@ -1293,7 +1293,6 @@ $repeatLabel    lda  $counterVar
             if(choice.values==null) {
                 // the else choice
                 translate(choice.statements)
-                jmp(endLabel)
             } else {
                 choiceBlocks.add(choiceLabel to choice.statements)
                 for (cv in choice.values!!) {
@@ -1313,10 +1312,11 @@ $repeatLabel    lda  $counterVar
             }
         }
         jmp(endLabel)
-        for(choiceBlock in choiceBlocks) {
-            out(choiceBlock.first)
-            translate(choiceBlock.second)
-            jmp(endLabel)
+        for(choiceBlock in choiceBlocks.withIndex()) {
+            out(choiceBlock.value.first)
+            translate(choiceBlock.value.second)
+            if(choiceBlock.index<choiceBlocks.size-1)
+                jmp(endLabel)
         }
         out(endLabel)
     }
diff --git a/compiler/src/prog8/compiler/astprocessing/StatementReorderer.kt b/compiler/src/prog8/compiler/astprocessing/StatementReorderer.kt
index 06a918ece..9cbf22aa7 100644
--- a/compiler/src/prog8/compiler/astprocessing/StatementReorderer.kt
+++ b/compiler/src/prog8/compiler/astprocessing/StatementReorderer.kt
@@ -253,6 +253,13 @@ internal class StatementReorderer(val program: Program,
     }
 
     override fun after(whenStatement: WhenStatement, parent: Node): Iterable<IAstModification> {
+        val lastChoiceValues = whenStatement.choices.lastOrNull()?.values
+        if(lastChoiceValues?.isNotEmpty()==true) {
+            val elseChoice = whenStatement.choices.indexOfFirst { it.values==null || it.values?.isEmpty()==true }
+            if(elseChoice>=0)
+                errors.err("else choice must be the last one", whenStatement.choices[elseChoice].position)
+        }
+
         val choices = whenStatement.choiceValues(program).sortedBy {
             it.first?.first() ?: Int.MAX_VALUE
         }
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index d58f014c3..cf98e7348 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -3,8 +3,7 @@ TODO
 
 For next compiler release (7.5)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-- fix: when statements generate unneeded branches to choice_end?
-
+...
 
 
 Blocked by an official Commander-x16 v39 release
diff --git a/examples/test.p8 b/examples/test.p8
index 1796054c5..c7d33e8ea 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -4,14 +4,18 @@
 
 main {
     sub start() {
-        ubyte @shared bb = @(cx16.r0)
-    }
+        ubyte @shared bb
 
-    sub count() -> ubyte {
-        repeat {
-            %asm {{
-                rts
-            }}
+        when bb {
+            0,1,2 -> {
+                bb ++
+            }
+            3 -> {
+                bb--
+            }
+            else -> {
+                bb=0
+            }
         }
     }
 }