From 5d1379ead19de7b6c5c6c2b0803e43fc33340a2a Mon Sep 17 00:00:00 2001
From: Zalim Bashorov <zalim.bashorov@jetbrains.com>
Date: Fri, 21 Jul 2023 14:40:57 +0200
Subject: [PATCH] [Wasm] Optimize `when` using br_table only if all conditions
 read and compare the same variable

#KT-60212 Fixed
---
 .../wasm/ir2wasm/OptimisedWhenGenerator.kt    | 11 +++++++
 compiler/testData/codegen/box/when/kt60212.kt | 32 +++++++++++++++++++
 2 files changed, 43 insertions(+)
 create mode 100644 compiler/testData/codegen/box/when/kt60212.kt

diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/OptimisedWhenGenerator.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/OptimisedWhenGenerator.kt
index 74c5cc6ad3b..b4f1ea53ce3 100644
--- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/OptimisedWhenGenerator.kt
+++ b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/OptimisedWhenGenerator.kt
@@ -44,6 +44,17 @@ internal fun BodyGenerator.tryGenerateOptimisedWhen(expression: IrWhen, symbols:
     if (extractedBranches.isEmpty()) return false
     val subject = extractedBranches[0].conditions[0].condition.getValueArgument(0) ?: return false
 
+    // Do the optimization only if all conditions read and compare the same var or val
+    // TODO: consider supporting other cases
+    if (subject !is IrGetValue) return false
+    val subjectValue = subject.symbol
+    val allConditionsReadsSameValue = !extractedBranches.all { branch ->
+        branch.conditions.all { whenCondition ->
+            (whenCondition.condition.getValueArgument(0) as? IrGetValue)?.symbol == subjectValue
+        }
+    }
+    if (allConditionsReadsSameValue) return false
+
     // Check all kinds are the same
     for (branch in extractedBranches) {
         //TODO: Support all primitive types
diff --git a/compiler/testData/codegen/box/when/kt60212.kt b/compiler/testData/codegen/box/when/kt60212.kt
new file mode 100644
index 00000000000..d473ab668a0
--- /dev/null
+++ b/compiler/testData/codegen/box/when/kt60212.kt
@@ -0,0 +1,32 @@
+fun box(): String {
+    testIfs()
+    testWhen()
+
+    return "OK"
+}
+
+// From KT-60212
+fun testIfs() {
+    val b1 = 0b11011111
+    if (b1 and 0b10000000 == 0) {
+        fail("A")
+    } else if (b1 and 0b11100000 == 0b11000000) {
+        return
+    } else {
+        fail("C")
+    }
+}
+
+// Snippet from kotlinx-io
+fun testWhen() {
+    val v = test(0xdf)
+    if (v != 1) fail("D")
+}
+
+fun test(b0: Int) = when {
+    b0 and 0x80 == 0 -> 0
+    b0 and 0xe0 == 0xc0 -> 1
+    b0 and 0xf0 == 0xe0 -> 2
+    b0 and 0xf8 == 0xf0 -> 3
+    else -> -1
+}