WASM: Implement asm/disasm support for exception handling tags

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

@@ -15,6 +15,7 @@ class WasmModule(
     val importedMemories: List<WasmMemory> = emptyList(),
     val importedTables: List<WasmTable> = emptyList(),
     val importedGlobals: List<WasmGlobal> = emptyList(),
+    val importedTags: List<WasmTag> = emptyList(),
 
     val definedFunctions: List<WasmFunction.Defined> = emptyList(),
     val tables: List<WasmTable> = emptyList(),
@@ -22,6 +23,7 @@ class WasmModule(
     val globals: List<WasmGlobal> = emptyList(),
     val exports: List<WasmExport<*>> = emptyList(),
     val elements: List<WasmElement> = emptyList(),
+    val tags: List<WasmTag> = emptyList(),
 
     val startFunction: WasmFunction? = null,
 
@@ -103,6 +105,15 @@ class WasmElement(
     }
 }
 
+class WasmTag(
+    val type: WasmFunctionType,
+    val importPair: WasmImportPair? = null
+) : WasmNamedModuleField() {
+    init {
+        assert(type.resultTypes.isEmpty()) { "Must have empty return as per current spec" }
+    }
+}
+
 class WasmLocal(
     val id: Int,
     val name: String,
@@ -128,6 +139,7 @@ sealed class WasmExport<T : WasmNamedModuleField>(
     class Table(name: String, field: WasmTable) : WasmExport<WasmTable>(name, field, 0x1, "table")
     class Memory(name: String, field: WasmMemory) : WasmExport<WasmMemory>(name, field, 0x2, "memory")
     class Global(name: String, field: WasmGlobal) : WasmExport<WasmGlobal>(name, field, 0x3, "global")
+    class Tag(name: String, field: WasmTag) : WasmExport<WasmTag>(name, field, 0x4, "tag")
 }
 
 sealed class WasmTypeDeclaration(

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

@@ -27,6 +27,7 @@ enum class WasmImmediateKind {
     DATA_IDX,
     TABLE_IDX,
     LABEL_IDX,
+    TAG_IDX,
     LABEL_IDX_VECTOR,
     ELEM_IDX,
 
@@ -78,6 +79,7 @@ sealed class WasmImmediate {
     }
 
     class LabelIdx(val value: Int) : WasmImmediate()
+    class TagIdx(val value: Int) : WasmImmediate()
     class LabelIdxVector(val value: List<Int>) : WasmImmediate()
     class ElemIdx(val value: WasmElement) : WasmImmediate()
 
@@ -306,6 +308,12 @@ enum class WasmOp(
     RETURN("return", 0x0F),
     CALL("call", 0x10, FUNC_IDX),
     CALL_INDIRECT("call_indirect", 0x11, listOf(TYPE_IDX, TABLE_IDX)),
+    TRY("try", 0x06, BLOCK_TYPE),
+    CATCH("catch", 0x07, TAG_IDX),
+    CATCH_ALL("catch_all", 0x19),
+    DELEGATE("delegate", 0x18, LABEL_IDX),
+    THROW("throw", 0x08, TAG_IDX),
+    RETHROW("rethrow", 0x09, LABEL_IDX),
 
     // Parametric
     DROP("drop", 0x1A),

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

@@ -21,10 +21,12 @@ fun WasmModule.calculateIds() {
     importedMemories.calculateIds()
     importedTables.calculateIds()
     importedGlobals.calculateIds()
+    importedTags.calculateIds()
     elements.calculateIds()
 
     definedFunctions.calculateIds(startIndex = importedFunctions.size)
     globals.calculateIds(startIndex = importedGlobals.size)
     memories.calculateIds(startIndex = importedMemories.size)
     tables.calculateIds(startIndex = importedTables.size)
+    tags.calculateIds(startIndex = importedTags.size)
 }

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/convertors/WasmBinaryToIR.kt

@@ -22,6 +22,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
     val importedMemories: MutableList<WasmMemory> = mutableListOf()
     val importedTables: MutableList<WasmTable> = mutableListOf()
     val importedGlobals: MutableList<WasmGlobal> = mutableListOf()
+    val importedTags: MutableList<WasmTag> = mutableListOf()
 
     val definedFunctions: MutableList<WasmFunction.Defined> = mutableListOf()
     val table: MutableList<WasmTable> = mutableListOf()
@@ -32,7 +33,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
     val elements: MutableList<WasmElement> = mutableListOf()
     val data: MutableList<WasmData> = mutableListOf()
     var dataCount: Boolean = false
-
+    val tags: MutableList<WasmTag> = mutableListOf()
 
     private fun <T> byIdx(l1: List<T>, l2: List<T>, index: Int): T {
         if (index < l1.size)
@@ -45,6 +46,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
     private fun elemByIdx(index: Int) = elements[index]
     private fun tableByIdx(index: Int) = byIdx(importedTables, table, index)
     private fun globalByIdx(index: Int) = byIdx(importedGlobals, globals, index)
+    private fun tagByIdx(index: Int) = byIdx(importedTags, tags, index)
 
     fun parseModule(): WasmModule {
         if (b.readUInt32() != 0x6d736100u)
@@ -119,6 +121,11 @@ class WasmBinaryToIR(val b: MyByteReader) {
                                         ).also { importsInOrder.add(it) }
                                     )
                                 }
+                                4 -> {
+                                    val tag = readTag(importPair)
+                                    importedTags.add(tag)
+                                    importsInOrder.add(tag)
+                                }
                                 else -> error(
                                     "Unsupported import kind $kind"
                                 )
@@ -162,6 +169,13 @@ class WasmBinaryToIR(val b: MyByteReader) {
                         }
                     }
 
+                    // Tag section
+                    13 -> {
+                        forEachVectorElement {
+                            tags.add(readTag())
+                        }
+                    }
+
                     // Globals section
                     6 -> {
                         forEachVectorElement {
@@ -190,6 +204,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
                                     1 -> WasmExport.Table(name, tableByIdx(index))
                                     2 -> WasmExport.Memory(name, memoryByIdx(index))
                                     3 -> WasmExport.Global(name, globalByIdx(index))
+                                    4 -> WasmExport.Tag(name, tagByIdx(index))
                                     else -> error("Invalid export kind $kind")
                                 }
                             )
@@ -319,6 +334,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
             importedMemories = importedMemories,
             importedTables = importedTables,
             importedGlobals = importedGlobals,
+            importedTags = importedTags,
             definedFunctions = definedFunctions,
             tables = table,
             memories = memory,
@@ -327,7 +343,8 @@ class WasmBinaryToIR(val b: MyByteReader) {
             startFunction = startFunction,
             elements = elements,
             data = data,
-            dataCount = dataCount
+            dataCount = dataCount,
+            tags = tags
         ).also {
             it.calculateIds()
         }
@@ -341,6 +358,13 @@ class WasmBinaryToIR(val b: MyByteReader) {
         )
     }
 
+    private fun readTag(importPair: WasmImportPair? = null): WasmTag {
+        val attribute = b.readByte()
+        check(attribute.toInt() == 0) { "as per spec" }
+        val type = functionTypes[b.readVarUInt32AsInt()]
+        return WasmTag(type, importPair)
+    }
+
     private fun readExpression(): MutableList<WasmInstr> =
         mutableListOf<WasmInstr>().also { readExpression(it) }
 
@@ -404,6 +428,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
                 WasmImmediateKind.DATA_IDX -> WasmImmediate.DataIdx(b.readVarUInt32AsInt())
                 WasmImmediateKind.TABLE_IDX -> WasmImmediate.TableIdx(b.readVarUInt32AsInt())
                 WasmImmediateKind.LABEL_IDX -> WasmImmediate.LabelIdx(b.readVarUInt32AsInt())
+                WasmImmediateKind.TAG_IDX -> WasmImmediate.TagIdx(b.readVarUInt32AsInt())
                 WasmImmediateKind.LABEL_IDX_VECTOR -> WasmImmediate.LabelIdxVector(mapVector { b.readVarUInt32AsInt() })
                 WasmImmediateKind.ELEM_IDX -> WasmImmediate.ElemIdx(elemByIdx(b.readVarUInt32AsInt()))
                 WasmImmediateKind.VAL_TYPE_VECTOR -> WasmImmediate.ValTypeVector(mapVector { readValueType() })

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/convertors/WasmIrToBinary.kt

@@ -41,6 +41,7 @@ class WasmIrToBinary(outputStream: OutputStream, val module: WasmModule) {
                         is WasmMemory -> appendMemory(it)
                         is WasmTable -> appendTable(it)
                         is WasmGlobal -> appendGlobal(it)
+                        is WasmTag -> appendTag(it)
                         else -> error("Unknown import kind ${it::class}")
                     }
                 }
@@ -64,7 +65,11 @@ class WasmIrToBinary(outputStream: OutputStream, val module: WasmModule) {
                 memories.forEach { appendMemory(it) }
             }
 
-            // Good
+            // tag section
+            appendSection(13u) {
+                appendVectorSize(tags.size)
+                tags.forEach { appendTag(it) }
+            }
 
             appendSection(6u) {
                 appendVectorSize(globals.size)
@@ -141,6 +146,7 @@ class WasmIrToBinary(outputStream: OutputStream, val module: WasmModule) {
             is WasmImmediate.DataIdx -> b.writeVarUInt32(x.value)
             is WasmImmediate.TableIdx -> b.writeVarUInt32(x.value.owner)
             is WasmImmediate.LabelIdx -> b.writeVarUInt32(x.value)
+            is WasmImmediate.TagIdx -> b.writeVarUInt32(x.value)
             is WasmImmediate.LabelIdxVector -> {
                 b.writeVarUInt32(x.value.size)
                 for (target in x.value) {
@@ -271,6 +277,18 @@ class WasmIrToBinary(outputStream: OutputStream, val module: WasmModule) {
         appendExpr(c.init)
     }
 
+    private fun appendTag(t: WasmTag) {
+        if (t.importPair != null) {
+            b.writeString(t.importPair.moduleName)
+            b.writeString(t.importPair.declarationName)
+            b.writeByte(4)
+            return
+        }
+        b.writeByte(0) // attribute
+        assert(t.type.id != null) { "Unlinked tag id" }
+        b.writeVarUInt32(t.type.id!!)
+    }
+
     private fun appendExpr(expr: Iterable<WasmInstr>) {
         expr.forEach { appendInstr(it) }
         appendInstr(WasmInstr(WasmOp.END))

wasm/wasm.ir/src/org/jetbrains/kotlin/wasm/ir/convertors/WasmIrToText.kt

@@ -60,13 +60,13 @@ class WasmIrToText : SExpressionBuilder() {
 
     private fun appendInstr(wasmInstr: WasmInstr) {
         val op = wasmInstr.operator
-        if (op == WasmOp.END || op == WasmOp.ELSE)
+        if (op == WasmOp.END || op == WasmOp.ELSE || op == WasmOp.CATCH)
             indent--
 
         newLine()
         stringBuilder.append(wasmInstr.operator.mnemonic)
 
-        if (op == WasmOp.BLOCK || op == WasmOp.LOOP || op == WasmOp.IF || op == WasmOp.ELSE)
+        if (op == WasmOp.BLOCK || op == WasmOp.LOOP || op == WasmOp.IF || op == WasmOp.ELSE || op == WasmOp.CATCH || op == WasmOp.TRY)
             indent++
 
         if (wasmInstr.operator in setOf(WasmOp.CALL_INDIRECT, WasmOp.TABLE_INIT)) {
@@ -100,6 +100,7 @@ class WasmIrToText : SExpressionBuilder() {
             is WasmImmediate.DataIdx -> appendElement(x.value.toString())
             is WasmImmediate.TableIdx -> appendElement(x.value.toString())
             is WasmImmediate.LabelIdx -> appendElement(x.value.toString())
+            is WasmImmediate.TagIdx -> appendElement(x.value.toString())
             is WasmImmediate.LabelIdxVector ->
                 x.value.forEach { appendElement(it.toString()) }
 
@@ -212,6 +213,7 @@ class WasmIrToText : SExpressionBuilder() {
                         is WasmMemory -> appendMemory(it)
                         is WasmTable -> appendTable(it)
                         is WasmGlobal -> appendGlobal(it)
+                        is WasmTag -> appendTag(it)
                         else -> error("Unknown import kind ${it::class}")
                     }
                 }
@@ -223,6 +225,7 @@ class WasmIrToText : SExpressionBuilder() {
                 elements.forEach { appendWasmElement(it) }
                 startFunction?.let { appendStartFunction(it) }
                 data.forEach { appendData(it) }
+                tags.forEach { appendTag(it) }
             }
         }
     }
@@ -400,6 +403,19 @@ class WasmIrToText : SExpressionBuilder() {
         }
     }
 
+    private fun appendTag(wasmTag: WasmTag) {
+        newLineList("tag") {
+            appendModuleFieldReference(wasmTag)
+
+            wasmTag.importPair?.appendImportPair()
+
+            sameLineList("param") {
+                wasmTag.type.parameterTypes.forEach { appendType(it) }
+            }
+            assert(wasmTag.type.resultTypes.isEmpty()) { "must be as per spec" }
+        }
+    }
+
     private fun appendLocal(local: WasmLocal) {
         newLineList(if (local.isParameter) "param" else "local") {
             appendLocalReference(local)
@@ -485,6 +501,7 @@ class WasmIrToText : SExpressionBuilder() {
             is WasmElement -> "elem"
             is WasmGlobal -> "g"
             is WasmTypeDeclaration -> "type"
+            is WasmTag -> "tag"
         }
 
         appendElement("\$${sanitizeWatIdentifier(field.name)}___${indexSpaceKind}_$id")