[Wasm] Improve interface method dispatch

- Use typed Wasm tables for each interface method to avoid runtime function type check - Use linear search by implemented interface rather than by individual virtual function signature
2020-12-13 20:55:40 +03:00
parent e7dc199ad7
commit b6ad1584c9
33 changed files with 318 additions and 172 deletions
@@ -126,7 +126,7 @@ class WasmSymbols(
    val wasmInterfaceId = getInternalFunction("wasmInterfaceId")

    val getVirtualMethodId = getInternalFunction("getVirtualMethodId")
-    val getInterfaceMethodId = getInternalFunction("getInterfaceMethodId")
+    val getInterfaceImplId = getInternalFunction("getInterfaceImplId")

    val isSubClass = getInternalFunction("isSubClass")
    val isInterface = getInternalFunction("isInterface")
@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.backend.common.ir.isOverridable
 import org.jetbrains.kotlin.backend.common.ir.returnType
 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.backend.wasm.WasmSymbols
-import org.jetbrains.kotlin.backend.wasm.lower.wasmSignature
 import org.jetbrains.kotlin.backend.wasm.utils.*
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
@@ -219,16 +218,19 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
                generateExpression(call.dispatchReceiver!!)
                body.buildConstI32(vfSlot)
                body.buildCall(context.referenceFunction(wasmSymbols.getVirtualMethodId))
+                body.buildCallIndirect(
+                    symbol = context.referenceFunctionType(function.symbol)
+                )
            } else {
-                val signatureId = context.referenceSignatureId(function.wasmSignature(backendContext.irBuiltIns))
                generateExpression(call.dispatchReceiver!!)
-                body.buildConstI32Symbol(signatureId)
-                body.buildCall(context.referenceFunction(wasmSymbols.getInterfaceMethodId))
+                body.buildConstI32Symbol(context.referenceInterfaceId(klass.symbol))
+                body.buildCall(context.referenceFunction(wasmSymbols.getInterfaceImplId))
+                body.buildCallIndirect(
+                    tableIdx = WasmSymbolIntWrapper(context.referenceInterfaceTable(function.symbol)),
+                    symbol = context.referenceFunctionType(function.symbol)
+                )
            }

