[WASM] When expressions optimisations for String and Int constant cases

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/IrWhenUtils.kt

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.common
+
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.types.isBoolean
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.isElseBranch
+import org.jetbrains.kotlin.ir.util.isTrueConst
+
+object IrWhenUtils {
+// psi2ir lowers multiple cases to nested conditions. For example,
+//
+// when (subject) {
+//   a, b, c -> action
+// }
+//
+// is lowered to
+//
+// if (if (subject == a)
+//       true
+//     else
+//       if (subject == b)
+//         true
+//       else
+//         subject == c) {
+//     action
+// }
+//
+// fir2ir lowers the same to an or sequence:
+//
+// if (((subject == a) || (subject == b)) || (subject = c)) action
+//
+// @return true if the conditions are equality checks of constants.
+    fun matchConditions(ororSymbol: IrFunctionSymbol, condition: IrExpression): ArrayList<IrCall>? {
+        if (condition is IrWhen && condition.origin == IrStatementOrigin.WHEN_COMMA) {
+            assert(condition.type.isBoolean()) { "WHEN_COMMA should always be a Boolean: ${condition.dump()}" }
+
+            val candidates = ArrayList<IrCall>()
+
+            // Match the following structure:
+            //
+            // when() {
+            //   cond_1 -> true
+            //   cond_2 -> true
+            //   ...
+            //   else -> cond_N
+            // }
+            //
+            // Namely, the structure which returns true if any one of the condition is true.
+            for (branch in condition.branches) {
+                candidates += if (isElseBranch(branch)) {
+                    assert(branch.condition.isTrueConst()) { "IrElseBranch.condition should be const true: ${branch.condition.dump()}" }
+                    matchConditions(ororSymbol, branch.result) ?: return null
+                } else {
+                    if (!branch.result.isTrueConst()) return null
+                    matchConditions(ororSymbol, branch.condition) ?: return null
+                }
+            }
+            return candidates.ifEmpty { null }
+        } else if (condition is IrCall && condition.symbol == ororSymbol) {
+            val candidates = ArrayList<IrCall>()
+            for (i in 0 until condition.valueArgumentsCount) {
+                val argument = condition.getValueArgument(i)!!
+                candidates += matchConditions(ororSymbol, argument) ?: return null
+            }
+            return candidates.ifEmpty { null }
+        } else if (condition is IrCall) {
+            return arrayListOf(condition)
+        }
+
+        return null
+    }
+}

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/SwitchGenerator.kt

@@ -5,12 +5,11 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
+import org.jetbrains.kotlin.backend.common.IrWhenUtils
 import org.jetbrains.kotlin.codegen.`when`.SwitchCodegen.Companion.preferLookupOverSwitch
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.util.hasAnnotation
-import org.jetbrains.kotlin.ir.util.isTrueConst
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Type
@@ -32,7 +31,7 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
             if (branch is IrElseBranch) {
                 elseExpression = branch.result
             } else {
-                val conditions = matchConditions(branch.condition) ?: return null
+                val conditions = IrWhenUtils.matchConditions(codegen.context.irBuiltIns.ororSymbol, branch.condition) ?: return null
                 val thenLabel = Label()
                 expressionToLabels.add(ExpressionToLabel(branch.result, thenLabel))
                 callToLabels += conditions.map { CallToLabel(it, thenLabel) }
@@ -234,71 +233,6 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
         removeIf { it.label !in reachableLabels }
     }
 
