[Wasm] Support main function

2021-10-13 12:08:54 +03:00
parent af18b10da9
commit 0f66f85cf9
22 changed files with 259 additions and 122 deletions
@@ -323,9 +323,8 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                val runner = """
                    export default WebAssembly.instantiateStreaming(fetch('${outputWasmFile.name}'), { runtime, js_code }).then((it) => {
                        wasmInstance = it.instance;
-                        
+                        wasmInstance.exports.__init?.();
                        wasmInstance.exports.startUnitTests?.();
-                        wasmInstance.exports.main?.();
                        
                        return it.instance.exports;
                    });
@@ -14,9 +14,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.impl.EmptyPackageFragmentDescriptor
 import org.jetbrains.kotlin.ir.*
-import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
-import org.jetbrains.kotlin.ir.backend.js.JsCommonCoroutineSymbols
-import org.jetbrains.kotlin.ir.backend.js.JsMapping
+import org.jetbrains.kotlin.ir.backend.js.*
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.backend.js.lower.JsInnerClassesSupport
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
@@ -35,8 +33,8 @@ import org.jetbrains.kotlin.name.Name
 class WasmBackendContext(
    val module: ModuleDescriptor,
    override val irBuiltIns: IrBuiltIns,
-    @Suppress("UNUSED_PARAMETER") symbolTable: SymbolTable,
-    @Suppress("UNUSED_PARAMETER") irModuleFragment: IrModuleFragment,
+    symbolTable: SymbolTable,
+    irModuleFragment: IrModuleFragment,
    val additionalExportedDeclarations: Set<FqName>,
    override val configuration: CompilerConfiguration,
 ) : JsCommonBackendContext {
@@ -90,15 +88,17 @@ class WasmBackendContext(
        }
    }

-
-    val startFunction = irFactory.buildFun {
-        name = Name.identifier("startFunction")
+    fun createInitFunction(identifier: String): IrSimpleFunction = irFactory.buildFun {
+        name = Name.identifier(identifier)
        returnType = irBuiltIns.unitType
    }.apply {
        body = irFactory.createBlockBody(UNDEFINED_OFFSET, UNDEFINED_OFFSET)
        internalPackageFragment.addChild(this)
    }

+    val fieldInitFunction = createInitFunction("fieldInit")
+    val mainCallsWrapperFunction = createInitFunction("mainCallsWrapper")
+
    override val sharedVariablesManager =
        WasmSharedVariablesManager(this, irBuiltIns, internalPackageFragment)

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.AddContinuationToNonL
 import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.JsSuspendFunctionsLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.RemoveInlineDeclarationsWithReifiedTypeParametersLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.WrapInlineDeclarationsWithReifiedTypeParametersLowering
+import org.jetbrains.kotlin.ir.backend.wasm.lower.generateMainFunctionCalls
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents

@@ -269,6 +270,12 @@ private val addContinuationToFunctionCallsLoweringPhase = makeWasmModulePhase(
    )
 )

+private val addMainFunctionCallsLowering = makeCustomWasmModulePhase(
+    ::generateMainFunctionCalls,
+    name = "GenerateMainFunctionCalls",
+    description = "Generate main function calls into start function",
+)
+
 private val defaultArgumentStubGeneratorPhase = makeWasmModulePhase(
    ::DefaultArgumentStubGenerator,
    name = "DefaultArgumentStubGenerator",
@@ -521,6 +528,7 @@ val wasmPhases = NamedCompilerPhase(

            addContinuationToNonLocalSuspendFunctionsLoweringPhase then
            addContinuationToFunctionCallsLoweringPhase then
+            addMainFunctionCallsLowering then

            tryCatchCanonicalization then
            returnableBlockLoweringPhase then
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.ir.builders.declarations.addFunction
 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.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.classifierOrFail
@@ -72,6 +73,10 @@ class WasmSymbols(
    override val returnIfSuspended =
        getInternalFunction("returnIfSuspended")

+    val coroutineEmptyContinuation: IrPropertySymbol = symbolTable.referenceProperty(
+        getProperty(FqName.fromSegments(listOf("kotlin", "wasm", "internal", "EmptyContinuation")))
+    )
+
    override val functionAdapter: IrClassSymbol
        get() = TODO()

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.backend.common.ir.isOverridableOrOverrides
 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.backend.wasm.lower.wasmSignature
 import org.jetbrains.kotlin.backend.wasm.utils.*
-import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
@@ -18,7 +17,6 @@ import org.jetbrains.kotlin.ir.backend.js.utils.findUnitGetInstanceFunction
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.declarations.*
 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
@@ -158,8 +156,13 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis

        context.defineFunction(declaration.symbol, function)

-        if (declaration == backendContext.startFunction)
-            context.setStartFunction(function)
+        val initPriority = when (declaration) {
+            backendContext.fieldInitFunction -> "0"
+            backendContext.mainCallsWrapperFunction -> "1"
+            else -> null
+        }
+        if (initPriority != null)
+            context.registerInitFunction(function, initPriority)

        if (declaration.isExported(backendContext)) {
            context.addExport(
@@ -82,7 +82,8 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {

    val jsFuns = mutableListOf<JsCodeSnippet>()

-    var startFunction: WasmFunction? = null
+    class FunWithPriority(val function: WasmFunction, val priority: String)
+    val initFunctions = mutableListOf<FunWithPriority>()

    val scratchMemAddr = WasmSymbol<Int>()
    val scratchMemSizeInBytes = 65_536
@@ -253,6 +254,15 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
            )
        }

+        val masterInitFunctionType = WasmFunctionType("__init_t", emptyList(), emptyList())
+        val masterInitFunction = WasmFunction.Defined("__init", masterInitFunctionType)
+        with(WasmIrExpressionBuilder(masterInitFunction.instructions)) {
+            initFunctions.sortedBy { it.priority }.forEach {
+                buildCall(WasmSymbol(it.function))
+            }
+        }
+        exports += WasmExport.Function("__init", masterInitFunction)
+
        interfaceMethodTables.defined.forEach { (function, table) ->
            val size = interfaceTableElementsLists[function]!!.size.toUInt()
            table.limits = WasmLimits(size, size)
@@ -271,16 +281,16 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
        val sortedRttGlobals = runtimeTypes.elements.sortedBy { (it.type as WasmRtt).depth }

        val module = WasmModule(
-            functionTypes = functionTypes.elements + tagFuncType,
+            functionTypes = functionTypes.elements + tagFuncType + masterInitFunctionType,
            gcTypes = gcTypes.elements,
            importsInOrder = importedFunctions,
            importedFunctions = importedFunctions,
-            definedFunctions = functions.elements.filterIsInstance<WasmFunction.Defined>(),
+            definedFunctions = functions.elements.filterIsInstance<WasmFunction.Defined>() + masterInitFunction,
            tables = listOf(table) + interfaceMethodTables.elements,
            memories = listOf(memory),
            globals = globals.elements + sortedRttGlobals,
            exports = exports,
-            startFunction = startFunction!!,
+            startFunction = null,  // Module is initialized via export call
            elements = listOf(elements) + interfaceTableElements,
            data = data,
            tags = listOf(tag)
@@ -23,7 +23,7 @@ interface WasmModuleCodegenContext : WasmBaseCodegenContext {
    fun defineInterfaceMethodTable(irFunction: IrFunctionSymbol, wasmTable: WasmTable)
    fun addJsFun(importName: String, jsCode: String)

-    fun setStartFunction(wasmFunction: WasmFunction)
+    fun registerInitFunction(wasmFunction: WasmFunction, priority: String)
    fun addExport(wasmExport: WasmExport<*>)

    fun registerVirtualFunction(irFunction: IrSimpleFunctionSymbol)
@@ -81,8 +81,8 @@ class WasmModuleCodegenContextImpl(
        wasmFragment.definedClassITableData.define(irClass, table)
    }

-    override fun setStartFunction(wasmFunction: WasmFunction) {
-        wasmFragment.startFunction = wasmFunction
+    override fun registerInitFunction(wasmFunction: WasmFunction, priority: String) {
+        wasmFragment.initFunctions += WasmCompiledModuleFragment.FunWithPriority(wasmFunction, priority)
    }

    override fun addExport(wasmExport: WasmExport<*>) {
@@ -29,8 +29,8 @@ import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 */
 class FieldInitializersLowering(val context: WasmBackendContext) : FileLoweringPass {
    override fun lower(irFile: IrFile) {
-        val builder = context.createIrBuilder(context.startFunction.symbol)
-        val startFunctionBody = context.startFunction.body as IrBlockBody
+        val builder = context.createIrBuilder(context.fieldInitFunction.symbol)
+        val startFunctionBody = context.fieldInitFunction.body as IrBlockBody

        irFile.acceptChildrenVoid(object : IrElementVisitorVoid {
            override fun visitElement(element: IrElement) {
@@ -0,0 +1,49 @@
+/*
+ * 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.ir.backend.wasm.lower
+
+import org.jetbrains.kotlin.backend.common.ir.createArrayOfExpression
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.backend.js.utils.JsMainFunctionDetector
+import org.jetbrains.kotlin.ir.backend.js.utils.isLoweredSuspendFunction
+import org.jetbrains.kotlin.ir.backend.js.utils.isStringArrayParameter
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+
+/**
+ * Find single most appropriate main function and call with empty arguments from
+ * [WasmBackendContext.mainCallsWrapperFunction].
+ */
+fun generateMainFunctionCalls(
+    context: WasmBackendContext,
+    module: IrModuleFragment
+) {
+    val mainFunction = JsMainFunctionDetector(context).getMainFunctionOrNull(module) ?: return
+    val generateArgv = mainFunction.valueParameters.firstOrNull()?.isStringArrayParameter() ?: false
+    val generateContinuation = mainFunction.isLoweredSuspendFunction(context)
+    with(context.createIrBuilder(context.mainCallsWrapperFunction.symbol)) {
+        val argv = if (generateArgv) {
+            context.createArrayOfExpression(UNDEFINED_OFFSET, UNDEFINED_OFFSET, context.irBuiltIns.stringType, emptyList())
+        } else {
+            null
+        }
+
+        val continuation =
+            if (generateContinuation) {
+                irCall(context.wasmSymbols.coroutineEmptyContinuation.owner.getter!!)
+            } else {
+                null
+            }
+
+        (context.mainCallsWrapperFunction.body as IrBlockBody).statements += irCall(mainFunction).also { call ->
+            listOfNotNull(argv, continuation).forEachIndexed { index: Int, arg: IrExpression -> call.putValueArgument(index, arg) }
+        }
+    }
+}