syntax check and optimization of 'when'

compiler/src/prog8/ast/processing/AstChecker.kt

@@ -912,6 +912,33 @@ internal class AstChecker(private val program: Program,
         return super.visit(arrayIndexedExpression)
     }
 
+    override fun visit(whenStatement: WhenStatement): IStatement {
+        val conditionType = whenStatement.condition.inferType(program)
+        if(conditionType !in IntegerDatatypes)
+            checkResult.add(SyntaxError("when condition must be an integer value", whenStatement.position))
+        val choiceValues = whenStatement.choiceValues(program)
+        val occurringValues = choiceValues.map {it.first}
+        val tally = choiceValues.associate { it.second to occurringValues.count { ov->it.first==ov} }
+        tally.filter { it.value>1 }.forEach {
+            checkResult.add(SyntaxError("choice value occurs multiple times", it.key.position))
+        }
+        return super.visit(whenStatement)
+    }
+
+    override fun visit(whenChoice: WhenChoice) {
+        if(whenChoice.value!=null) {
+            val constvalue = whenChoice.value.constValue(program)
+            if (constvalue == null)
+                checkResult.add(SyntaxError("value of a when choice must be a constant", whenChoice.position))
+            else if (constvalue.type !in IntegerDatatypes)
+                checkResult.add(SyntaxError("value of a when choice must be an integer", whenChoice.position))
+        } else {
+            if(whenChoice !== (whenChoice.parent as WhenStatement).choices.last())
+                checkResult.add(SyntaxError("else choice must be the last one", whenChoice.position))
+        }
+        super.visit(whenChoice)
+    }
+
     private fun checkFunctionOrLabelExists(target: IdentifierReference, statement: IStatement): IStatement? {
         val targetStatement = target.targetStatement(program.namespace)
         if(targetStatement is Label || targetStatement is Subroutine || targetStatement is BuiltinFunctionStatementPlaceholder)

compiler/src/prog8/ast/processing/StatementReorderer.kt

@@ -1,5 +1,6 @@
 package prog8.ast.processing
 
+import kotlin.comparisons.nullsLast
 import prog8.ast.*
 import prog8.ast.base.DataType
 import prog8.ast.base.FatalAstException
@@ -22,6 +23,7 @@ internal class StatementReorderer(private val program: Program): IAstModifyingVi
     //
     // - the 'start' subroutine in the 'main' block will be moved to the top immediately following the directives.
     // - all other subroutines will be moved to the end of their block.
+    // - sorts the choices in when statement.
     //
     // Also, makes sure any value assignments get the proper type casts if needed to cast them into the target variable's type.
     // (this includes function call arguments)
@@ -273,4 +275,11 @@ internal class StatementReorderer(private val program: Program): IAstModifyingVi
         }
         return super.visit(typecast)
     }
+
+    override fun visit(whenStatement: WhenStatement): IStatement {
+        // sort the choices in low-to-high value order
+        whenStatement.choices
+                .sortWith(compareBy<WhenChoice, Int?>(nullsLast(), {it.value?.constValue(program)?.asIntegerValue}))
+        return super.visit(whenStatement)
+    }
 }

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

@@ -6,6 +6,8 @@ import prog8.ast.expressions.*
 import prog8.ast.processing.IAstModifyingVisitor
 import prog8.ast.processing.IAstVisitor
 import prog8.compiler.HeapValues
+import kotlin.math.max
+import kotlin.math.min
 
 
 class BuiltinFunctionStatementPlaceholder(val name: String, override val position: Position) : IStatement {
@@ -677,6 +679,18 @@ class WhenStatement(val condition: IExpression,
         choices.forEach { it.linkParents(this) }
     }
 
+    fun choiceValues(program: Program): List<Pair<Int?, WhenChoice>> {
+        // only gives sensible results when the choices are all valid (constant integers)
+        return choices
+                .map {
+                    val cv = it.value?.constValue(program)
+                    if(cv==null)
+                        null to it
+                    else
+                        cv.asNumericValue!!.toInt() to it
+                    }
+    }
+
     override fun accept(visitor: IAstVisitor) = visitor.visit(this)
     override fun accept(visitor: IAstModifyingVisitor) = visitor.visit(this)
 }
