[WASM] Lower Unit to Void in order to mark statements with no result

.idea/dictionaries/igor.xml

@@ -4,6 +4,9 @@
       <w>addr</w>
       <w>descr</w>
       <w>exprs</w>
+      <w>voidification</w>
+      <w>voidifier</w>
+      <w>voidify</w>
     </words>
   </dictionary>
 </component>

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

@@ -420,6 +420,12 @@ private val expressionBodyTransformer = makeWasmModulePhase(
     description = "Replace IrExpressionBody with IrBlockBody"
 )
 
+private val unitToVoidLowering = makeWasmModulePhase(
+    ::UnitToVoidLowering,
+    name = "UnitToVoidLowering",
+    description = "Replace some Unit's with Void's"
+)
+
 val wasmPhases = NamedCompilerPhase(
     name = "IrModuleLowering",
     description = "IR module lowering",
@@ -493,6 +499,7 @@ val wasmPhases = NamedCompilerPhase(
             fieldInitializersLoweringPhase then
             genericReturnTypeLowering then
             expressionBodyTransformer then
+            unitToVoidLowering then
 
             // Replace builtins before autoboxing
             builtInsLoweringPhase0 then

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -71,6 +72,10 @@ class WasmSymbols(
 
     val wasmUnreachable = getInternalFunction("wasm_unreachable")
 
+    val consumeAnyIntoVoid = getInternalFunction("consumeAnyIntoVoid")
+    val voidClass = getIrClass(FqName("kotlin.wasm.internal.Void"))
+    val voidType by lazy { voidClass.defaultType }
+
     val equalityFunctions = mapOf(
         context.irBuiltIns.booleanType to getInternalFunction("wasm_i32_eq"),
         context.irBuiltIns.byteType to getInternalFunction("wasm_i32_eq"),

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

@@ -38,7 +38,7 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
         buildInstr(WasmOp.GET_UNIT, WasmImmediate.FuncIdx(context.referenceFunction(unitGetInstance.symbol)))
     }
 
-    // Generates code for the given IR element and *always* leaves something on the stack
+    // Generates code for the given IR element. Leaves something on the stack unless expression was of the type Void.
     private fun generateExpression(elem: IrElement) {
         assert(elem is IrExpression || elem is IrVariable) { "Unsupported statement kind" }
 
@@ -49,12 +49,14 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
         }
     }
 
-    // Generates code for the given IR element but *never* leaves anything on the stack
+    // Generates code for the given IR element but *never* leaves anything on the stack.
     private fun generateStatement(statement: IrElement) {
         assert(statement is IrExpression || statement is IrVariable) { "Unsupported statement kind" }
 
         generateExpression(statement)
-        body.buildDrop()
+        if (statement is IrExpression && statement.type != wasmSymbols.voidType) {
+            body.buildDrop()
+        }
     }
 
     override fun visitElement(element: IrElement) {
@@ -408,9 +410,10 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
 
     override fun visitContainerExpression(expression: IrContainerExpression) {
         val statements = expression.statements
+
         if (statements.isEmpty()) {
-            assert(expression.type == irBuiltIns.unitType) { "Empty block with non-unit return type" }
-            body.buildGetUnit()
+            if (expression.type == irBuiltIns.unitType)
+                body.buildGetUnit()
             return
         }
 
@@ -419,6 +422,12 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
         }
 
         generateExpression(statements.last())
+
+        // 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()
+        }
     }
 
     override fun visitBreak(jump: IrBreak) {
@@ -553,14 +562,12 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
     override fun visitVariable(declaration: IrVariable) {
         context.defineLocal(declaration.symbol)
         if (declaration.initializer == null) {
-            body.buildGetUnit()
             return
         }
         val init = declaration.initializer!!
         generateExpression(init)
         val varName = context.referenceLocal(declaration.symbol)
         body.buildSetLocal(varName)
-        body.buildGetUnit()
     }
 
     // Return true if function is recognized as intrinsic.

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

@@ -21,6 +21,8 @@ class WasmTypeTransformer(
     val context: WasmBaseCodegenContext,
     val builtIns: IrBuiltIns
 ) {
+    val symbols = context.backendContext.wasmSymbols
+
     fun IrType.toWasmResultType(): WasmType? =
         when (this) {
             builtIns.unitType,
@@ -37,6 +39,9 @@ class WasmTypeTransformer(
             builtIns.nothingType ->
                 WasmUnreachableType
 
+            symbols.voidType ->
+                null
+
             else ->
                 toWasmValueType()
         }
@@ -85,6 +90,9 @@ class WasmTypeTransformer(
             builtIns.nothingType ->
                 WasmExternRef
 
+            symbols.voidType ->
+                error("Void type can't be used as a value")
+
             // this also handles builtIns.stringType
             else -> {
                 val klass = this.getClass()

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

@@ -0,0 +1,143 @@
+/*
+ * Copyright 2010-2021 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.lower.AbstractValueUsageTransformer
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.backend.js.JsStatementOrigins
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+
+// NOTE: This is an optional lowering.
+// In the wasm backend we treat Unit type as a regular object (except for the function return values). This means that all expressions
+// with Unit return type will actually produce a real Unit instance. However sometimes Unit type signifies that the result of the
+// expression will never be used and for such cases we want to be able to avoid generating redundant Unit objects.
+// In order to simplify the reasoning we introduce special type 'Void'. It means absence of any value, so there are never any real
+// objects of type Void. The only way to produce a Void type is by calling intrinsic 'consumeAnyIntoVoid'. In the wasm code it will turn
+// into a single drop instruction.
+// This lowering visits all statements with Unit return type and tries to replace it with Void type. In the result of
+// such statements it emits 'consumeAnyIntoVoid(previous_result)'. In other words it tries to propagate 'consumeAnyIntoVoid' as
+// deep as possible, changing statement return type to Void on each nesting level it crosses. For example, from this code:
+//    block_body {
+//        block: Unit {
+//            if (some_condition): Unit {
+//                foo(): Unit
+//            } else {
+//                IMPLICIT_COERCION_TO_UNIT(bar(): Int)
+//            }
+//            foo.bar = 10
+//        }
+//        return 1
+//    }
+// We will get this code:
+//    block_body {
+//        block: Void {
+//            if (some_condition): Void {
+//                consumeAnyIntoVoid(foo(): Unit): Void
+//            } else {
+//                consumeAnyIntoVoid(bar(): Int): Void
+//            }
+//            foo.bar = 10
+//        }
+//        return 1
+//    }
+// NOTE: In order to reduce the amount of new nodes in the IR we only do this transformation if it involves 'when' or 'try/catch' statements
+// because WASM backend handles other cases well enough.
+// As a further optimization we can directly mark 'call', 'set_value' and 'set_field' with Void return type and handle them as special
+// cases in the backend.
+
+class UnitToVoidLowering(val context: WasmBackendContext) : FileLoweringPass, AbstractValueUsageTransformer(context.irBuiltIns) {
+    val builtIns = context.irBuiltIns
+    val symbols = context.wasmSymbols
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun IrExpression.useAsStatement(): IrExpression {
+        if (type != builtIns.unitType)
+            return this
+
+        if (shouldVoidify(this))
+            return voidify(this) as IrExpression
+        return this
+    }
+
+    // Checks if this expression is an interesting target for the voidification.
+    private fun shouldVoidify(expr: IrStatement): Boolean {
+        if (expr !is IrExpression)
+            return true
+
+        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
+            is IrTypeOperatorCall ->
+                return expr.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT
+            else -> return false
+        }
+    }
+
+    // Takes expression of any type and unconditionally turns it into an expression of a void type.
+    private fun voidify(expr: IrStatement): IrStatement {
+        // Declarations are implicitly void
+        if (expr !is IrExpression)
+            return expr
+
+        // Don't voidify 'Nothing' because it can be casted to anything, including void.
+        if (expr.type == symbols.voidType || expr.type == builtIns.nothingType)
+            return expr
+
+        // Try to look through one of the known expression types.
+        when (expr) {
+            is IrContainerExpression -> {
+                expr.type = symbols.voidType
+                expr.statements.lastOrNull()?.let { last ->
+                    expr.statements[expr.statements.lastIndex] = voidify(last)
+                }
+                return expr
+            }
+            is IrWhen -> {
+                expr.type = symbols.voidType
+                expr.branches.forEachIndexed { i, branch ->
+                    expr.branches[i].result = voidify(branch.result) as IrExpression
+                }
+                return expr
+            }
+            is IrTry -> {
+                expr.type = symbols.voidType
+                expr.tryResult = voidify(expr.tryResult) as IrExpression
+                expr.catches.forEachIndexed { i, catch ->
+                    expr.catches[i].result = voidify(catch.result) as IrExpression
+                }
+                return expr
+            }
+            is IrTypeOperatorCall -> {
+                if (expr.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT)
+                    return voidify(expr.argument)
+            }
+        }
+
+        return IrCallImpl(expr.startOffset, expr.endOffset, symbols.voidType, symbols.consumeAnyIntoVoid, 0, 1).apply {
+            putValueArgument(0, expr)
+        }
+    }
+}

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

@@ -67,3 +67,12 @@ internal fun <T, R> boxIntrinsic(x: T): R =
 @ExcludedFromCodegen
 internal fun <T, R> unboxIntrinsic(x: T): R =
     implementedAsIntrinsic
+
+// Represents absence of a value. Should never be used as a real object. See UnitToVoidLowering.kt for more info.
+@ExcludedFromCodegen
+internal class Void private constructor()
+
+// This is the only way to introduce Void type.
+@WasmOp(WasmOp.DROP)
+internal fun consumeAnyIntoVoid(a: Any?): Void =
+    implementedAsIntrinsic