[Wasm/WASI] Implementation of a callback on root exported function exit

Fixed #KT-64486

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt

@@ -567,6 +567,12 @@ private val propertyAccessorInlinerLoweringPhase = makeIrModulePhase(
     description = "[Optimization] Inline property accessors"
 )
 
+private val invokeOnExportedFunctionExitLowering = makeIrModulePhase(
+    ::InvokeOnExportedFunctionExitLowering,
+    name = "InvokeOnExportedFunctionExitLowering",
+    description = "This pass is needed to call exported function exit callback for WASI",
+)
+
 private val expressionBodyTransformer = makeIrModulePhase(
     ::ExpressionBodyTransformer,
     name = "ExpressionBodyTransformer",
@@ -673,6 +679,8 @@ val loweringList = listOf(
     addContinuationToFunctionCallsLoweringPhase,
     addMainFunctionCallsLowering,
 
+    invokeOnExportedFunctionExitLowering,
+
     unhandledExceptionLowering,
 
     tryCatchCanonicalization,

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt

@@ -384,6 +384,15 @@ class WasmSymbols(
 
     val jsRelatedSymbols get() = jsRelatedSymbolsIfNonWasi ?: error("Cannot access to js related std in wasi mode")
 
+
+    private val invokeOnExportedFunctionExitIfWasi =
+        when (context.configuration.get(JSConfigurationKeys.WASM_TARGET, WasmTarget.JS) == WasmTarget.WASI) {
+            true -> getInternalFunction("invokeOnExportedFunctionExit")
+            else -> null
+        }
+
+    val invokeOnExportedFunctionExit get() = invokeOnExportedFunctionExitIfWasi ?: error("Cannot access to wasi related std in js mode")
+
     private fun findClass(memberScope: MemberScope, name: Name): ClassDescriptor =
         memberScope.getContributedClassifier(name, NoLookupLocation.FROM_BACKEND) as ClassDescriptor
 

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/InvokeOnExportedFunctionExitLowering.kt

@@ -0,0 +1,127 @@
+/*
+ * Copyright 2010-2022 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.backend.wasm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.lower.irNot
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
+import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExported
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.buildVariable
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
+import org.jetbrains.kotlin.ir.util.toIrConst
+import org.jetbrains.kotlin.js.config.JSConfigurationKeys
+import org.jetbrains.kotlin.js.config.WasmTarget
+import org.jetbrains.kotlin.name.Name
+
+// This pass needed to call coroutines event loop run after exported functions calls
+// @WasmExport
+// fun someExportedMethod() {
+//     println("hello world")
+// }
+//
+// converts into
+//
+// @WasmExport
+// fun someExportedMethod() {
+//     val currentIsNotFirstWasmExportCall = isNotFirstWasmExportCall
+//     try {
+//         isNotFirstWasmExportCall = true
+//         println("hello world")
+//     } finally {
+//         isNotFirstWasmExportCall = currentIsNotFirstWasmExportCall
+//         if (!currentIsNotFirstWasmExportCall) invokeOnExportedFunctionExit()
+//     }
+// }
+
+internal class InvokeOnExportedFunctionExitLowering(val context: WasmBackendContext) : FileLoweringPass {
+    private val invokeOnExportedFunctionExit get() = context.wasmSymbols.invokeOnExportedFunctionExit
+    private val irBooleanType = context.wasmSymbols.irBuiltIns.booleanType
+    private val isNotFirstWasmExportCallGetter = context.wasmSymbols.isNotFirstWasmExportCall.owner.getter!!.symbol
+    private val isNotFirstWasmExportCallSetter = context.wasmSymbols.isNotFirstWasmExportCall.owner.setter!!.symbol
+
+    private fun processExportFunction(irFunction: IrFunction) {
+        val body = irFunction.body ?: return
+        if (body is IrBlockBody && body.statements.isEmpty()) return
+
+        val bodyType = when (body) {
+            is IrExpressionBody -> body.expression.type
+            is IrBlockBody -> context.irBuiltIns.unitType
+            else -> TODO(this::class.qualifiedName!!)
+        }
+
+        with(context.createIrBuilder(irFunction.symbol)) {
+            val currentIsNotFirstWasmExportCall = buildVariable(
+                parent = irFunction,
+                startOffset = UNDEFINED_OFFSET, endOffset = UNDEFINED_OFFSET,
+                origin = IrDeclarationOrigin.IR_TEMPORARY_VARIABLE,
+                name = Name.identifier("currentIsNotFirstWasmExportCall"),
+                type = irBooleanType
+            )
+
+            currentIsNotFirstWasmExportCall.initializer =
+                irGet(irBooleanType, null, isNotFirstWasmExportCallGetter)
+
+            val tryBody = irComposite {
+                +irSet(irBooleanType, null, isNotFirstWasmExportCallSetter, true.toIrConst(irBooleanType))
+                when (body) {
+                    is IrBlockBody -> body.statements.forEach { +it }
+                    is IrExpressionBody -> +body.expression
+                    else -> TODO(this::class.qualifiedName!!)
+                }
+            }
+
+            val finally = irComposite(resultType = context.irBuiltIns.unitType) {
+                +irSet(
+                    type = irBooleanType,
+                    receiver = null,
+                    setterSymbol = isNotFirstWasmExportCallSetter,
+                    value = irGet(currentIsNotFirstWasmExportCall, irBooleanType)
+                )
+                +irIfThen(
+                    type = context.irBuiltIns.unitType,
+                    condition = irNot(irGet(currentIsNotFirstWasmExportCall, irBooleanType)),
+                    thenPart = irCall(invokeOnExportedFunctionExit),
+                )
+            }
+
+            val tryWrap = irTry(
+                type = bodyType,
+                tryResult = tryBody,
+                catches = emptyList(),
+                finallyExpression = finally
+            )
+
+            when (body) {
+                is IrExpressionBody -> body.expression = irComposite {
+                    +currentIsNotFirstWasmExportCall
+                    +tryWrap
+                }
+                is IrBlockBody -> with(body.statements) {
+                    clear()
+                    add(currentIsNotFirstWasmExportCall)
+                    add(tryWrap)
+                }
+                else -> TODO(this::class.qualifiedName!!)
+            }
+        }
+    }
+
+    override fun lower(irFile: IrFile) {
+        if (context.isWasmJsTarget) return
+        for (declaration in irFile.declarations) {
+            if (declaration is IrFunction && (declaration.isExported() || context.mainCallsWrapperFunction == declaration)) {
+                processExportFunction(declaration)
+            }
+        }
+    }
+}

compiler/testData/codegen/boxWasmWasi/onExportedFunctionExit.kt

@@ -0,0 +1,33 @@
+// MODULE: main
+// FILE: main.kt
+
+var onExportedFunctionExitCounter = 0
+
+@WasmExport()
+@Suppress("OPT_IN_USAGE_ERROR") // Opt-in kotlin.wasm.internal.InternalWasmApi is private to prevent user use of this internal API
+fun initializeOnExportedFunctionExit() {
+    kotlin.wasm.internal.onExportedFunctionExit = { onExportedFunctionExitCounter++ }
+}
+
+@WasmExport()
+fun getCounter() = onExportedFunctionExitCounter
+
+@WasmExport()
+fun fooRecursive(n: Int) {
+    if (n > 0)
+        fooRecursive(n - 1)
+}
+
+fun box() = "OK"
+
+// FILE: entry.mjs
+
+import main from "./index.mjs"
+
+if (main.getCounter() !== 0) throw "Err0"
+main.initializeOnExportedFunctionExit()
+if (main.getCounter() !== 1) throw "Err1"  // +1 from getCounter (should initializeOnExportedFunctionExit also increse the counter?)
+main.fooRecursive(0)
+if (main.getCounter() !== 3) throw "Err2"  // +2 from fooRecursive(0) and getCounter
+main.fooRecursive(10)
+if (main.getCounter() !== 5) throw "Err3"  // +2 from fooRecursive(10) and getCounter

libraries/stdlib/wasm/wasi/src/kotlin/internal/internalCallback.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2023 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 kotlin.wasm.internal
+
+@RequiresOptIn
+private annotation class InternalWasmApi
+
+/*
+This is internal API which is not intended to use on user-side.
+ */
+@InternalWasmApi
+public var onExportedFunctionExit: (() -> Unit)? = null
+
+internal fun invokeOnExportedFunctionExit() {
+    @OptIn(InternalWasmApi::class)
+    onExportedFunctionExit?.invoke()
+}

libraries/tools/kotlin-stdlib-docs/build.gradle.kts

@@ -74,6 +74,7 @@ fun createStdLibVersionedDocTask(version: String, isLatest: Boolean) =
                 "kotlin.native.internal",
                 "kotlin.jvm.functions",
                 "kotlin.coroutines.jvm.internal",
+                "kotlin.wasm.internal",
         )
 
         val kotlinLanguageVersion = version

wasm/wasm.tests/test/org/jetbrains/kotlin/wasm/test/handlers/WasiBoxRunner.kt

@@ -32,6 +32,7 @@ class WasiBoxRunner(
         val outputDirBase = testServices.getWasmTestOutputDirectory()
 
         val originalFile = testServices.moduleStructure.originalTestDataFiles.first()
+        val collectedJsArtifacts = collectJsArtifacts(originalFile)
 
         val debugMode = DebugMode.fromSystemProperty("kotlin.wasm.debugMode")
         val startUnitTests = RUN_UNIT_TESTS in testServices.moduleStructure.allDirectives
@@ -67,6 +68,7 @@ class WasiBoxRunner(
             writeCompilationResult(res, dir, baseFileName)
 
             File(dir, "test.mjs").writeText(testWasi)
+            val (jsFilePaths) = collectedJsArtifacts.saveJsArtifacts(dir)
 
             if (debugMode >= DebugMode.DEBUG) {
                 val path = dir.absolutePath
@@ -84,8 +86,8 @@ class WasiBoxRunner(
                     debugMode = debugMode,
                     disableExceptions = false,
                     failsIn = failsIn,
-                    entryMjs = "test.mjs",
-                    jsFilePaths = emptyList(),
+                    entryMjs = collectedJsArtifacts.entryPath ?: "test.mjs",
+                    jsFilePaths = jsFilePaths,
                     workingDirectory = dir
                 )
             }

wasm/wasm.tests/tests-gen/org/jetbrains/kotlin/wasm/test/K1WasmWasiCodegenBoxTestGenerated.java

@@ -25,6 +25,12 @@ public class K1WasmWasiCodegenBoxTestGenerated extends AbstractK1WasmWasiCodegen
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/boxWasmWasi"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.WASM, true);
     }
 
+    @Test
+    @TestMetadata("onExportedFunctionExit.kt")
+    public void testOnExportedFunctionExit() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmWasi/onExportedFunctionExit.kt");
+    }
+
     @Test
     @TestMetadata("simpleWasi.kt")
     public void testSimpleWasi() throws Exception {