@@ -692,6 +706,10 @@ class WhenChoice(val value: IExpression?,           // if null,  this is the 'el
         this.parent = parent
     }
 
+    override fun toString(): String {
+        return "Choice($value at $position)"
+    }
+
     fun accept(visitor: IAstVisitor) = visitor.visit(this)
     fun accept(visitor: IAstModifyingVisitor) = visitor.visit(this)
 }

compiler/src/prog8/compiler/Compiler.kt

@@ -212,6 +212,7 @@ internal class Compiler(private val program: Program) {
                 is Subroutine -> translate(stmt)
                 is NopStatement -> {}
                 is InlineAssembly -> translate(stmt)
+                is WhenStatement -> translate(stmt)
                 else -> TODO("translate statement $stmt to stackvm")
             }
         }
@@ -2080,6 +2081,30 @@ internal class Compiler(private val program: Program) {
             throw CompilerException("cannot take memory pointer $addrof")
     }
 
+    private fun translate(whenstmt: WhenStatement) {
+        val conditionDt = whenstmt.condition.inferType(program)
+        translate(whenstmt.condition)
+        if(whenstmt.choices.isEmpty()) {
+            when(conditionDt) {
+                in ByteDatatypes -> prog.instr(Opcode.DISCARD_BYTE)
+                in WordDatatypes -> prog.instr(Opcode.DISCARD_WORD)
+                else -> throw CompilerException("when condition must be integer")
+            }
+            return
+        }
+
+        // TODO compare
+
+        for(choice in whenstmt.choiceValues(program)) {
+            if(choice.first==null) {
+                // the else clause
+                translate(choice.second.statements)
+            }
+        }
+
+        TODO("whenstmt $whenstmt with choice values ${whenstmt.choiceValues(program).map{it.first}}")
+    }
+
     private fun translateAsmInclude(args: List<DirectiveArg>, source: Path) {
         val scopeprefix = if(args[1].str!!.isNotBlank()) "${args[1].str}\t.proc\n" else ""
         val scopeprefixEnd = if(args[1].str!!.isNotBlank()) "\t.pend\n" else ""

compiler/src/prog8/optimizer/StatementOptimizer.kt

@@ -435,6 +435,15 @@ internal class StatementOptimizer(private val program: Program, private val opti
         return nopStatement
     }
 
+    override fun visit(whenStatement: WhenStatement): IStatement {
+        val choices = whenStatement.choices.toList()
+        for(choice in choices) {
+            if(choice.statements.containsNoCodeNorVars())
+                whenStatement.choices.remove(choice)
+        }
+        return super.visit(whenStatement)
+    }
+
     private fun hasContinueOrBreak(scope: INameScope): Boolean {
 
         class Searcher: IAstModifyingVisitor

examples/test.p8

@@ -5,20 +5,30 @@
 ~ main {
 
     sub start() {
-        float fl
-        ubyte ub
+        A=10
+        Y=22
 
-        when A+Y {
-            fl -> A=2
-            fl+3.3 -> A=3
-            ub -> A=4
-            ub+2 -> A=5
-            4 -> {
-                Y=7
+        when 4+A+Y {
+            10 -> {
+                c64scr.print("ten")
+            }
+            5 -> c64scr.print("five")
+            30 -> c64scr.print("thirty")
+            99 -> c64scr.print("nn")
+            55 -> {
+                ; should be optimized away
+            }
+            56 -> {
+                ; should be optimized away
+            }
+            57 -> {
+                ; should be optimized away
+            }
+            else -> {
+                c64scr.print("!??!\n")
+                c64scr.print("!??!!??!\n")
+                c64scr.print("!??!!??!!?!\n")
             }
-            else -> A=99
-            5 -> Y=5        ; @todo error; else already seen
         }
     }
-
 }