automatically convert multi-compare expression (if X==1 or X==2..) to contaiment check if X in [1,2,..]

compiler/src/prog8/compiler/astprocessing/VariousCleanups.kt

@@ -91,13 +91,47 @@ internal class VariousCleanups(val program: Program, val errors: IErrorReporter,
     }
 
     override fun before(expr: BinaryExpression, parent: Node): Iterable<IAstModification> {
+        // try to replace a multi-comparison expression (if x==1 or x==2 or x==3 ... ) by a simple containment check.
+        // but only if the containment check is the top-level expression.
+        if(parent is BinaryExpression)
+            return noModifications
         if(expr.operator == "or") {
-            val leftBinExpr = expr.left as? BinaryExpression
+            val leftBinExpr1 = expr.left as? BinaryExpression
-            val rightBinExpr = expr.right as? BinaryExpression
+            val rightBinExpr1 = expr.right as? BinaryExpression
-            if(leftBinExpr!=null && leftBinExpr.operator=="==" && rightBinExpr!=null && rightBinExpr.operator=="==") {
+
-                if(leftBinExpr.right is NumericLiteral && rightBinExpr.right is NumericLiteral) {
+            if(rightBinExpr1?.operator=="==" && rightBinExpr1.right is NumericLiteral && leftBinExpr1!=null) {
-                    if(leftBinExpr.left isSameAs rightBinExpr.left)
+                val needle = rightBinExpr1.left
-                        errors.warn("consider using 'in' or 'when' to test for multiple values", expr.position)
+                val values = mutableListOf(rightBinExpr1.right as NumericLiteral)
+
+                fun isMultiComparisonRecurse(expr: BinaryExpression): Boolean {
+                    if(expr.operator=="==") {
+                        if(expr.right is NumericLiteral && expr.left isSameAs needle) {
+                            values.add(expr.right as NumericLiteral)
+                            return true
+                        }
+                        return false
+                    }
+                    if(expr.operator!="or")
+                        return false
+                    val leftBinExpr = expr.left as? BinaryExpression
+                    val rightBinExpr = expr.right as? BinaryExpression
+                    if(leftBinExpr==null || rightBinExpr==null || rightBinExpr.right !is NumericLiteral || !rightBinExpr.left.isSameAs(needle))
+                        return false
+                    if(rightBinExpr.operator=="==")
+                        values.add(rightBinExpr.right as NumericLiteral)
+                    else
+                        return false
+                    return isMultiComparisonRecurse(leftBinExpr)
+                }
+
+                if(isMultiComparisonRecurse(leftBinExpr1)) {
+                    // replace it!
+                    val valueCopies = values.sortedBy { it.number }.map { it.copy() }
+                    val elementType = needle.inferType(program).getOrElse { throw FatalAstException("invalid needle dt") }
+                    val arrayType = ElementToArrayTypes.getValue(elementType)
+                    val valuesArray = ArrayLiteral(InferredTypes.InferredType.known(arrayType), valueCopies.toTypedArray(), expr.position)
+                    val containment = ContainmentCheck(needle, valuesArray, expr.position)
+                    return listOf(IAstModification.ReplaceNode(expr, containment, parent))
                 }
             }
         }

compiler/test/TestOptimization.kt

@@ -3,10 +3,12 @@ package prog8tests
 import io.kotest.assertions.fail
 import io.kotest.assertions.withClue
 import io.kotest.core.spec.style.FunSpec
+import io.kotest.matchers.collections.shouldStartWith
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.shouldNotBe
 import io.kotest.matchers.string.shouldContain
 import io.kotest.matchers.string.shouldNotBeBlank
+import io.kotest.matchers.string.shouldStartWith
 import io.kotest.matchers.types.instanceOf
 import io.kotest.matchers.types.shouldBeSameInstanceAs
 import prog8.ast.ParentSentinel
@@ -169,7 +171,6 @@ class TestOptimization: FunSpec({
                 }
             }"""
         val result = compileText(C64Target(), true, source, writeAssembly = false)!!
-        printProgram(result.program)
         // expected:
 //        word llw
 //        llw = 300
@@ -424,7 +425,6 @@ class TestOptimization: FunSpec({
                 }
             }"""
         val result = compileText(C64Target(), optimize=true, src, writeAssembly=false)!!
-        printProgram(result.program)
         result.program.entrypoint.statements.size shouldBe 3
         val ifstmt = result.program.entrypoint.statements[0] as IfElse
         ifstmt.truepart.statements.size shouldBe 1
@@ -609,7 +609,6 @@ class TestOptimization: FunSpec({
         }
         """
         val result = compileText(C64Target(), optimize=true, src, writeAssembly=false)!!
-        printProgram(result.program)
         /* expected result:
         uword yy
         yy = 20
@@ -649,8 +648,6 @@ class TestOptimization: FunSpec({
         yy = 0
         xx += 10
          */
-        printProgram(result.program)
-
         val stmts = result.program.entrypoint.statements
         stmts.size shouldBe 7
         stmts.filterIsInstance<VarDecl>().size shouldBe 2
@@ -665,4 +662,101 @@ class TestOptimization: FunSpec({
         (xxValue.left as? IdentifierReference)?.nameInSource shouldBe listOf("xx")
         xxValue.right shouldBe NumericLiteral(DataType.UBYTE, 10.0, Position.DUMMY)
     }
+
+    test("multi-comparison replaced by containment check") {
+        val src="""
+        main {
+            sub start() {
+                ubyte source=99
+                ubyte thingy=42
+        
+                if source==3 or source==4 or source==99 or source==1
+                    thingy++
+            }
+        }"""
+        val result = compileText(C64Target(), optimize=true, src, writeAssembly=false)!!
+        /*
+        expected result:
+        ubyte[] auto_heap_var = [1,4,99,3]
+        ubyte source
+        source = 99
+        ubyte thingy
+        thingy = 42
+        if source in auto_heap_var
+            thingy++
+         */
+        val stmts = result.program.entrypoint.statements
+        stmts.size shouldBe 6
+        val ifStmt = stmts[5] as IfElse
+        val containment = ifStmt.condition as ContainmentCheck
+        (containment.element as IdentifierReference).nameInSource shouldBe listOf("source")
+        (containment.iterable as IdentifierReference).nameInSource.single() shouldStartWith "auto_heap_value"
+        val arrayDecl = stmts[0] as VarDecl
+        arrayDecl.isArray shouldBe true
+        arrayDecl.arraysize?.constIndex() shouldBe 4
+        val arrayValue = arrayDecl.value as ArrayLiteral
+        arrayValue.type shouldBe InferredTypes.InferredType.known(DataType.ARRAY_UB)
+        arrayValue.value shouldBe listOf(
+            NumericLiteral.optimalInteger(1, Position.DUMMY),
+            NumericLiteral.optimalInteger(3, Position.DUMMY),
+            NumericLiteral.optimalInteger(4, Position.DUMMY),
+            NumericLiteral.optimalInteger(99, Position.DUMMY))
+    }
+
+    test("invalid multi-comparison (not all equals) not replaced") {
+        val src="""
+        main {
+            sub start() {
+                ubyte source=99
+                ubyte thingy=42
+        
+                if source==3 or source==4 or source!=99 or source==1
+                    thingy++
+            }
+        }"""
+        val result = compileText(C64Target(), optimize=true, src, writeAssembly=false)!!
+        printProgram(result.program)
+        val stmts = result.program.entrypoint.statements
+        stmts.size shouldBe 5
+        val ifStmt = stmts[4] as IfElse
+        ifStmt.condition shouldBe instanceOf<BinaryExpression>()
+    }
+
+    test("invalid multi-comparison (not all same needle) not replaced") {
+        val src="""
+        main {
+            sub start() {
+                ubyte source=99
+                ubyte thingy=42
+        
+                if source==3 or source==4 or thingy==99 or source==1
+                    thingy++
+            }
+        }"""
+        val result = compileText(C64Target(), optimize=true, src, writeAssembly=false)!!
+        printProgram(result.program)
+        val stmts = result.program.entrypoint.statements
+        stmts.size shouldBe 5
+        val ifStmt = stmts[4] as IfElse
+        ifStmt.condition shouldBe instanceOf<BinaryExpression>()
+    }
+
+    test("invalid multi-comparison (not all or) not replaced") {
+        val src="""
+        main {
+            sub start() {
+                ubyte source=99
+                ubyte thingy=42
+        
+                if source==3 or source==4 and source==99 or source==1
+                    thingy++
+            }
+        }"""
+        val result = compileText(C64Target(), optimize=true, src, writeAssembly=false)!!
+        printProgram(result.program)
+        val stmts = result.program.entrypoint.statements
+        stmts.size shouldBe 5
+        val ifStmt = stmts[4] as IfElse
+        ifStmt.condition shouldBe instanceOf<BinaryExpression>()
+    }
 })

docs/source/todo.rst

@@ -3,12 +3,15 @@ TODO
 
 For next release
 ^^^^^^^^^^^^^^^^
-- pipe operator: allow non-unary function calls in the pipe that specify the other argument(s) in the calls.
+- fix compiler crash when comparing strings
 - make it possible to inline non-asmsub routines that just contain a single statement (return, functioncall, assignment)
-  but this requires all identifiers in the inlined expression to be changed to fully scoped names.
+  Only if the arguments are simple expressions, and the inlined subroutine cannot contain further nested subroutines!
-  If we can do that why not perhaps also able to inline multi-line subroutines? Why would it be limited to just 1 line? Maybe to protect against code size bloat.
+  This requires all identifiers in the inlined expression to be changed to fully scoped names (because their scope changes).
-  Inlined subroutines cannot contain further nested subroutines!
+  If we can do that why not perhaps also able to inline multi-line subroutines?
+  Why would it be limited to just 1 line? Maybe to protect against code size bloat.
   Once this works, look for library subroutines that should be inlined.
+- vm: add way more instructions operating directly on memory instead of only registers
+- add McCarthy evaluation to shortcircuit and/or expressions. First do ifs by splitting them up? Then do expressions that compute a value?
 
 ...
 
@@ -24,7 +27,6 @@ Future Things and Ideas
 ^^^^^^^^^^^^^^^^^^^^^^^
 Compiler:
 
-- vm: add way more instructions operating directly on memory instead of only registers
 - vm: codeGen: various TODOs to tweak code
 - vm: somehow deal with asmsubs otherwise the vm IR can't fully encode all of prog8
 - vm: make registers typed? so that it's immediately obvious what type they represent. Much like regular variables in memory.
@@ -40,7 +42,6 @@ Compiler:
 - make everything an expression? (get rid of Statements. Statements are expressions with void return types?).
 - simplifyConditionalExpression() should not split expression if it still results in stack-based evaluation, but how does it know?
 - simplifyConditionalExpression() sometimes introduces needless assignment to r9 tempvar (what scenarios?)
-- add McCarthy evaluation to shortcircuit and/or expressions. First do ifs by splitting them up? Then do expressions that compute a value?
 - make it possible to use cpu opcodes such as 'nop' as variable names by prefixing all asm vars with something such as ``p8v_``? Or not worth it (most 3 letter opcodes as variables are nonsensical anyway)
   then we can get rid of the instruction lists in the machinedefinitions as well?
 - [problematic due to using 64tass:] add a compiler option to not remove unused subroutines. this allows for building library programs. But this won't work with 64tass's .proc ...
@@ -58,12 +59,12 @@ Libraries:
 - optimize several inner loops in gfx2 even further?
 - add modes 2 and 3 to gfx2 (lowres 4 color and 16 color)?
 - add a flood fill routine to gfx2?
-- diskio: use cx16 MACPTR() to load stuff faster? (see it's use in X16edit to fast load blocks)
+- diskio: use cx16 MACPTR() to load stuff faster? (see its use in X16edit to fast load blocks)
   note that it might fail on non sdcard files so have to make graceful degradation
-- add a diskio.f_seek() routine for the Cx16 that uses its seek dos api? (only if that's stable)
 
 Expressions:
 
+- pipe operator: (targets other than 'Virtual'): allow non-unary function calls in the pipe that specify the other argument(s) in the calls.
 - rethink the whole "isAugmentable" business.  Because the way this is determined, should always also be exactly mirrorred in the AugmentableAssignmentAsmGen or you'll get a crash at code gen time.
   note: new ast PtAssignment already has no knowledge about this anymore.
 - can we get rid of pieces of asmgen.AssignmentAsmGen by just reusing the AugmentableAssignment ? generated code should not suffer
@@ -84,5 +85,4 @@ Optimizations:
 - AssignmentAsmGen.assignExpression() -> better code gen for assigning boolean comparison expressions
 - when a for loop's loopvariable isn't referenced in the body, and the iterations are known, replace the loop by a repeatloop
   but we have no efficient way right now to see if the body references a variable.
-- automatically convert if statements that test for multiple values (if X==1 or X==2..) to if X in [1,2,..] statements, instead of just a warning.
 - introduce byte-index operator to avoid index multiplications in loops over arrays? see github issue #4

examples/test.p8

@@ -7,23 +7,20 @@
 
 main {
 
-    sub add(ubyte first, ubyte second) -> ubyte {
-        return first+second
-    }
-
-    sub mul(ubyte first, ubyte second) -> ubyte {
-        return first*second
-    }
-
     sub start() {
+        str thing = "????"
+
+        if thing=="bmap" {
+            txt.print("gottem")
+        }
+
         ; TODO: test with builtin function using multiple args (such as mkword)
-        ubyte source=99
+;        ubyte value = add(3,4) |> add(10) |> mul(2) |> math.sin8u()     ; TODO should not work yet on vm codegen, but it compiles.... :/
-        ubyte value = add(3,4) |> add(10) |> mul(2) |> math.sin8u()     ; TODO should not work yet on vm codegen, but it compiles.... :/
+;        txt.print_ub(value)
-        txt.print_ub(value)
+;        txt.nl()
-        txt.nl()
+;        uword wvalue = add(3,4) |> add($30) |> mkword($ea)
-        uword wvalue = add(3,4) |> add($30) |> mkword($ea)
+;        txt.print_uwhex(wvalue, true)
-        txt.print_uwhex(wvalue, true)
+;        txt.nl()
-        txt.nl()
 
         ; expected output: aaabbb aaa bbb