From 95396eb61a06981a0ee43c26f4dcab730463694d Mon Sep 17 00:00:00 2001 From: Igor Yakovlev Date: Thu, 14 Jul 2022 13:04:45 +0200 Subject: [PATCH] [WASM] When expressions optimisations for String and Int constant cases --- .../kotlin/backend/common/IrWhenUtils.kt | 78 ++++ .../backend/jvm/codegen/SwitchGenerator.kt | 70 +--- .../kotlin/backend/wasm/WasmLoweringPhases.kt | 8 + .../kotlin/backend/wasm/WasmSymbols.kt | 11 +- .../backend/wasm/ir2wasm/BodyGenerator.kt | 4 + .../wasm/ir2wasm/OptimisedWhenGenerator.kt | 340 ++++++++++++++++++ .../ir2wasm/WasmFunctionCodegenContext.kt | 2 +- .../ir2wasm/WasmFunctionCodegenContextImpl.kt | 8 +- .../WasmStringSwitchOptimizerLowering.kt | 290 +++++++++++++++ .../kotlin/wasm/internal/ExternalWrapper.kt | 1 - .../kotlin/wasm/ir/WasmExpressionBuilder.kt | 6 + 11 files changed, 742 insertions(+), 76 deletions(-) create mode 100644 compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/IrWhenUtils.kt create mode 100644 compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/OptimisedWhenGenerator.kt create mode 100644 compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmStringSwitchOptimizerLowering.kt diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/IrWhenUtils.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/IrWhenUtils.kt new file mode 100644 index 00000000000..82df29b9569 --- /dev/null +++ b/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? { + 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() + + // 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() + 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 + } +} \ No newline at end of file diff --git a/compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/SwitchGenerator.kt b/compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/SwitchGenerator.kt index 00029a53184..6e6da9a1f79 100644 --- a/compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/SwitchGenerator.kt +++ b/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? { - 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() - - // 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() - 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?, diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt index 27727882437..ef1e0bd22ae 100644 --- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt +++ b/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 diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt index 517251faacb..4ef4f190730 100644 --- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt +++ b/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)) diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt index 6053bb45b4e..0ea017e390f 100644 --- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt +++ b/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 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 new file mode 100644 index 00000000000..8badc3350a3 --- /dev/null +++ b/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(val condition: IrCall, val const: IrConst) +private class ExtractedWhenBranch(val conditions: List>, 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>() + + // Parse when structure. Note that the condition can be nested. See matchConditions() for details. + var noMultiplyConditionBranches = true + val seenConditions = mutableSetOf() //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 + } + + 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() + 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>) { + val sortedCaseToBranchIndex = mutableListOf>() + 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): List>? { + 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>() + for (condition in conditions) { + if (condition.symbol != firstConditionSymbol) return null + @Suppress("UNCHECKED_CAST") + val conditionConst = condition.getValueArgument(1) as? IrConst ?: 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>, + elseExpression: IrExpression?, + resultType: WasmType?, + nothingType: IrType, +) { + val sortedCaseToBranchIndex = mutableListOf>() + 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 BodyGenerator.createBinaryTable( + selectorLocal: WasmLocal, + resultType: WasmType?, + sortedCases: List>, + 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>, + elseExpression: IrExpression?, + shift: Int, + brTable: List, + 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) +} \ No newline at end of file diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContext.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContext.kt index 9f8767ea538..bc5acc2af52 100644 --- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContext.kt +++ b/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 diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContextImpl.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContextImpl.kt index 2f1caf37c9a..05395d0a373 100644 --- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmFunctionCodegenContextImpl.kt +++ b/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 diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmStringSwitchOptimizerLowering.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmStringSwitchOptimizerLowering.kt new file mode 100644 index 00000000000..46421ed90f2 --- /dev/null +++ b/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, + buckets: Map>, + transformedWhen: IrWhen, + ): List = 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() + 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, + selectorsType: IrType, + elseBranchExpression: IrExpression? + ): IrWhen { + val allBucketsWhenBranches = mutableListOf() + 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, + buckets: Map>, + elseBranchIndex: Int, + ): List = 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() + 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() + 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() + 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 + } +} + diff --git a/libraries/stdlib/wasm/internal/kotlin/wasm/internal/ExternalWrapper.kt b/libraries/stdlib/wasm/internal/kotlin/wasm/internal/ExternalWrapper.kt index a9f36fde475..8cdc0aaaed9 100644 --- a/libraries/stdlib/wasm/internal/kotlin/wasm/internal/ExternalWrapper.kt +++ b/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 diff --git a/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmExpressionBuilder.kt b/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmExpressionBuilder.kt index ca675725b7f..330ef72f5f6 100644 --- a/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmExpressionBuilder.kt +++ b/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)