WASM: Stupidify unit handling

Now every IR instruction returns an actual Unit instance.
All Unit's are actual objects passed around except for function boundaries.

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -42890,6 +42890,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @Test
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @Test
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/AutoboxingTransformer.kt

@@ -192,21 +192,6 @@ class AutoboxingTransformer(context: JsCommonBackendContext) : AbstractValueUsag
     }
 }
 
-class UnitMaterializationPass(val context: JsCommonBackendContext) : AbstractValueUsageTransformer(context.irBuiltIns), BodyLoweringPass {
-    override fun lower(irBody: IrBody, container: IrDeclaration) {
-        irBody.transformChildrenVoid()
-    }
-
-    override fun IrExpression.useAsValue(value: IrValueDeclaration): IrExpression {
-        if (this.type == irBuiltIns.unitType && !this.isGetUnit(irBuiltIns)) {
-            val unitValue = JsIrBuilder.buildGetObjectValue(type, irBuiltIns.unitClass)
-            return JsIrBuilder.buildComposite(type, listOf(this, unitValue))
-        }
-        return this
-    }
-}
-
-
 private tailrec fun IrExpression.isGetUnit(irBuiltIns: IrBuiltIns): Boolean =
     when (this) {
         is IrContainerExpression ->

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

@@ -328,12 +328,6 @@ private val wasmNullSpecializationLowering = makeWasmModulePhase(
     description = "Specialize assigning Nothing? values to other types."
 )
 
-private val unitMaterializationPass = makeWasmModulePhase(
-    { context -> UnitMaterializationPass(context) },
-    name = "UnitLoweringPass",
-    description = "Materialize unit instances"
-)
-
 private val staticMembersLoweringPhase = makeWasmModulePhase(
     ::StaticMembersLowering,
     name = "StaticMembersLowering",
@@ -512,6 +506,5 @@ val wasmPhases = NamedCompilerPhase(
             wasmThrowDebugLoweringPhase then
             staticMembersLoweringPhase then
             wasmNullSpecializationLowering then
-            unitMaterializationPass then
             validateIrAfterLowering
 )

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

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.backend.wasm.WasmSymbols
 import org.jetbrains.kotlin.backend.wasm.utils.*
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.backend.js.utils.findUnitGetInstanceFunction
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
@@ -33,20 +33,24 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
     private val wasmSymbols: WasmSymbols = backendContext.wasmSymbols
     private val irBuiltIns: IrBuiltIns = backendContext.irBuiltIns
 
+    private val unitGetInstance by lazy { backendContext.findUnitGetInstanceFunction() }
+    fun WasmExpressionBuilder.buildGetUnit() {
+        buildCall(context.referenceFunction(unitGetInstance.symbol))
+    }
+
     override fun visitElement(element: IrElement) {
         error("Unexpected element of type ${element::class}")
     }
 
-    val unitGetInstance by lazy { backendContext.mapping.objectToGetInstanceFunction[irBuiltIns.unitClass.owner]!! }
-
-    override fun visitGetObjectValue(expression: IrGetObjectValue) {
-        require(expression.symbol == irBuiltIns.unitClass)
-        body.buildCall(context.referenceFunction(unitGetInstance.symbol))
-    }
-
     override fun visitTypeOperator(expression: IrTypeOperatorCall) {
-        require(expression.operator == IrTypeOperator.REINTERPRET_CAST) { "Other types of casts must be lowered" }
+        when (expression.operator) {
-        generateExpression(expression.argument)
+            IrTypeOperator.REINTERPRET_CAST -> generateExpression(expression.argument)
+            IrTypeOperator.IMPLICIT_COERCION_TO_UNIT -> {
+                statementToWasmInstruction(expression.argument)
+                body.buildGetUnit()
+            }
+            else -> assert(false) { "Other types of casts must be lowered" }
+        }
     }
 
     override fun <T> visitConst(expression: IrConst<T>) {
@@ -120,6 +124,8 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
             generateExpression(expression.value)
             body.buildSetGlobal(context.referenceGlobal(expression.symbol))
         }
+
+        body.buildGetUnit()
     }
 
     override fun visitGetValue(expression: IrGetValue) {
@@ -129,6 +135,7 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
     override fun visitSetValue(expression: IrSetValue) {
         generateExpression(expression.value)
         body.buildSetLocal(context.referenceLocal(expression.symbol))
+        body.buildGetUnit()
     }
 
     override fun visitCall(expression: IrCall) {
@@ -176,6 +183,7 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
 
         // Don't delegate constructors of Any to Any.
         if (klass.defaultType.isAny()) {
+            body.buildGetUnit()
             return
         }
 
@@ -208,6 +216,8 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
         val function: IrFunction = call.symbol.owner.realOverrideTarget
 
         if (tryToGenerateIntrinsicCall(call, function)) {
+            if (function.returnType == irBuiltIns.unitType)
+                body.buildGetUnit()
             return
         }
 
@@ -250,6 +260,10 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
                 body.buildRefCast()
             }
         }
+
+        // Unit types don't cross function boundaries
+        if (function.returnType == irBuiltIns.unitType && function.symbol != unitGetInstance.symbol)
+            body.buildGetUnit()
     }
 
     private fun generateTypeRTT(type: IrType) {
@@ -327,28 +341,32 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
 
     override fun visitContainerExpression(expression: IrContainerExpression) {
         val statements = expression.statements
-        if (statements.isEmpty()) return
+        if (statements.isEmpty()) {
+            assert(expression.type == irBuiltIns.unitType) { "Empty block with non-unit return type" }
+            body.buildGetUnit()
+            return
+        }
 
         statements.dropLast(1).forEach {
             statementToWasmInstruction(it)
         }
 
-        generateExpression(statements.last() as IrExpression)
+        generateExpression(statements.last())
     }
 
     override fun visitBreak(jump: IrBreak) {
+        assert(jump.type == irBuiltIns.nothingType)
         body.buildBr(context.referenceLoopLevel(jump.loop, LoopLabelType.BREAK))
     }
 
     override fun visitContinue(jump: IrContinue) {
+        assert(jump.type == irBuiltIns.nothingType)
         body.buildBr(context.referenceLoopLevel(jump.loop, LoopLabelType.CONTINUE))
     }
 
     override fun visitReturn(expression: IrReturn) {
-        if (
+        if (expression.returnTargetSymbol.owner.returnType(backendContext) == irBuiltIns.unitType &&
-            expression.value.type == irBuiltIns.unitType &&
+            expression.returnTargetSymbol.owner != unitGetInstance) {
-            expression.returnTargetSymbol.owner.returnType(backendContext) == irBuiltIns.unitType
-        ) {
             statementToWasmInstruction(expression.value)
         } else {
             generateExpression(expression.value)
@@ -372,6 +390,8 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
     override fun visitWhen(expression: IrWhen) {
         val resultType = context.transformBlockResultType(expression.type)
         var ifCount = 0
+        var seenElse = false
+
         for (branch in expression.branches) {
             if (!isElseBranch(branch)) {
                 generateExpression(branch.condition)
@@ -387,10 +407,18 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
                 if (expression.type == irBuiltIns.nothingType) {
                     body.buildUnreachable()
                 }
+                seenElse = true
                 break
             }
         }
 
+        // Always generate the last else to make verifier happy. If this when expression is exhaustive we will never reach the last else.
+        // 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)
+        }
+
         repeat(ifCount) { body.buildEnd() }
     }
 
@@ -420,6 +448,8 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
         body.buildBrIf(wasmLoop)
         body.buildEnd()
         body.buildEnd()
+
+        body.buildGetUnit()
     }
 
     override fun visitWhileLoop(loop: IrWhileLoop) {
@@ -449,68 +479,40 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
         body.buildBr(wasmLoop)
         body.buildEnd()
         body.buildEnd()
+
+        body.buildGetUnit()
     }
 
-    fun generateExpression(expression: IrExpression) {
+    // Generates code for the given IR element and *always* leaves something on the stack
-        expression.acceptVoid(this)
+    fun generateExpression(elem: IrElement) {
+        assert(elem is IrExpression || elem is IrVariable) { "Unsupported statement kind" }
 
-        if (expression.type == irBuiltIns.nothingType) {
+        elem.acceptVoid(this)
+
+        if (elem is IrExpression && elem.type == irBuiltIns.nothingType) {
             body.buildUnreachable()
         }
     }
 
-    fun statementToWasmInstruction(statement: IrStatement) {
+    // Generates code for the given IR element but *never* leaves anything on the stack
-        if (statement is IrVariable) {
+    fun statementToWasmInstruction(statement: IrElement) {
-            context.defineLocal(statement.symbol)
+        assert(statement is IrExpression || statement is IrVariable) { "Unsupported statement kind" }
-            val init = statement.initializer ?: return
+
-            generateExpression(init)
+        generateExpression(statement)
-            val varName = context.referenceLocal(statement.symbol)
+        body.buildDrop()
-            body.buildSetLocal(varName)
+    }
+
+    override fun visitVariable(declaration: IrVariable) {
+        context.defineLocal(declaration.symbol)
+        if (declaration.initializer == null) {
+            body.buildGetUnit()
             return
         }
-
+        val init = declaration.initializer!!
-        if (statement is IrContainerExpression) {
+        generateExpression(init)
-            statement.statements.forEach { it ->
+        val varName = context.referenceLocal(declaration.symbol)
-                statementToWasmInstruction(it)
+        body.buildSetLocal(varName)
-            }
+        body.buildGetUnit()
-        } else if (statement is IrWhen) {
-            var ifCount = 0
-            for (branch in statement.branches) {
-                if (!isElseBranch(branch)) {
-                    generateExpression(branch.condition)
-                    body.buildIf(label = null, resultType = null)
-                    statementToWasmInstruction(branch.result)
-                    body.buildElse()
-                    ifCount++
-                } else {
-                    statementToWasmInstruction(branch.result)
-                    break
-                }
-            }
-
-            repeat(ifCount) { body.buildEnd() }
-        } else {
-            generateExpression(statement as IrExpression)
-
-            var needDrop = true
-            if (statement.type == irBuiltIns.nothingType)
-                needDrop = false
-
-            if (statement is IrSetValue || statement is IrBreakContinue || statement is IrSetField || statement is IrLoop || statement is IrDelegatingConstructorCall) {
-                needDrop = false
-            }
-
-            if (statement is IrCall || statement is IrConstructorCall) {
-                val unitGetInstanceCall = statement is IrCall && statement.symbol.owner == unitGetInstance
-                if (statement.type == irBuiltIns.unitType && !unitGetInstanceCall) {
-                    needDrop = false
-                }
-            }
-
-            if (needDrop) {
-                body.buildInstr(WasmOp.DROP)
-            }
-        }
     }
 
     // Return true if function is recognized as intrinsic.

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

@@ -71,10 +71,12 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
         val irParameters = declaration.getEffectiveValueParameters()
         // TODO: Exported types should be transformed in a separate lowering by creating shim functions for each export.
         val resultType =
-            if (declaration.isExported(context.backendContext))
+            when {
-                context.transformExportedResultType(declaration.returnType)
+                declaration.isExported(context.backendContext) -> context.transformExportedResultType(declaration.returnType)
-            else
+                // Unit_getInstance returns true Unit reference instead of "void"
-                context.transformResultType(declaration.returnType)
+                declaration == unitGetInstanceFunction -> context.transformType(declaration.returnType)
+                else -> context.transformResultType(declaration.returnType)
+            }
         val wasmFunctionType =
             WasmFunctionType(
                 name = watName,
@@ -88,17 +90,10 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
                 },
                 resultTypes = listOfNotNull(
                     run {
-                        val type = if (declaration == unitGetInstanceFunction) {
+                        if (importedName != null && resultType is WasmRefNullType)
-                            // Unit_getInstance returns true Unit reference instead of "void"
-                            context.transformType(declaration.returnType)
-                        } else {
-                            resultType
-                        }
-
-                        if (importedName != null && type is WasmRefNullType)
                             WasmEqRef
                         else
-                            type
+                            resultType
                     }
                 )
             )
@@ -148,9 +143,6 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
                     bodyBuilder.statementToWasmInstruction(statement)
                 }
 
-            is IrExpressionBody ->
-                bodyBuilder.generateExpression(body.expression)
-
             else -> error("Unexpected body $body")
         }
 

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

@@ -33,9 +33,6 @@ class WasmTypeTransformer(
 
     fun IrType.toWasmBlockResultType(): WasmType? =
         when (this) {
-            builtIns.unitType ->
-                null
-
             // TODO: Lower blocks with Nothing type?
             builtIns.nothingType ->
                 WasmUnreachableType

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

@@ -38,7 +38,6 @@ class WasmTypeOperatorLowering(val context: WasmBackendContext) : FileLoweringPa
 class WasmBaseTypeOperatorTransformer(val context: WasmBackendContext) : IrElementTransformerVoidWithContext() {
     private val symbols = context.wasmSymbols
     private val builtIns = context.irBuiltIns
-    private val unitGetInstance = context.findUnitGetInstanceFunction()
 
     private lateinit var builder: DeclarationIrBuilder
 
@@ -49,7 +48,7 @@ class WasmBaseTypeOperatorTransformer(val context: WasmBackendContext) : IrEleme
         return when (expression.operator) {
             IrTypeOperator.IMPLICIT_CAST -> lowerImplicitCast(expression)
             IrTypeOperator.IMPLICIT_DYNAMIC_CAST -> error("Dynamic casts are not supported in Wasm backend")
-            IrTypeOperator.IMPLICIT_COERCION_TO_UNIT -> builder.irComposite(resultType = builtIns.unitType) { +expression.argument }
+            IrTypeOperator.IMPLICIT_COERCION_TO_UNIT -> expression
             IrTypeOperator.IMPLICIT_INTEGER_COERCION -> lowerIntegerCoercion(expression)
             IrTypeOperator.IMPLICIT_NOTNULL -> lowerImplicitCast(expression)
             IrTypeOperator.INSTANCEOF -> lowerInstanceOf(expression, inverted = false)

compiler/testData/codegen/box/unit/sillyThings.kt

@@ -0,0 +1,67 @@
+// Some silly things you can do with unit types.
+
+fun foo() {
+}
+
+fun zoot(): String {
+    return "str"
+}
+
+class Blumbs {
+    var t: Unit = Unit
+}
+
+fun varfoo(vararg t: Unit) {
+}
+
+fun foo1() {
+    if (zoot() == "str")
+        return
+    return
+}
+
+fun box(): String {
+    // Only check that this code can be compiled and verified.
+
+    1
+
+    {}
+
+    val tmp = Unit
+    {
+
+    }()
+
+    val tmp2 = {
+
+    }()
+
+    val tmp3 = Blumbs()
+    tmp3.t = Unit
+
+    val tmp4 = Blumbs()
+    tmp4.t = foo()
+
+    val units = arrayOf(Unit, Unit, Unit)
+    varfoo(Unit, *units, foo(), Unit)
+
+    val tmp5 = foo() as? Blumbs
+
+    val tmp6 = if (zoot() == "str") Unit else foo()
+
+    // From spec
+    val a /* : () -> Unit */ = {
+        if (true) 42
+        // CSB with no last expression
+        // Type is defined to be `kotlin.Unit`
+    }
+
+    val b: () -> Unit = {
+        if (true) 42 else -42
+        // CSB with last expression of type `kotlin.Int`
+        // Type is expected to be `kotlin.Unit`
+        // Coercion to kotlin.Unit applied
+    }
+
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -42734,6 +42734,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @Test
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @Test
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -42890,6 +42890,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @Test
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @Test
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -34462,6 +34462,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {
             runTest("compiler/testData/codegen/box/unit/unitClassObject.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -28667,6 +28667,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {
             runTest("compiler/testData/codegen/box/unit/unitClassObject.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -28073,6 +28073,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {
             runTest("compiler/testData/codegen/box/unit/unitClassObject.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -28013,6 +28013,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {
             runTest("compiler/testData/codegen/box/unit/unitClassObject.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -16095,6 +16095,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
             runTest("compiler/testData/codegen/box/unit/nullableUnitInWhen3.kt");
         }
 
+        @TestMetadata("sillyThings.kt")
+        public void testSillyThings() throws Exception {
+            runTest("compiler/testData/codegen/box/unit/sillyThings.kt");
+        }
+
         @TestMetadata("unitClassObject.kt")
         public void testUnitClassObject() throws Exception {
             runTest("compiler/testData/codegen/box/unit/unitClassObject.kt");