From 1db45faba215eeaba2a3bf0efe3da4bd3a9d4bab Mon Sep 17 00:00:00 2001
From: Igor Laevsky <igor.laevsky@jetbrains.com>
Date: Wed, 8 Sep 2021 20:18:09 +0300
Subject: [PATCH] [WASM] NFC. Fold instructions from the wasm expression
 builder

---
 .../src/org/jetbrains/kotlin/wasm/ir/Utils.kt     | 15 +++++++++++++++
 .../kotlin/wasm/ir/WasmExpressionBuilder.kt       |  6 ------
 .../kotlin/wasm/ir/WasmIrExpressionBuilder.kt     | 13 +++++++++----
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/Utils.kt b/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/Utils.kt
index e4fde56df06..4a70db2bfd3 100644
--- a/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/Utils.kt
+++ b/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/Utils.kt
@@ -30,3 +30,18 @@ fun WasmModule.calculateIds() {
     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)
+
+    return null
+}
+
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 4138b387c4b..14b56953518 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
@@ -9,8 +9,6 @@ abstract class WasmExpressionBuilder {
     abstract fun buildInstr(op: WasmOp, vararg immediates: WasmImmediate)
     abstract var numberOfNestedBlocks: Int
 
-    abstract val lastInstr: WasmOp?
-
     fun buildConstI32(value: Int) {
         buildInstr(WasmOp.I32_CONST, WasmImmediate.ConstI32(value))
     }
@@ -32,10 +30,6 @@ abstract class WasmExpressionBuilder {
     }
 
     fun buildUnreachable() {
-        // Unreachable is not needed
-        if (lastInstr == WasmOp.UNREACHABLE || lastInstr == WasmOp.RETURN)
-            return
-
         buildInstr(WasmOp.UNREACHABLE)
     }
 
diff --git a/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmIrExpressionBuilder.kt b/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmIrExpressionBuilder.kt
index 56ac1c9a674..f1e1f652ce0 100644
--- a/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmIrExpressionBuilder.kt
+++ b/wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/WasmIrExpressionBuilder.kt
@@ -10,7 +10,15 @@ class WasmIrExpressionBuilder(
 ) : WasmExpressionBuilder() {
 
     override fun buildInstr(op: WasmOp, vararg immediates: WasmImmediate) {
-        expression.add(WasmInstr(op, immediates.toList()))
+        val nextInstr = WasmInstr(op, immediates.toList())
+        val foldedInstrs = foldWasmInstructions(expression.lastOrNull(), nextInstr)
+
+        if (foldedInstrs == null) {
+            expression += nextInstr
+        } else {
+            expression.removeLastOrNull()
+            expression += foldedInstrs
+        }
     }
 
 
@@ -19,9 +27,6 @@ class WasmIrExpressionBuilder(
             assert(value >= 0) { "end without matching block" }
             field = value
         }
-
-    override val lastInstr: WasmOp?
-        get() = expression.lastOrNull()?.operator
 }
 
 inline fun buildWasmExpression(body: WasmExpressionBuilder.() -> Unit): MutableList<WasmInstr> {