-            body.buildCallIndirect(
-                symbol = context.referenceFunctionType(function.symbol)
-            )
        } else {
            // Static function call
            body.buildCall(context.referenceFunction(function.symbol))
@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.backend.wasm.ir2wasm
 import org.jetbrains.kotlin.backend.common.ir.isOverridableOrOverrides
 import org.jetbrains.kotlin.backend.wasm.lower.WasmSignature
 import org.jetbrains.kotlin.backend.wasm.lower.wasmSignature
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
@@ -78,6 +80,23 @@ class ClassMetadata(
        virtualMethods.map { it.signature }.toSet()
 }

+class InterfaceMetadata(
+    val iface: IrClass,
+    irBuiltIns: IrBuiltIns
+) {
+    val methods: List<VirtualMethodMetadata> =
+        iface.declarations
+            .filterIsInstance<IrSimpleFunction>()
+            .filter { !it.isFakeOverride && it.visibility != DescriptorVisibilities.PRIVATE && it.modality != Modality.FINAL }
+            .map {
+                VirtualMethodMetadata(
+                    it,
+                    it.wasmSignature(irBuiltIns)
+                )
+            }
+}
+
+
 class VirtualMethodMetadata(
    val function: IrSimpleFunction,
    val signature: WasmSignature
@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
 import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
@@ -186,6 +187,14 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
        }

        if (declaration.isInterface) {
+            val metadata = InterfaceMetadata(declaration, irBuiltIns)
+            for (method in metadata.methods) {
+                val methodSymbol = method.function.symbol
+                val table = WasmTable(
+                    elementType = WasmRefNullType(WasmHeapType.Type(context.referenceFunctionType(methodSymbol)))
+                )
+                context.defineInterfaceMethodTable(methodSymbol, table)
+            }
            context.registerInterface(symbol)
        } else {
            val nameStr = declaration.fqNameWhenAvailable.toString()
@@ -235,6 +244,28 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
            context.defineRTT(symbol, rtt)
            context.registerClass(symbol)
            context.generateTypeInfo(symbol, binaryDataStruct(metadata))
+
+            // New type info model
+            if (declaration.modality != Modality.ABSTRACT) {
+                context.generateInterfaceTable(symbol, interfaceTable(metadata))
+                for (i in metadata.interfaces) {
+                    val interfaceImplementation = InterfaceImplementation(i.symbol, declaration.symbol)
+                    // TODO: Cache it
+                    val interfaceMetadata = InterfaceMetadata(i, irBuiltIns)
+                    val table = interfaceMetadata.methods.associate { method ->
+                        val classMethod: VirtualMethodMetadata =
+                            metadata.virtualMethods
+                                .find { it.signature == method.signature }  // TODO: Use map
+                                ?: error("Cannot find class implementation of method ${method.signature} in class ${declaration.fqNameWhenAvailable}")
+
+                        method.function.symbol as IrFunctionSymbol to context.referenceFunction(classMethod.function.symbol)
+                    }
+                    context.registerInterfaceImplementationMethod(
+                        interfaceImplementation,
+                        table
+                    )
+                }
+            }
        }

        for (member in declaration.declarations) {
@@ -252,20 +283,6 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
        val superTypeField =
            ConstantDataIntField("Super class", superClassSymbol)

-        val interfacesArray = ConstantDataIntArray(
-            "data",
-            classMetadata.interfaces.map { context.referenceInterfaceId(it.symbol) }
-        )
-        val interfacesArraySize = ConstantDataIntField(
-            "size",
-            interfacesArray.value.size
-        )
-
-        val implementedInterfacesArrayWithSize = ConstantDataStruct(
-            "Implemented interfaces array",
-            listOf(interfacesArraySize, interfacesArray)
-        )
-
        val vtableSizeField = ConstantDataIntField(
            "V-table length",
            classMetadata.virtualMethods.size
@@ -282,29 +299,48 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
            }
        )

-        val signaturesArray = ConstantDataIntArray(
-            "Signatures",
-            classMetadata.virtualMethods.map {
-                if (it.function.modality == Modality.ABSTRACT) {
-                    WasmSymbol(invalidIndex)
-                } else {
-                    context.referenceSignatureId(it.signature)
-                }
-            }
+        val interfaceTablePtr = ConstantDataIntField(
+            "interfaceTablePtr",
+            context.referenceInterfaceTableAddress(classMetadata.klass.symbol)
        )

        return ConstantDataStruct(
            "Class TypeInfo: ${classMetadata.klass.fqNameWhenAvailable} ",
            listOf(
                superTypeField,
+                interfaceTablePtr,
                vtableSizeField,
                vtableArray,
-                signaturesArray,
-                implementedInterfacesArrayWithSize,
            )
        )
    }

+
+    private fun interfaceTable(classMetadata: ClassMetadata): ConstantDataStruct {
+        val interfaces = classMetadata.interfaces
+        val size = ConstantDataIntField("size", interfaces.size)
+        val interfaceIds = ConstantDataIntArray(
+            "interfaceIds",
+            interfaces.map { context.referenceInterfaceId(it.symbol) },
+        )
+        val interfaceImplementationIds = ConstantDataIntArray(
+            "interfaceImplementationId",
+            interfaces.map {
+                context.referenceInterfaceImplementationId(InterfaceImplementation(it.symbol, classMetadata.klass.symbol))
+            },
+        )
+
+        return ConstantDataStruct(
+            "Class interface table: ${classMetadata.klass.fqNameWhenAvailable} ",
+            listOf(
+                size,
+                interfaceIds,
+                interfaceImplementationIds,
+            )
+        )
+    }
+
+
    override fun visitField(declaration: IrField) {
        // Member fields are generated as part of struct type
        if (!declaration.isStatic) return
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.backend.wasm.ir2wasm

 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
-import org.jetbrains.kotlin.wasm.ir.*
 import org.jetbrains.kotlin.backend.wasm.lower.WasmSignature
 import org.jetbrains.kotlin.ir.declarations.IrField
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
@@ -15,6 +14,7 @@ import org.jetbrains.kotlin.ir.symbols.IrFieldSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.wasm.ir.*

 interface WasmBaseCodegenContext {
    val backendContext: WasmBackendContext
@@ -31,6 +31,8 @@ interface WasmBaseCodegenContext {

    fun referenceSignatureId(signature: WasmSignature): WasmSymbol<Int>

+    fun referenceInterfaceTable(irFunction: IrFunctionSymbol): WasmSymbol<WasmTable>
+
    fun referenceStringLiteral(string: String): WasmSymbol<Int>

    fun transformType(irType: IrType): WasmType
@@ -5,13 +5,13 @@

 package org.jetbrains.kotlin.backend.wasm.ir2wasm

-import org.jetbrains.kotlin.wasm.ir.*
 import org.jetbrains.kotlin.backend.wasm.lower.WasmSignature
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
 import org.jetbrains.kotlin.ir.declarations.IrExternalPackageFragment
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.getPackageFragment
+import org.jetbrains.kotlin.wasm.ir.*

 class WasmCompiledModuleFragment {
    val functions =
@@ -44,6 +44,27 @@ class WasmCompiledModuleFragment {

    val typeInfo =
        ReferencableAndDefinable<IrClassSymbol, ConstantDataElement>()
+
+    // Wasm table for each method of each interface.
+    val interfaceMethodTables =
+        ReferencableAndDefinable<IrFunctionSymbol, WasmTable>()
+
+    // Defined class interface tables
+    val definedClassITableData =
+        ReferencableAndDefinable<IrClassSymbol, ConstantDataElement>()
+
+    // Address of class interface table in linear memory
+    val referencedClassITableAddresses =
+        ReferencableElements<IrClassSymbol, Int>()
+
+    // Sequential number of an implementation (class, object, etc.) for a particular interface
+    // Used as index in table for interface method dispatch
+    val referencedInterfaceImplementationId =
+        ReferencableElements<InterfaceImplementation, Int>()
+
+    val interfaceImplementationsMethods =
+        LinkedHashMap<InterfaceImplementation, Map<IrFunctionSymbol, WasmSymbol<WasmFunction>>>()
+
    val exports = mutableListOf<WasmExport<*>>()

    class JsCodeSnippet(val importName: String, val jsCode: String)
@@ -92,27 +113,42 @@ class WasmCompiledModuleFragment {
        bind(runtimeTypes.unbound, runtimeTypes.defined)

        val klassIds = mutableMapOf<IrClassSymbol, Int>()
-        var classId = 0
+        var currentDataSectionAddress = 0
        for (typeInfoElement in typeInfo.elements) {
            val ir = typeInfo.wasmToIr.getValue(typeInfoElement)
-            klassIds[ir] = classId
-            classId += typeInfoElement.sizeInBytes
+            klassIds[ir] = currentDataSectionAddress
+            currentDataSectionAddress += typeInfoElement.sizeInBytes
+        }
+
+        val interfaceTableAddresses = mutableMapOf<IrClassSymbol, Int>()
+        for (typeInfoElement in definedClassITableData.elements) {
+            val ir = definedClassITableData.wasmToIr.getValue(typeInfoElement)
+            interfaceTableAddresses[ir] = currentDataSectionAddress
+            currentDataSectionAddress += typeInfoElement.sizeInBytes
        }

        bind(classIds.unbound, klassIds)
+        bind(referencedClassITableAddresses.unbound, interfaceTableAddresses)
        bindIndices(virtualFunctionId.unbound, virtualFunctions)
        bindIndices(signatureId.unbound, signatures.toList())
        bindIndices(interfaceId.unbound, interfaces)
        bindIndices(stringLiteralId.unbound, stringLiterals)

-        val data = typeInfo.elements.map {
-            val ir = typeInfo.wasmToIr.getValue(it)
-            val id = klassIds.getValue(ir)
-            val offset = mutableListOf<WasmInstr>()
-            WasmIrExpressionBuilder(offset).buildConstI32(id)
-            WasmData(WasmDataMode.Active(0, offset), it.toBytes())
+        val interfaceImplementationIds = mutableMapOf<InterfaceImplementation, Int>()
+        val numberOfInterfaceImpls = mutableMapOf<IrClassSymbol, Int>()
+        for (interfaceImplementation in interfaceImplementationsMethods.keys) {
+            val prev = numberOfInterfaceImpls.getOrPut(interfaceImplementation.irInterface) { 0 }
+            interfaceImplementationIds[interfaceImplementation] = prev
+            numberOfInterfaceImpls[interfaceImplementation.irInterface] = prev + 1
        }

+        bind(referencedInterfaceImplementationId.unbound, interfaceImplementationIds)
+        bind(interfaceMethodTables.unbound, interfaceMethodTables.defined)
+
+        val data =
+            typeInfo.buildData(address = { klassIds.getValue(it) }) +
+                    definedClassITableData.buildData(address = { interfaceTableAddresses.getValue(it) })
+
        val logTypeInfo = false
        if (logTypeInfo) {
            println("Signatures: ")
@@ -151,7 +187,38 @@ class WasmCompiledModuleFragment {
            WasmElement.Mode.Active(table, offsetExpr)
        )

-        val typeInfoSize = classId
+        val interfaceTableElementsLists = interfaceMethodTables.defined.keys.associateWith {
+            mutableMapOf<Int, WasmSymbol<WasmFunction>>()
+        }
+
+        interfaceImplementationIds.forEach { ii: InterfaceImplementation, implId: Int ->
+            for ((interfaceFunction: IrFunctionSymbol, wasmFunction: WasmSymbol<WasmFunction>) in interfaceImplementationsMethods[ii]!!) {
+                interfaceTableElementsLists[interfaceFunction]!![implId] = wasmFunction
+            }
+        }
+
+        val interfaceTableElements = interfaceTableElementsLists.map { (interfaceFunction, methods) ->
+            val type = interfaceMethodTables.defined[interfaceFunction]!!.elementType
+            val functions = MutableList(methods.size) { idx ->
+                val wasmFunc = methods[idx]!!
+                val expression = buildWasmExpression {
+                    buildInstr(WasmOp.REF_FUNC, WasmImmediate.FuncIdx(wasmFunc))
+                }
+                WasmTable.Value.Expression(expression)
+            }
+            WasmElement(
+                type,
+                values = functions,
+                WasmElement.Mode.Active(interfaceMethodTables.defined[interfaceFunction]!!, offsetExpr)
+            )
+        }
+
+        interfaceMethodTables.defined.forEach { (function, table) ->
+            val size = interfaceTableElementsLists[function]!!.size.toUInt()
+            table.limits = WasmLimits(size, size)
+        }
+
+        val typeInfoSize = currentDataSectionAddress
        val memorySizeInPages = (typeInfoSize / 65_536) + 1
        val memory = WasmMemory(WasmLimits(memorySizeInPages.toUInt(), memorySizeInPages.toUInt()))

@@ -166,12 +233,12 @@ class WasmCompiledModuleFragment {
            importsInOrder = importedFunctions,
            importedFunctions = importedFunctions,
            definedFunctions = functions.elements.filterIsInstance<WasmFunction.Defined>(),
-            tables = listOf(table),
+            tables = listOf(table) + interfaceMethodTables.elements,
            memories = listOf(memory),
            globals = globals.elements + sortedRttGlobals,
            exports = exports,
            startFunction = startFunction!!,
-            elements = listOf(elements),
+            elements = listOf(elements) + interfaceTableElements,
            data = data
        )
        module.calculateIds()
@@ -207,4 +274,19 @@ fun <IrSymbolType> bindIndices(
            error("Can't link symbol with indices ${irSymbolDebugDump(irSymbol)}")
        wasmSymbol.bind(index)
    }
-}
+}
+
+inline fun WasmCompiledModuleFragment.ReferencableAndDefinable<IrClassSymbol, ConstantDataElement>.buildData(address: (IrClassSymbol) -> Int): List<WasmData> {
+    return elements.map {
+        val id = address(wasmToIr.getValue(it))
+        val offset = mutableListOf<WasmInstr>()
+        WasmIrExpressionBuilder(offset).buildConstI32(id)
+        WasmData(WasmDataMode.Active(0, offset), it.toBytes())
+    }
+}
+
+
+data class InterfaceImplementation(
+    val irInterface: IrClassSymbol,
+    val irClass: IrClassSymbol
+)
@@ -5,13 +5,11 @@

 package org.jetbrains.kotlin.backend.wasm.ir2wasm

-import org.jetbrains.kotlin.wasm.ir.*
-import org.jetbrains.kotlin.backend.wasm.ir2wasm.ConstantDataElement
-import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmBaseCodegenContext
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFieldSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.wasm.ir.*

 /**
 * Interface for generating WebAssembly module.
@@ -22,6 +20,7 @@ interface WasmModuleCodegenContext : WasmBaseCodegenContext {
    fun defineGcType(irClass: IrClassSymbol, wasmType: WasmTypeDeclaration)
    fun defineRTT(irClass: IrClassSymbol, wasmGlobal: WasmGlobal)
    fun defineFunctionType(irFunction: IrFunctionSymbol, wasmFunctionType: WasmFunctionType)
+    fun defineInterfaceMethodTable(irFunction: IrFunctionSymbol, wasmTable: WasmTable)
    fun addJsFun(importName: String, jsCode: String)

    fun setStartFunction(wasmFunction: WasmFunction)
@@ -32,4 +31,13 @@ interface WasmModuleCodegenContext : WasmBaseCodegenContext {
    fun registerClass(irClass: IrClassSymbol)

    fun generateTypeInfo(irClass: IrClassSymbol, typeInfo: ConstantDataElement)
+    fun generateInterfaceTable(irClass: IrClassSymbol, table: ConstantDataElement)
+
+    fun registerInterfaceImplementationMethod(
+        interfaceImplementation: InterfaceImplementation,
+        table: Map<IrFunctionSymbol, WasmSymbol<WasmFunction>>,
+    )
+
+    fun referenceInterfaceImplementationId(interfaceImplementation: InterfaceImplementation): WasmSymbol<Int>
+    fun referenceInterfaceTableAddress(irClass: IrClassSymbol): WasmSymbol<Int>
 }
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.backend.wasm.ir2wasm

 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
-import org.jetbrains.kotlin.wasm.ir.*
 import org.jetbrains.kotlin.backend.wasm.lower.WasmSignature
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.declarations.IrField
@@ -20,6 +19,7 @@ import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.types.isNothing
 import org.jetbrains.kotlin.ir.util.isFunction
 import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.wasm.ir.*


 class WasmModuleCodegenContextImpl(
@@ -68,6 +68,10 @@ class WasmModuleCodegenContextImpl(
        wasmFragment.typeInfo.define(irClass, typeInfo)
    }

+    override fun generateInterfaceTable(irClass: IrClassSymbol, table: ConstantDataElement) {
+        wasmFragment.definedClassITableData.define(irClass, table)
+    }
+
    override fun setStartFunction(wasmFunction: WasmFunction) {
        wasmFragment.startFunction = wasmFunction
    }
@@ -108,6 +112,23 @@ class WasmModuleCodegenContextImpl(
        wasmFragment.functionTypes.define(irFunction, wasmFunctionType)
    }

+    override fun defineInterfaceMethodTable(irFunction: IrFunctionSymbol, wasmTable: WasmTable) {
+        wasmFragment.interfaceMethodTables.define(irFunction, wasmTable)
+    }
+
+    override fun referenceInterfaceImplementationId(
+        interfaceImplementation: InterfaceImplementation
+    ): WasmSymbol<Int> =
+        wasmFragment.referencedInterfaceImplementationId.reference(interfaceImplementation)
+
+
+    override fun registerInterfaceImplementationMethod(
+        interfaceImplementation: InterfaceImplementation,
+        table: Map<IrFunctionSymbol, WasmSymbol<WasmFunction>>
+    ) {
+        wasmFragment.interfaceImplementationsMethods[interfaceImplementation] = table
+    }
+
    private val classMetadataCache = mutableMapOf<IrClassSymbol, ClassMetadata>()
    override fun getClassMetadata(irClass: IrClassSymbol): ClassMetadata =
        classMetadataCache.getOrPut(irClass) {
@@ -143,6 +164,12 @@ class WasmModuleCodegenContextImpl(
    override fun referenceClassId(irClass: IrClassSymbol): WasmSymbol<Int> =
        wasmFragment.classIds.reference(irClass)

+    override fun referenceInterfaceTableAddress(irClass: IrClassSymbol): WasmSymbol<Int> {
+        if (irClass.owner.modality == Modality.ABSTRACT) return WasmSymbol(-1)
+        return wasmFragment.referencedClassITableAddresses.reference(irClass)
+    }
+
+
    override fun referenceInterfaceId(irInterface: IrClassSymbol): WasmSymbol<Int> {
        // HACK to substitute kotlin.Function5 with kotlin.wasm.internal.Function5
        val defaultType = irInterface.defaultType
@@ -164,6 +191,10 @@ class WasmModuleCodegenContextImpl(
        return wasmFragment.signatureId.reference(signature)
    }

+    override fun referenceInterfaceTable(irFunction: IrFunctionSymbol): WasmSymbol<WasmTable> {
+        return wasmFragment.interfaceMethodTables.reference(irFunction)
+    }
+
    override fun getStructFieldRef(field: IrField): WasmSymbol<Int> {
        val klass = field.parentAsClass
        val metadata = getClassMetadata(klass.symbol)