[WASM] Remove unreachable code from compiled wasm

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

@@ -7,9 +7,9 @@ package org.jetbrains.kotlin.backend.wasm.ir2wasm
 
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.IrLoop
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.wasm.ir.WasmExpressionBuilder
-import org.jetbrains.kotlin.wasm.ir.WasmInstr
 import org.jetbrains.kotlin.wasm.ir.WasmLocal
 
 enum class LoopLabelType { BREAK, CONTINUE }
@@ -23,6 +23,9 @@ interface WasmFunctionCodegenContext : WasmBaseCodegenContext {
     fun referenceLocal(index: Int): WasmLocal
     fun referenceLocal(type: SyntheticLocalType): WasmLocal
 
+    fun defineNonLocalReturnLevel(block: IrReturnableBlockSymbol, level: Int)
+    fun referenceNonLocalReturnLevel(block: IrReturnableBlockSymbol): Int
+
     fun defineLoopLevel(irLoop: IrLoop, labelType: LoopLabelType, level: Int)
     fun referenceLoopLevel(irLoop: IrLoop, labelType: LoopLabelType): Int
 

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.expressions.IrLoop
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.wasm.ir.*
@@ -29,6 +30,7 @@ class WasmFunctionCodegenContextImpl(
     private val wasmLocals = LinkedHashMap<IrValueSymbol, WasmLocal>()
     private val wasmSyntheticLocals = LinkedHashMap<SyntheticLocalType, WasmLocal>()
     private val loopLevels = LinkedHashMap<Pair<IrLoop, LoopLabelType>, Int>()
+    private val nonLocalReturnLevels = LinkedHashMap<IrReturnableBlockSymbol, Int>()
 
     override fun defineLocal(irValueDeclaration: IrValueSymbol) {
         assert(irValueDeclaration !in wasmLocals) { "Redefinition of local" }
@@ -70,6 +72,14 @@ class WasmFunctionCodegenContextImpl(
         }
     }
 
+    override fun defineNonLocalReturnLevel(block: IrReturnableBlockSymbol, level: Int) {
+        nonLocalReturnLevels[block] = level
+    }
+
+    override fun referenceNonLocalReturnLevel(block: IrReturnableBlockSymbol): Int {
+        return nonLocalReturnLevels.getValue(block)
+    }
+
     override fun defineLoopLevel(irLoop: IrLoop, labelType: LoopLabelType, level: Int) {
         val loopKey = Pair(irLoop, labelType)
         assert(loopKey !in loopLevels) { "Redefinition of loop" }

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

@@ -29,29 +29,4 @@ fun WasmModule.calculateIds() {
     memories.calculateIds(startIndex = importedMemories.size)
     tables.calculateIds(startIndex = importedTables.size)
     tags.calculateIds(startIndex = importedTags.size)
-}
+}
-
-// This is used to perform simple peephole-like optimizations from the WasmExpressionBuilder.
-// Takes two adjacent instructions and returns an array with 0, 1 or 2 instructions to replace the original ones.
-// Returns null if no action is required.
-fun foldWasmInstructions(prev: WasmInstr?, next: WasmInstr): List<WasmInstr>? {
-    if (prev == null)
-        return null
-
-    // Unreachable is not needed after another unreachable or return
-    if (next.operator == WasmOp.UNREACHABLE && prev.operator in listOf(WasmOp.UNREACHABLE, WasmOp.RETURN))
-        return listOf(prev)
-
-    if (next.operator == WasmOp.DROP) {
-        // simple pure instruction + drop -> nothing
-        if (prev.operator == WasmOp.GET_UNIT || prev.operator == WasmOp.REF_NULL)
-            return listOf()
-
-        // return + drop -> return
-        if (prev.operator == WasmOp.RETURN)
-            return listOf(prev)
-    }
-
-    return null
-}
-

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

@@ -60,8 +60,8 @@ abstract class WasmExpressionBuilder {
     }
 
     fun buildBlock(resultType: WasmType? = null): Int {
-        buildInstr(WasmOp.BLOCK, WasmImmediate.BlockType.Value(resultType))
         numberOfNestedBlocks++
+        buildInstr(WasmOp.BLOCK, WasmImmediate.BlockType.Value(resultType))
         return numberOfNestedBlocks
     }
 

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

@@ -5,20 +5,86 @@
 
 package org.jetbrains.kotlin.wasm.ir
 
+private fun WasmOp.isOutCfgNode() = when (this) {
+    WasmOp.UNREACHABLE, WasmOp.RETURN, WasmOp.THROW, WasmOp.RETHROW, WasmOp.BR, WasmOp.BR_TABLE -> true
+    else -> false
+}
+
+private fun WasmOp.isInCfgNode() = when (this) {
+    WasmOp.ELSE, WasmOp.CATCH, WasmOp.CATCH_ALL -> true
+    else -> false
+}
+
+private fun WasmOp.pureStacklessInstruction() = when (this) {
+    WasmOp.GET_UNIT, WasmOp.REF_NULL, WasmOp.I32_CONST, WasmOp.I64_CONST, WasmOp.F32_CONST, WasmOp.F64_CONST, WasmOp.LOCAL_GET, WasmOp.GLOBAL_GET -> true
+    else -> false
+}
+
 class WasmIrExpressionBuilder(
     val expression: MutableList<WasmInstr>
 ) : WasmExpressionBuilder() {
 
-    override fun buildInstr(op: WasmOp, vararg immediates: WasmImmediate) {
+    private val lastInstr: WasmInstr?
-        val nextInstr = WasmInstr(op, immediates.toList())
+        get() = expression.lastOrNull()
-        val foldedInstrs = foldWasmInstructions(expression.lastOrNull(), nextInstr)
+    private var eatEverythingUntilLevel: Int? = null
 
-        if (foldedInstrs == null) {
+    private fun addInstruction(op: WasmOp, immediates: Array<out WasmImmediate>) {
-            expression += nextInstr
+        val newInstruction = WasmInstr(op, immediates.toList())
-        } else {
+        expression.add(newInstruction)
-            expression.removeLastOrNull()
+    }
-            expression += foldedInstrs
+
+    private fun getCurrentEatLevel(op: WasmOp): Int? {
+        val eatLevel = eatEverythingUntilLevel ?: return null
+        if (numberOfNestedBlocks == eatLevel && op.isInCfgNode()) {
+            eatEverythingUntilLevel = null
+            return null
         }
+        if (numberOfNestedBlocks < eatLevel) {
+            eatEverythingUntilLevel = null
+            return null
+        }
+        return eatLevel
+    }
+
+    override fun buildInstr(op: WasmOp, vararg immediates: WasmImmediate) {
+        val currentEatUntil = getCurrentEatLevel(op)
+        if (currentEatUntil != null) {
+            if (currentEatUntil <= numberOfNestedBlocks) return
+        } else {
+            if (op.isOutCfgNode()) {
+                eatEverythingUntilLevel = numberOfNestedBlocks
+                addInstruction(op, immediates)
+                return
+            }
+        }
+
+        val lastInstruction = lastInstr
+        if (lastInstruction == null) {
+            addInstruction(op, immediates)
+            return
+        }
+        val lastOperator = lastInstruction.operator
+
+        // droppable instructions + drop/unreachable -> nothing
+        if ((op == WasmOp.DROP || op == WasmOp.UNREACHABLE) && lastOperator.pureStacklessInstruction()) {
+            expression.removeLast()
+            return
+        }
+
+        // local set and local get for the same argument -> local tee
+        if (lastOperator == WasmOp.LOCAL_SET && op == WasmOp.LOCAL_GET) {
+            val localSetNumber = (lastInstruction.immediates.firstOrNull() as? WasmImmediate.LocalIdx)?.value
+            if (localSetNumber != null) {
+                val localGetNumber = (immediates.firstOrNull() as? WasmImmediate.LocalIdx)?.value
+                if (localGetNumber == localSetNumber) {
+                    expression.removeLast()
+                    addInstruction(WasmOp.LOCAL_TEE, immediates)
+                    return
+                }
+            }
+        }
+
+        addInstruction(op, immediates)
     }
 
     override var numberOfNestedBlocks: Int = 0
@@ -32,4 +98,4 @@ inline fun buildWasmExpression(body: WasmExpressionBuilder.() -> Unit): MutableL
     val res = mutableListOf<WasmInstr>()
     WasmIrExpressionBuilder(res).body()
     return res
-}
+}