-    // psi2ir lowers multiple cases to nested conditions. For example,
-    //
-    // when (subject) {
-    //   a, b, c -> action
-    // }
-    //
-    // is lowered to
-    //
-    // if (if (subject == a)
-    //       true
-    //     else
-    //       if (subject == b)
-    //         true
-    //       else
-    //         subject == c) {
-    //     action
-    // }
-    //
-    // fir2ir lowers the same to an or sequence:
-    //
-    // if (((subject == a) || (subject == b)) || (subject = c)) action
-    //
-    // @return true if the conditions are equality checks of constants.
-    private fun matchConditions(condition: IrExpression): ArrayList<IrCall>? {
-        if (condition is IrWhen && condition.origin == IrStatementOrigin.WHEN_COMMA) {
-            assert(condition.type.isBoolean()) { "WHEN_COMMA should always be a Boolean: ${condition.dump()}" }
-
-            val candidates = ArrayList<IrCall>()
-
-            // Match the following structure:
-            //
-            // when() {
-            //   cond_1 -> true
-            //   cond_2 -> true
-            //   ...
-            //   else -> cond_N
-            // }
-            //
-            // Namely, the structure which returns true if any one of the condition is true.
-            for (branch in condition.branches) {
-                candidates += if (branch is IrElseBranch) {
-                    assert(branch.condition.isTrueConst()) { "IrElseBranch.condition should be const true: ${branch.condition.dump()}" }
-                    matchConditions(branch.result) ?: return null
-                } else {
-                    if (!branch.result.isTrueConst())
-                        return null
-                    matchConditions(branch.condition) ?: return null
-                }
-            }
-
-            return if (candidates.isNotEmpty()) candidates else return null
-        } else if (condition is IrCall && condition.symbol == codegen.context.irBuiltIns.ororSymbol) {
-            val candidates = ArrayList<IrCall>()
-            for (i in 0 until condition.valueArgumentsCount) {
-                val argument = condition.getValueArgument(i)!!
-                candidates += matchConditions(argument) ?: return null
-            }
-            return if (candidates.isNotEmpty()) candidates else return null
-        } else if (condition is IrCall) {
-            return arrayListOf(condition)
-        }
-
-        return null
-    }
-
     abstract inner class Switch(
         val subject: IrExpression,
         val elseExpression: IrExpression?,

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt

@@ -142,6 +142,12 @@ private val tailrecLoweringPhase = makeWasmModulePhase(
     description = "Replace `tailrec` call sites with equivalent loop"
 )
 
+private val wasmStringSwitchOptimizerLowering = makeWasmModulePhase(
+    ::WasmStringSwitchOptimizerLowering,
+    name = "!!!",
+    description = "!!!"
+)
+
 private val complexExternalDeclarationsToTopLevelFunctionsLowering = makeWasmModulePhase(
     ::ComplexExternalDeclarationsToTopLevelFunctionsLowering,
     name = "ComplexExternalDeclarationsToTopLevelFunctionsLowering",
@@ -599,6 +605,8 @@ val wasmPhases = NamedCompilerPhase(
             delegateToPrimaryConstructorLoweringPhase then
             // Common prefix ends
 
+            wasmStringSwitchOptimizerLowering then
+
             complexExternalDeclarationsToTopLevelFunctionsLowering then
             complexExternalDeclarationsUsagesLowering then
 

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.wasm
 
 import org.jetbrains.kotlin.backend.common.ir.Symbols
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrCall
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -15,10 +16,11 @@ import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.backend.js.ReflectionSymbols
-import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
-import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
-import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
-import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
+import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.constructors
@@ -315,7 +317,6 @@ class WasmSymbols(
         return symbolTable.referenceSimpleFunction(tmp.single())
     }
 
-
     private fun getInternalFunction(name: String) = getFunction(name, wasmInternalPackage)
 
     private fun getIrClass(fqName: FqName): IrClassSymbol = symbolTable.referenceClass(getClass(fqName))

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt

@@ -542,6 +542,10 @@ class BodyGenerator(
     }
 
     override fun visitWhen(expression: IrWhen) {
+        if (tryGenerateOptimisedWhen(expression, context.backendContext.wasmSymbols)) {
+            return
+        }
+
         val resultType = context.transformBlockResultType(expression.type)
         var ifCount = 0
         var seenElse = false

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/OptimisedWhenGenerator.kt

@@ -0,0 +1,340 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.wasm.ir2wasm
+
+import org.jetbrains.kotlin.backend.common.IrWhenUtils
+import org.jetbrains.kotlin.backend.wasm.WasmSymbols
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.isElseBranch
+import org.jetbrains.kotlin.ir.visitors.acceptVoid
+import org.jetbrains.kotlin.wasm.ir.*
+
+private class ExtractedWhenCondition<T>(val condition: IrCall, val const: IrConst<T>)
+private class ExtractedWhenBranch<T>(val conditions: List<ExtractedWhenCondition<T>>, val expression: IrExpression)
+
+internal fun BodyGenerator.tryGenerateOptimisedWhen(expression: IrWhen, symbols: WasmSymbols): Boolean {
+    if (expression.branches.size <= 2) return false
+
+    var elseExpression: IrExpression? = null
+    val extractedBranches = mutableListOf<ExtractedWhenBranch<Any>>()
+
+    // Parse when structure. Note that the condition can be nested. See matchConditions() for details.
+    var noMultiplyConditionBranches = true
+    val seenConditions = mutableSetOf<Any>() //to filter out equal conditions
+    for (branch in expression.branches) {
+        if (isElseBranch(branch)) {
+            elseExpression = branch.result
+        } else {
+            val conditions = IrWhenUtils.matchConditions(symbols.irBuiltIns.ororSymbol, branch.condition) ?: return false
+            val extractedConditions = tryExtractEqEqNumberConditions(symbols, conditions) ?: return false
+            val filteredExtractedConditions = extractedConditions.filter { it.const.value !in seenConditions }
+            seenConditions.addAll(extractedConditions.map { it.const.value })
+            if (filteredExtractedConditions.isNotEmpty()) {
+                noMultiplyConditionBranches = noMultiplyConditionBranches && filteredExtractedConditions.size == 1
+                extractedBranches.add(ExtractedWhenBranch(filteredExtractedConditions, branch.result))
+            }
+        }
+    }
+    if (extractedBranches.isEmpty()) return false
+    val subject = extractedBranches[0].conditions[0].condition.getValueArgument(0) ?: return false
+
+    // Check all kinds are the same
+    for (branch in extractedBranches) {
+        //TODO: Support all primitive types
+        if (!branch.conditions.all { it.const.kind.equals(IrConstKind.Int) }) return false
+    }
+
+    val intBranches = extractedBranches.map { branch ->
+        @Suppress("UNCHECKED_CAST")
+        branch as ExtractedWhenBranch<Int>
+    }
+
+    val maxValue = intBranches.maxOf { branch -> branch.conditions.maxOf { it.const.value } }
+    val minValue = intBranches.minOf { branch -> branch.conditions.minOf { it.const.value } }
+    if (minValue == maxValue) return false
+
+    val selectorLocal = context.referenceLocal(SyntheticLocalType.TABLE_SWITCH_SELECTOR)
+    subject.accept(this, null)
+    body.buildSetLocal(selectorLocal)
+
+    val resultType = context.transformBlockResultType(expression.type)
+    //int overflow or load is too small then make table switch
+    val tableSize = maxValue - minValue
+    if (tableSize <= 0 || tableSize > seenConditions.size * 2) {
+        if (noMultiplyConditionBranches) {
+            createBinaryTable(
+                selectorLocal = selectorLocal,
+                intBranches = intBranches,
+                elseExpression = elseExpression,
+                resultType = resultType,
+                nothingType = symbols.irBuiltIns.nothingType
+            )
+        } else {
+            createBinaryTable(selectorLocal, intBranches)
+            body.buildSetLocal(selectorLocal)
+            genTableIntSwitch(
+                selectorLocal = selectorLocal,
+                resultType = resultType,
+                branches = intBranches,
+                elseExpression = elseExpression,
+                shift = 0,
+                brTable = intBranches.indices.toList(),
+                nothingType = symbols.irBuiltIns.nothingType
+            )
+        }
+    } else {
+        val brTable = mutableListOf<Int>()
+        for (i in minValue..maxValue) {
+            val branchIndex = intBranches.indexOfFirst { branch -> branch.conditions.any { it.const.value == i } }
+            val brIndex = if (branchIndex != -1) branchIndex else intBranches.size
+            brTable.add(brIndex)
+        }
+        genTableIntSwitch(
+            selectorLocal = selectorLocal,
+            resultType = resultType,
+            branches = intBranches,
+            elseExpression = elseExpression,
+            shift = minValue,
+            brTable = brTable,
+            nothingType = symbols.irBuiltIns.nothingType
+        )
+    }
+    return true
+}
+
+private fun IrExpression.safeAccept(bodyGenerator: BodyGenerator, nothingType: IrType) {
+    acceptVoid(bodyGenerator)
+    if (type == nothingType) {
+        bodyGenerator.body.buildUnreachable()
+    }
+}
+
+/**
+ * Create binary search for when that emit branches index in leafs
+ * when (a) {
+ *   123 -> expr1
+ *   456 -> expr2
+ *   else -> elseExpr
+ * }
+ * crates binary search linked to index of branch
+ *   IF (a < 456) {
+ *      IF (a == 123)
+ *        #expr1
+ *      ELSE
+ *        #else
+ *      END IF
+ *   ELSE
+ *      IF (1 == 456)
+ *        #expr2
+ *      ELSE
+ *        #else
+ *      END IF
+ *   END IF
+ * }
+ */
+private fun BodyGenerator.createBinaryTable(selectorLocal: WasmLocal, intBranches: List<ExtractedWhenBranch<Int>>) {
+    val sortedCaseToBranchIndex = mutableListOf<Pair<Int, Int>>()
+    intBranches.flatMapIndexedTo(sortedCaseToBranchIndex) { index, branch -> branch.conditions.map { it.const.value to index } }
+    sortedCaseToBranchIndex.sortBy { it.first }
+
+    val thenBody = { result: Int ->
+        body.buildConstI32(result)
+    }
+    val elseBody: () -> Unit = {
+        body.buildConstI32(intBranches.size)
+    }
+    createBinaryTable(selectorLocal, WasmI32, sortedCaseToBranchIndex, 0, sortedCaseToBranchIndex.size, thenBody, elseBody)
+}
+
+private fun tryExtractEqEqNumberConditions(symbols: WasmSymbols, conditions: List<IrCall>): List<ExtractedWhenCondition<Any>>? {
+    if (conditions.isEmpty()) return null
+
+    val firstCondition = conditions[0]
+    val firstConditionSymbol = firstCondition.symbol
+        .takeIf { it in symbols.equalityFunctions.values }
+        ?: return null
+    if (firstCondition.valueArgumentsCount != 2) return null
+
+    // All conditions has the same eqeq
+    if (conditions.any { it.symbol != firstConditionSymbol }) return null
+
+    val result = mutableListOf<ExtractedWhenCondition<Any>>()
+    for (condition in conditions) {
+        if (condition.symbol != firstConditionSymbol) return null
+        @Suppress("UNCHECKED_CAST")
+        val conditionConst = condition.getValueArgument(1) as? IrConst<Any> ?: return null
+        result.add(ExtractedWhenCondition(condition, conditionConst))
+    }
+
+    return result
+}
+
+/**
+ * Create binary search for when that emit when expressions in leafs
+ * when (a) {
+ *   123 -> expr1
+ *   456 -> expr2
+ *   else -> elseExpr
+ * }
+ * crates binary search linked to index of branch
+ * BLOCK
+ *   IF (a < 456) {
+ *      IF (a == 123)
+ *        #expr1
+ *        GOTO END BLOCK
+ *      END IF
+ *   ELSE
+ *      IF (1 == 456)
+ *        #expr2
+ *        GOTO END BLOCK
+ *      END IF
+ *   END IF
+ *   elseExpr
+ * END BLOCK
+ * }
+*/
+private fun BodyGenerator.createBinaryTable(
+    selectorLocal: WasmLocal,
+    intBranches: List<ExtractedWhenBranch<Int>>,
+    elseExpression: IrExpression?,
+    resultType: WasmType?,
+    nothingType: IrType,
+) {
+    val sortedCaseToBranchIndex = mutableListOf<Pair<Int, IrExpression>>()
+    intBranches.mapTo(sortedCaseToBranchIndex) { branch -> branch.conditions[0].const.value to branch.expression }
+    sortedCaseToBranchIndex.sortBy { it.first }
+
+    body.buildBlock("when_block", resultType) { currentBlock ->
+        val thenBody = { result: IrExpression ->
+            result.safeAccept(this, nothingType)
+            body.buildBr(currentBlock)
+        }
+        val elseBody: () -> Unit = {
+            resultType?.let { generateDefaultInitializerForType(it, body) }
+        }
+        createBinaryTable(
+            selectorLocal = selectorLocal,
+            resultType = resultType,
+            sortedCases = sortedCaseToBranchIndex,
+            fromIncl = 0,
+            toExcl = sortedCaseToBranchIndex.size,
+            thenBody = thenBody,
+            elseBody = elseBody
+        )
+
+        if (resultType != null) {
+            body.buildDrop()
+        }
+        elseExpression?.safeAccept(this, nothingType)
+    }
+}
+
+private fun <T> BodyGenerator.createBinaryTable(
+    selectorLocal: WasmLocal,
+    resultType: WasmType?,
+    sortedCases: List<Pair<Int, T>>,
+    fromIncl: Int,
+    toExcl: Int,
+    thenBody: (T) -> Unit,
+    elseBody: () -> Unit
+) {
+    val size = toExcl - fromIncl
+    if (size == 1) {
+        val (case, result) = sortedCases[fromIncl]
+        body.buildGetLocal(selectorLocal)
+        body.buildConstI32(case)
+        body.buildInstr(WasmOp.I32_EQ)
+        body.buildIf("binary_tree_branch", resultType)
+        thenBody(result)
+        body.buildElse()
+        elseBody()
+        body.buildEnd()
+        return
+    }
+
+    val border = fromIncl + size / 2
+
+    body.buildGetLocal(selectorLocal)
+    body.buildConstI32(sortedCases[border].first)
+    body.buildInstr(WasmOp.I32_LT_S)
+    body.buildIf("binary_tree_node", resultType)
+    createBinaryTable(selectorLocal, resultType, sortedCases, fromIncl, border, thenBody, elseBody)
+    body.buildElse()
+    createBinaryTable(selectorLocal, resultType, sortedCases, border, toExcl, thenBody, elseBody)
+    body.buildEnd()
+}
+
+
+/**
+ * Create table switch with expressions
+ * when (a) {
+ *   0 -> expr1
+ *   1 -> expr2
+ *   else -> elseExpr
+ * }
+ * crates binary search linked to index of branch
+ * BLOCK FOR ELSE
+ *   BLOCK1
+ *     BLOCK2
+ *       BLOCK FOR BRTABLE
+ *         BRTABLE 0 1 2
+ *       END BLOCK FOR BRTABLE
+ *       expr1
+ *       GOTO END BLOCK FOR ELSE
+ *     END BLOCK1
+ *     expr2
+ *     GOTO END BLOCK FOR ELSE
+ *   END BLOCK2
+ *   elseExpr
+ * END BLOCK FOR ELSE
+ */
+private fun BodyGenerator.genTableIntSwitch(
+    selectorLocal: WasmLocal,
+    resultType: WasmType?,
+    branches: List<ExtractedWhenBranch<Int>>,
+    elseExpression: IrExpression?,
+    shift: Int,
+    brTable: List<Int>,
+    nothingType: IrType,
+) {
+    val baseBlockIndex = body.numberOfNestedBlocks
+    //expressions + else branch + br_table
+    repeat(branches.size + 2) {
+        body.buildBlock(resultType)
+    }
+
+    resultType?.let { generateDefaultInitializerForType(it, body) } //stub value
+    body.buildGetLocal(selectorLocal)
+    if (shift != 0) {
+        body.buildConstI32(shift)
+        body.buildInstr(WasmOp.I32_SUB)
+    }
+    body.buildInstr(
+        WasmOp.BR_TABLE,
+        WasmImmediate.LabelIdxVector(brTable),
+        WasmImmediate.LabelIdx(branches.size)
+    )
+    body.buildEnd()
+
+    for (expression in branches) {
+        if (resultType != null) {
+            body.buildDrop()
+        }
+        expression.expression.safeAccept(this, nothingType)
+
+        body.buildBr(baseBlockIndex + 1)
+        body.buildEnd()
+    }
+
+    if (resultType != null) {
+        body.buildDrop()
+    }
+    elseExpression?.safeAccept(this, nothingType)
+
+    body.buildEnd()
+    check(baseBlockIndex == body.numberOfNestedBlocks)
+}

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContext.kt

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.wasm.ir.WasmInstr
 import org.jetbrains.kotlin.wasm.ir.WasmLocal
 
 enum class LoopLabelType { BREAK, CONTINUE }
-enum class SyntheticLocalType { IS_INTERFACE_PARAMETER }
+enum class SyntheticLocalType { IS_INTERFACE_PARAMETER, TABLE_SWITCH_SELECTOR }
 
 interface WasmFunctionCodegenContext : WasmBaseCodegenContext {
     val irFunction: IrFunction

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContextImpl.kt

@@ -53,11 +53,17 @@ class WasmFunctionCodegenContextImpl(
         return wasmFunction.locals[index]
     }
 
+    private val SyntheticLocalType.wasmType
+        get() = when (this) {
+            SyntheticLocalType.IS_INTERFACE_PARAMETER -> WasmRefNullType(WasmHeapType.Type(referenceGcType(backendContext.irBuiltIns.anyClass)))
+            SyntheticLocalType.TABLE_SWITCH_SELECTOR -> WasmI32
+        }
+
     override fun referenceLocal(type: SyntheticLocalType): WasmLocal = wasmSyntheticLocals.getOrPut(type) {
         WasmLocal(
             wasmFunction.locals.size,
             type.name,
-            WasmRefNullType(WasmHeapType.Type(referenceGcType(backendContext.irBuiltIns.anyClass))),
+            type.wasmType,
             isParameter = false
         ).also {
             wasmFunction.locals += it

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmStringSwitchOptimizerLowering.kt

@@ -0,0 +1,290 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.wasm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.IrWhenUtils
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.buildVariable
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOriginImpl
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.interpreter.toIrConst
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.isNullable
+import org.jetbrains.kotlin.ir.util.getSimpleFunction
+import org.jetbrains.kotlin.ir.util.isElseBranch
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.name.Name
+
+private object OPTIMISED_WHEN_SUBJECT : IrDeclarationOriginImpl("OPTIMISED_WHEN_SUBJECT")
+
+class WasmStringSwitchOptimizerLowering(
+    private val context: WasmBackendContext
+) : FileLoweringPass, IrElementTransformerVoidWithContext() {
+    private val symbols = context.wasmSymbols
+
+    private val stringHashCode by lazy {
+        symbols.irBuiltIns.stringClass.getSimpleFunction("hashCode")!!
+    }
+
+    private val intType: IrType = symbols.irBuiltIns.intType
+    private val booleanType: IrType = symbols.irBuiltIns.booleanType
+
+    private fun IrBlockBuilder.createEqEqForIntVariable(tempIntVariable: IrVariable, value: Int) =
+        irCall(context.irBuiltIns.eqeqSymbol, booleanType).also {
+            it.putValueArgument(0, irGet(tempIntVariable))
+            it.putValueArgument(1, value.toIrConst(intType))
+        }
+
+    private fun asEqCall(expression: IrExpression): IrCall? =
+        (expression as? IrCall)?.takeIf { it.symbol == context.irBuiltIns.eqeqSymbol }
+
+    private class MatchedCase(val condition: IrCall, val branchIndex: Int)
+    private class BucketSelector(val hashCode: Int, val selector: IrExpression)
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    private fun tryMatchCaseToNullableStringConstant(condition: IrExpression): IrConst<*>? {
+        val eqCall = asEqCall(condition) ?: return null
+        if (eqCall.valueArgumentsCount < 2) return null
+        val constantReceiver = eqCall.getValueArgument(1) as? IrConst<*> ?: return null
+        return when (constantReceiver.kind) {
+            IrConstKind.String, IrConstKind.Null -> constantReceiver
+            else -> null
+        }
+    }
+
+    private fun IrBlockBuilder.addHashCodeVariable(firstEqCall: IrCall): IrVariable {
+        val subject = firstEqCall.getValueArgument(0)!!
+        val subjectType = subject.type
+
+        val whenSubject = buildVariable(
+            scope.getLocalDeclarationParent(),
+            startOffset,
+            endOffset,
+            OPTIMISED_WHEN_SUBJECT,
+            Name.identifier("tmp_when_subject"),
+            subjectType,
+        )
+
+        whenSubject.initializer = subject
+        +whenSubject
+        firstEqCall.putValueArgument(0, irGet(whenSubject))
+
+        val tmpIntWhenSubject = buildVariable(
+            scope.getLocalDeclarationParent(),
+            startOffset,
+            endOffset,
+            OPTIMISED_WHEN_SUBJECT,
+            Name.identifier("tmp_int_when_subject"),
+            intType,
+        )
+
+        val getHashCode = irCall(stringHashCode, intType).also {
+            it.dispatchReceiver = irGet(whenSubject)
+        }
+
+        val hashCode: IrExpression = if (subjectType.isNullable()) {
+            val stringIsNull = irCall(context.irBuiltIns.eqeqeqSymbol, booleanType).also {
+                it.putValueArgument(0, irGet(whenSubject))
+                it.putValueArgument(1, irNull(subjectType))
+            }
+            irIfThenElse(intType, stringIsNull, 0.toIrConst(intType), getHashCode)
+        } else {
+            getHashCode
+        }
+
+        tmpIntWhenSubject.initializer = hashCode
+        +tmpIntWhenSubject
+
+        return tmpIntWhenSubject
+    }
+
+    /**
+     * Create simple 1-element buckets (for when without else block and commas)
+     * when(a) {
+     *  "123" -> 123
+     *  "456" -> 456
+     *  "789" -> 789
+     *  }
+     *  into the integer when's collections of
+     *  48690 -> if(a == "123") -> 123
+     *  51669 -> if(a == "456") -> 456
+     *  54648 -> if(a == "789") -> 789
+     */
+    private fun IrBlockBuilder.createSimpleBucketSelectors(
+        stringConstantToMatchedCase: Map<String?, MatchedCase>,
+        buckets: Map<Int, List<String?>>,
+        transformedWhen: IrWhen,
+    ): List<BucketSelector> = buckets.entries.map { bucket ->
+        val selector = if (bucket.value.size == 1) {
+            val bucketCase = bucket.value[0]
+            val matchedCase = stringConstantToMatchedCase.getValue(bucketCase)
+            irIfThen(
+                type = transformedWhen.type,
+                condition = matchedCase.condition,
+                thenPart = transformedWhen.branches[matchedCase.branchIndex].result,
+            )
+        } else {
+            val bucketBranches = mutableListOf<IrBranch>()
+            bucket.value.mapTo(bucketBranches) { bucketCase ->
+                val matchedCase = stringConstantToMatchedCase.getValue(bucketCase)
+                irBranch(matchedCase.condition, transformedWhen.branches[matchedCase.branchIndex].result)
+            }
+            irWhen(transformedWhen.type, bucketBranches)
+        }
+        BucketSelector(bucket.key, selector)
+    }
+
+    private fun IrBlockBuilder.createWhenForBucketSelectors(
+        tempIntVariable: IrVariable,
+        bucketsSelectors: List<BucketSelector>,
+        selectorsType: IrType,
+        elseBranchExpression: IrExpression?
+    ): IrWhen {
+        val allBucketsWhenBranches = mutableListOf<IrBranch>()
+        bucketsSelectors.mapTo(allBucketsWhenBranches) { bucketSelector ->
+            val condition = createEqEqForIntVariable(tempIntVariable, bucketSelector.hashCode)
+            irBranch(condition, bucketSelector.selector)
+        }
+        if (elseBranchExpression != null) {
+            allBucketsWhenBranches.add(irElseBranch(elseBranchExpression))
+        }
+        return irWhen(selectorsType, allBucketsWhenBranches)
+    }
+
+    /**
+     * Create multi-element buckets for every hashCode
+     * 48690 -> when(a) {
+     *   "123" -> 0
+     *   "ARcZguv123" -> 1
+     *   else -> 3
+     * }
+     * 51669 -> when(a) {
+     *   "456" -> 0
+     *   else -> 3
+     * }
+     * 54648 -> when(a) {
+     *   "789" -> 1
+     *   else -> 3
+     * }
+     * else -> 3
+     */
+    private fun IrBlockBuilder.createBucketSelectors(
+        stringConstantToMatchedCase: Map<String?, MatchedCase>,
+        buckets: Map<Int, List<String?>>,
+        elseBranchIndex: Int,
+    ): List<BucketSelector> = buckets.entries.map { bucket ->
+        val selector = if (bucket.value.size == 1) {
+            val bucketCase = bucket.value[0]
+            val matchedCase = stringConstantToMatchedCase.getValue(bucketCase)
+            irIfThenElse(
+                type = intType,
+                condition = matchedCase.condition,
+                thenPart = matchedCase.branchIndex.toIrConst(intType),
+                elsePart = elseBranchIndex.toIrConst(intType)
+            )
+        } else {
+            val bucketBranches = mutableListOf<IrBranch>()
+            bucket.value.mapTo(bucketBranches) { bucketCase ->
+                val matchedCase = stringConstantToMatchedCase.getValue(bucketCase)
+                irBranch(matchedCase.condition, matchedCase.branchIndex.toIrConst(intType))
+            }
+            bucketBranches.add(irElseBranch(elseBranchIndex.toIrConst(intType)))
+            irWhen(intType, bucketBranches)
+        }
+        BucketSelector(bucket.key, selector)
+    }
+
+    private fun IrBlockBuilder.createTransformedWhen(tempIntVariable: IrVariable, transformedWhen: IrWhen): IrWhen {
+        val mainResultsBranches = mutableListOf<IrBranch>()
+        transformedWhen.branches.mapIndexedTo(mainResultsBranches) { index, branch ->
+            if (!isElseBranch(branch)) {
+                irBranch(createEqEqForIntVariable(tempIntVariable, index), branch.result)
+            } else {
+                branch
+            }
+        }
+        return irWhen(transformedWhen.type, mainResultsBranches)
+    }
+
+    override fun visitWhen(expression: IrWhen): IrExpression {
+        val visitedWhen = super.visitWhen(expression) as IrWhen
+        if (visitedWhen.branches.size <= 2) return visitedWhen
+
+        var firstEqCall: IrCall? = null
+        var isSimpleWhen = true //simple when is when without else block and commas
+        val stringConstantToMatchedCase = mutableMapOf<String?, MatchedCase>()
+        visitedWhen.branches.forEachIndexed { branchIndex, branch ->
+            if (!isElseBranch(branch)) {
+                val conditions = IrWhenUtils.matchConditions(context.irBuiltIns.ororSymbol, branch.condition) ?: return visitedWhen
+                if (conditions.isEmpty()) return visitedWhen
+
+                isSimpleWhen = isSimpleWhen && conditions.size == 1
+                for (condition in conditions) {
+                    val matchedStringConstant = tryMatchCaseToNullableStringConstant(condition) ?: return visitedWhen
+                    val matchedString = matchedStringConstant.value as? String
+                    if (matchedString !in stringConstantToMatchedCase) {
+                        stringConstantToMatchedCase[matchedString] = MatchedCase(condition, branchIndex)
+                        firstEqCall = firstEqCall ?: asEqCall(condition)
+                    }
+                }
+            } else {
+                isSimpleWhen = false
+            }
+        }
+
+        if (firstEqCall == null || stringConstantToMatchedCase.size < 2) return visitedWhen
+
+        val convertedBlock = context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol).run {
+            irBlock(resultType = visitedWhen.type) {
+                val tempIntVariable = addHashCodeVariable(firstEqCall!!)
+
+                val buckets = stringConstantToMatchedCase.keys.groupBy { it.hashCode() }
+
+                if (isSimpleWhen) {
+                    val bucketsSelectors = createSimpleBucketSelectors(
+                        stringConstantToMatchedCase = stringConstantToMatchedCase,
+                        buckets = buckets,
+                        transformedWhen = expression
+                    )
+                    +createWhenForBucketSelectors(
+                        tempIntVariable = tempIntVariable,
+                        bucketsSelectors = bucketsSelectors,
+                        selectorsType = expression.type,
+                        elseBranchExpression = null
+                    )
+                } else {
+                    val elseBranchIndex = expression.branches.size
+                    val bucketsSelectors = createBucketSelectors(
+                        stringConstantToMatchedCase = stringConstantToMatchedCase,
+                        buckets = buckets,
+                        elseBranchIndex = elseBranchIndex
+                    )
+                    val caseSelectorWhen = createWhenForBucketSelectors(
+                        tempIntVariable = tempIntVariable,
+                        bucketsSelectors = bucketsSelectors,
+                        selectorsType = intType,
+                        elseBranchExpression = elseBranchIndex.toIrConst(intType)
+                    )
+                    +irSet(tempIntVariable, caseSelectorWhen)
+                    +createTransformedWhen(tempIntVariable, expression)
+                }
+            }
+        }
+
+        return convertedBlock
+    }
+}
+

libraries/stdlib/wasm/internal/kotlin/wasm/internal/ExternalWrapper.kt

@@ -105,7 +105,6 @@ internal fun ExternalInterfaceType.externAsWasmAnyref(): anyref =
 internal fun Any?.asWasmExternRef(): ExternalInterfaceType =
     implementedAsIntrinsic
 
-
 @JsFun("(ref) => ref == null")
 internal external fun isNullish(ref: ExternalInterfaceType): Boolean
 

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmExpressionBuilder.kt

@@ -59,6 +59,12 @@ abstract class WasmExpressionBuilder {
         buildInstr(WasmOp.ELSE)
     }
 
+    fun buildBlock(resultType: WasmType? = null): Int {
+        buildInstr(WasmOp.BLOCK, WasmImmediate.BlockType.Value(resultType))
+        numberOfNestedBlocks++
+        return numberOfNestedBlocks
+    }
+
     fun buildEnd() {
         numberOfNestedBlocks--
         buildInstr(WasmOp.END)