[WASM] Native support of returnable blocks

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

@@ -144,8 +144,8 @@ private val tailrecLoweringPhase = makeWasmModulePhase(
 
 private val wasmStringSwitchOptimizerLowering = makeWasmModulePhase(
     ::WasmStringSwitchOptimizerLowering,
-    name = "!!!",
-    description = "!!!"
+    name = "WasmStringSwitchOptimizerLowering",
+    description = "Replace when with constant string cases to binary search by string hashcodes"
 )
 
 private val complexExternalDeclarationsToTopLevelFunctionsLowering = makeWasmModulePhase(
@@ -403,13 +403,6 @@ private val tryCatchCanonicalization = makeWasmModulePhase(
     prerequisite = setOf(functionInliningPhase)
 )
 
-private val returnableBlockLoweringPhase = makeWasmModulePhase(
-    ::ReturnableBlockLowering,
-    name = "ReturnableBlockLowering",
-    description = "Replace returnable block with do-while loop",
-    prerequisite = setOf(functionInliningPhase)
-)
-
 private val bridgesConstructionPhase = makeWasmModulePhase(
     ::WasmBridgesConstruction,
     name = "BridgesConstruction",
@@ -627,7 +620,6 @@ val wasmPhases = NamedCompilerPhase(
             unhandledExceptionLowering then
 
             tryCatchCanonicalization then
-            returnableBlockLoweringPhase then
 
             forLoopsLoweringPhase then
             propertyLazyInitLoweringPhase then

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

@@ -3,8 +3,6 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-@file:OptIn(ExperimentalUnsignedTypes::class)
-
 package org.jetbrains.kotlin.backend.wasm.ir2wasm
 
 import org.jetbrains.kotlin.backend.common.ir.returnType
@@ -21,6 +19,7 @@ import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
@@ -30,7 +29,8 @@ import org.jetbrains.kotlin.wasm.ir.*
 
 class BodyGenerator(
     val context: WasmFunctionCodegenContext,
-    private val hierarchyDisjointUnions: DisjointUnions<IrClassSymbol>
+    private val hierarchyDisjointUnions: DisjointUnions<IrClassSymbol>,
+    private val isGetUnitFunction: Boolean,
 ) : IrElementVisitorVoid {
     val body: WasmExpressionBuilder = context.bodyGen
 
@@ -45,12 +45,12 @@ class BodyGenerator(
     }
 
     // Generates code for the given IR element. Leaves something on the stack unless expression was of the type Void.
-    private fun generateExpression(elem: IrElement) {
+    internal fun generateExpression(elem: IrElement) {
         assert(elem is IrExpression || elem is IrVariable) { "Unsupported statement kind" }
 
         elem.acceptVoid(this)
 
-        if (elem is IrExpression && elem.type == irBuiltIns.nothingType) {
+        if (elem is IrExpression && elem.type.isNothing()) {
             body.buildUnreachable()
         }
     }
@@ -296,7 +296,7 @@ class BodyGenerator(
         }
 
         if (tryToGenerateIntrinsicCall(call, function)) {
-            if (function.returnType == irBuiltIns.unitType)
+            if (function.returnType.isUnit())
                 body.buildGetUnit()
             return
         }
@@ -346,7 +346,7 @@ class BodyGenerator(
         }
 
         // Unit types don't cross function boundaries
-        if (function.returnType == irBuiltIns.unitType)
+        if (function.returnType.isUnit())
             body.buildGetUnit()
     }
 
@@ -498,21 +498,36 @@ class BodyGenerator(
         val statements = expression.statements
 
         if (statements.isEmpty()) {
-            if (expression.type == irBuiltIns.unitType)
+            if (expression.type == irBuiltIns.unitType) {
                 body.buildGetUnit()
+            }
             return
         }
 
-        statements.dropLast(1).forEach {
-            generateStatement(it)
+        if (expression is IrReturnableBlock) {
+            context.defineNonLocalReturnLevel(
+                expression.symbol,
+                body.buildBlock(context.transformBlockResultType(expression.type))
+            )
         }
 
-        generateExpression(statements.last())
+        statements.forEachIndexed { i, statement ->
+            if (i != statements.lastIndex) {
+                generateStatement(statement)
+            } else {
+                if (statement is IrExpression) {
+                    generateWithExpectedType(statement, expression.type)
+                } else {
+                    generateStatement(statement)
+                    if (expression.type != wasmSymbols.voidType) {
+                        body.buildGetUnit()
+                    }
+                }
+            }
+        }
 
-        // This handles cases where the last statement of a block is declaration which doesn't produce any value,
-        // but the block itself marked with the unit type.
-        if (statements.last() !is IrExpression && expression.type != wasmSymbols.voidType) {
-            body.buildGetUnit()
+        if (expression is IrReturnableBlock) {
+            body.buildEnd()
         }
     }
 
@@ -526,13 +541,16 @@ class BodyGenerator(
         body.buildBr(context.referenceLoopLevel(jump.loop, LoopLabelType.CONTINUE))
     }
 
-    override fun visitReturn(expression: IrReturn) {
-        if (expression.returnTargetSymbol.owner.returnType(backendContext) == irBuiltIns.unitType &&
-            expression.returnTargetSymbol.owner != unitGetInstance
-        ) {
+    private fun visitFunctionReturn(expression: IrReturn) {
+        val returnType = expression.returnTargetSymbol.owner.returnType(backendContext)
+        if (returnType == irBuiltIns.unitType && expression.returnTargetSymbol.owner != unitGetInstance) {
             generateStatement(expression.value)
         } else {
-            generateExpression(expression.value)
+            if (isGetUnitFunction) {
+                generateExpression(expression.value)
+            } else {
+                generateWithExpectedType(expression.value, returnType)
+            }
         }
 
         if (context.irFunction is IrConstructor) {
@@ -542,6 +560,88 @@ class BodyGenerator(
         body.buildInstr(WasmOp.RETURN)
     }
 
+    internal fun generateWithExpectedType(expression: IrExpression, expectedType: IrType) {
+        val actualType = expression.type
+
+        if (expectedType == wasmSymbols.voidType) {
+            generateStatement(expression)
+            return
+        }
+
+        if (expectedType.isUnit() && !actualType.isUnit()) {
+            generateStatement(expression)
+            body.buildGetUnit()
+            return
+        }
+
+        generateExpression(expression)
+        recoverToExpectedType(actualType = actualType, expectedType = expectedType)
+    }
+
+    //TODO: This method needed because of IR has type inconsistency. We need to discover why is it and fix
+    private fun recoverToExpectedType(actualType: IrType, expectedType: IrType) {
+        // TYPE -> NOTHING -> FALSE
+        if (expectedType.isNothing()) {
+            body.buildUnreachable()
+            return
+        }
+
+        // NOTHING -> TYPE -> TRUE
+        if (actualType.isNothing()) return
+
+        // NOTHING? -> TYPE? -> TRUE
+        if (actualType.isNullableNothing() && expectedType.isNullable()) return
+
+        val expectedClassErased = expectedType.getRuntimeClass(irBuiltIns)
+
+        // TYPE -> EXTERNAL -> TRUE
+        if (expectedClassErased.isExternal) return
+
+        val actualClassErased = actualType.getRuntimeClass(irBuiltIns)
+        val expectedTypeErased = expectedClassErased.defaultType
+        val actualTypeErased = actualClassErased.defaultType
+
+        // TYPE -> TYPE -> TRUE
+        if (expectedTypeErased == actualTypeErased) return
+
+        // NOT_NOTHING_TYPE -> NOTHING -> FALSE
+        if (expectedTypeErased.isNothing() && !actualTypeErased.isNothing()) {
+            body.buildUnreachable()
+            return
+        }
+
+        // TYPE -> BASE -> TRUE
+        if (actualClassErased.isSubclassOf(expectedClassErased)) {
+            return
+        }
+
+        val expectedIsPrimitive = expectedTypeErased.isPrimitiveType()
+        val actualIsPrimitive = actualTypeErased.isPrimitiveType()
+
+        // PRIMITIVE -> REF -> FALSE
+        // REF -> PRIMITIVE -> FALSE
+        if (expectedIsPrimitive != actualIsPrimitive) {
+            body.buildUnreachable()
+            return
+        }
+
+        // REF -> REF -> REF_CAST
+        if (!expectedIsPrimitive) {
+            generateRefCast(actualTypeErased, expectedTypeErased)
+        }
+    }
+
+    override fun visitReturn(expression: IrReturn) {
+        val nonLocalReturnSymbol = expression.returnTargetSymbol as? IrReturnableBlockSymbol
+        if (nonLocalReturnSymbol != null) {
+            generateWithExpectedType(expression.value, nonLocalReturnSymbol.owner.type)
+            body.buildBr(context.referenceNonLocalReturnLevel(nonLocalReturnSymbol))
+            body.buildUnreachable()
+        } else {
+            visitFunctionReturn(expression)
+        }
+    }
+
     override fun visitWhen(expression: IrWhen) {
         if (tryGenerateOptimisedWhen(expression, context.backendContext.wasmSymbols)) {
             return
@@ -555,17 +655,11 @@ class BodyGenerator(
             if (!isElseBranch(branch)) {
                 generateExpression(branch.condition)
                 body.buildIf(null, resultType)
-                generateExpression(branch.result)
-                if (expression.type == irBuiltIns.nothingType) {
-                    body.buildUnreachable()
-                }
+                generateWithExpectedType(branch.result, expression.type)
                 body.buildElse()
                 ifCount++
             } else {
-                generateExpression(branch.result)
-                if (expression.type == irBuiltIns.nothingType) {
-                    body.buildUnreachable()
-                }
+                generateWithExpectedType(branch.result, expression.type)
                 seenElse = true
                 break
             }
@@ -575,7 +669,15 @@ class BodyGenerator(
         // If it's not exhaustive it must be used as a statement (per kotlin spec) and so the result value of the last else will never be used.
         if (!seenElse && resultType != null) {
             assert(expression.type != irBuiltIns.nothingType)
-            generateDefaultInitializerForType(resultType, body)
+            if (expression.type.isUnit()) {
+                if (isGetUnitFunction) {
+                    generateDefaultInitializerForType(resultType, body)
+                } else {
+                    body.buildGetUnit()
+                }
+            } else {
+                body.buildUnreachable()
+            }
         }
 
         repeat(ifCount) { body.buildEnd() }

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

@@ -126,7 +126,11 @@ class DeclarationGenerator(
         }
 
         val exprGen = functionCodegenContext.bodyGen
-        val bodyBuilder = BodyGenerator(functionCodegenContext, hierarchyDisjointUnions)
+        val bodyBuilder = BodyGenerator(
+            context = functionCodegenContext,
+            hierarchyDisjointUnions = hierarchyDisjointUnions,
+            isGetUnitFunction = declaration == unitGetInstanceFunction
+        )
 
         require(declaration.body is IrBlockBody) { "Only IrBlockBody is supported" }
         declaration.body?.acceptVoid(bodyBuilder)

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

@@ -9,8 +9,8 @@ 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.types.isUnit
 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>)
@@ -58,7 +58,7 @@ internal fun BodyGenerator.tryGenerateOptimisedWhen(expression: IrWhen, symbols:
     if (minValue == maxValue) return false
 
     val selectorLocal = context.referenceLocal(SyntheticLocalType.TABLE_SWITCH_SELECTOR)
-    subject.accept(this, null)
+    generateExpression(subject)
     body.buildSetLocal(selectorLocal)
 
     val resultType = context.transformBlockResultType(expression.type)
@@ -71,7 +71,7 @@ internal fun BodyGenerator.tryGenerateOptimisedWhen(expression: IrWhen, symbols:
                 intBranches = intBranches,
                 elseExpression = elseExpression,
                 resultType = resultType,
-                nothingType = symbols.irBuiltIns.nothingType
+                expectedType = expression.type,
             )
         } else {
             createBinaryTable(selectorLocal, intBranches)
@@ -83,7 +83,7 @@ internal fun BodyGenerator.tryGenerateOptimisedWhen(expression: IrWhen, symbols:
                 elseExpression = elseExpression,
                 shift = 0,
                 brTable = intBranches.indices.toList(),
-                nothingType = symbols.irBuiltIns.nothingType
+                expectedType = expression.type,
             )
         }
     } else {
@@ -100,19 +100,12 @@ internal fun BodyGenerator.tryGenerateOptimisedWhen(expression: IrWhen, symbols:
             elseExpression = elseExpression,
             shift = minValue,
             brTable = brTable,
-            nothingType = symbols.irBuiltIns.nothingType
+            expectedType = expression.type,
         )
     }
     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) {
@@ -202,7 +195,7 @@ private fun BodyGenerator.createBinaryTable(
     intBranches: List<ExtractedWhenBranch<Int>>,
     elseExpression: IrExpression?,
     resultType: WasmType?,
-    nothingType: IrType,
+    expectedType: IrType,
 ) {
     val sortedCaseToBranchIndex = mutableListOf<Pair<Int, IrExpression>>()
     intBranches.mapTo(sortedCaseToBranchIndex) { branch -> branch.conditions[0].const.value to branch.expression }
@@ -210,26 +203,31 @@ private fun BodyGenerator.createBinaryTable(
 
     body.buildBlock("when_block", resultType) { currentBlock ->
         val thenBody = { result: IrExpression ->
-            result.safeAccept(this, nothingType)
+            generateWithExpectedType(result, expectedType)
             body.buildBr(currentBlock)
         }
-        val elseBody: () -> Unit = {
-            resultType?.let { generateDefaultInitializerForType(it, body) }
-        }
         createBinaryTable(
             selectorLocal = selectorLocal,
-            resultType = resultType,
+            resultType = null,
             sortedCases = sortedCaseToBranchIndex,
             fromIncl = 0,
             toExcl = sortedCaseToBranchIndex.size,
             thenBody = thenBody,
-            elseBody = elseBody
+            elseBody = { }
         )
 
-        if (resultType != null) {
-            body.buildDrop()
+        if (elseExpression != null) {
+            generateWithExpectedType(elseExpression, expectedType)
+        } else {
+            // default else block
+            if (resultType != null) {
+                if (expectedType.isUnit()) {
+                    body.buildGetUnit()
+                } else {
+                    body.buildUnreachable()
+                }
+            }
         }
-        elseExpression?.safeAccept(this, nothingType)
     }
 }
 
@@ -299,7 +297,7 @@ private fun BodyGenerator.genTableIntSwitch(
     elseExpression: IrExpression?,
     shift: Int,
     brTable: List<Int>,
-    nothingType: IrType,
+    expectedType: IrType,
 ) {
     val baseBlockIndex = body.numberOfNestedBlocks
     //expressions + else branch + br_table
@@ -324,16 +322,18 @@ private fun BodyGenerator.genTableIntSwitch(
         if (resultType != null) {
             body.buildDrop()
         }
-        expression.expression.safeAccept(this, nothingType)
+        generateWithExpectedType(expression.expression, expectedType)
 
         body.buildBr(baseBlockIndex + 1)
         body.buildEnd()
     }
 
-    if (resultType != null) {
-        body.buildDrop()
+    if (elseExpression != null) {
+        if (resultType != null) {
+            body.buildDrop()
+        }
+        generateWithExpectedType(elseExpression, expectedType)
     }
-    elseExpression?.safeAccept(this, nothingType)
 
     body.buildEnd()
     check(baseBlockIndex == body.numberOfNestedBlocks)

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

@@ -80,19 +80,17 @@ class UnitToVoidLowering(val context: WasmBackendContext) : FileLoweringPass, Ab
         if (expr.type == symbols.voidType || expr.type == builtIns.nothingType)
             return false
 
-        when (expr) {
-            is IrContainerExpression -> {
-                if (expr.statements.isEmpty())
-                    return false
-                return shouldVoidify(expr.statements.last())
-            }
-            is IrWhen ->
-                return true
-            is IrTry ->
-                return true
+        return when (expr) {
+            is IrContainerExpression ->
+                expr.statements.isNotEmpty() && expr !is IrReturnableBlock && shouldVoidify(expr.statements.last())
+
+            is IrWhen, is IrTry ->
+                true
+
             is IrTypeOperatorCall ->
-                return expr.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT
-            else -> return false
+                expr.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT
+
+            else -> false
         }
     }
 

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.backend.wasm.lower
 
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
@@ -22,11 +23,11 @@ import org.jetbrains.kotlin.ir.types.makeNotNull
  *
  * Wasm GC doesn't have a nullref type anymore.
  */
-class WasmNullCoercingLowering(context: JsCommonBackendContext) : AbstractValueUsageLowering(context) {
+class WasmNullCoercingLowering(private val contextx: WasmBackendContext) : AbstractValueUsageLowering(contextx) {
     override fun IrExpression.useExpressionAsType(actualType: IrType, expectedType: IrType): IrExpression =
-        if (actualType.makeNotNull().isNothing() && actualType.isNullable() && !expectedType.makeNotNull().isNothing())
+        if (actualType.makeNotNull().isNothing() && actualType.isNullable() && !expectedType.makeNotNull().isNothing() && expectedType != contextx.wasmSymbols.voidType)
             JsIrBuilder.buildComposite(
-                type,
+                expectedType,
                 listOf(this, IrConstImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, expectedType, IrConstKind.Null, null))
             )
         else

compiler/testData/codegen/box/inlineClasses/resultRunCatchingOrElse.kt

@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: WASM
-// WASM_MUTE_REASON: TYPE_ISSUES
 // WITH_STDLIB
 // WORKS_WHEN_VALUE_CLASS
 // LANGUAGE: +ValueClasses