From 35a5c9ea719fb2c921a789b563ce0c9d0238b90c Mon Sep 17 00:00:00 2001
From: Svyatoslav Kuzmich <svyatoslav.kuzmich@jetbrains.com>
Date: Wed, 8 Sep 2021 13:05:07 +0300
Subject: [PATCH] [Wasm] Support exporting functions via @kotlin.js.JsExport

Previously we exported main function only.
Extend this feature to arbitrary number of annotated functions.
Type restrictions and conversions are similar to imported functions.
---
 .../backend/wasm/ir2wasm/BodyGenerator.kt      | 18 ++++++++++++++++--
 .../wasm/ir2wasm/DeclarationGenerator.kt       | 18 +++++++++---------
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt
index e23e1e8e4cd..083f08f3a75 100644
--- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt
+++ b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt
@@ -172,7 +172,21 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
     }
 
     override fun visitGetValue(expression: IrGetValue) {
-        body.buildGetLocal(context.referenceLocal(expression.symbol))
+        val valueSymbol = expression.symbol
+        body.buildGetLocal(context.referenceLocal(valueSymbol))
+
+        val valueDeclaration = valueSymbol.owner
+        if (valueSymbol.owner is IrValueParameter) {
+            val parent = valueDeclaration.parent
+            if (parent is IrFunction && parent.isExported(backendContext)) {
+                val type = context.transformType(valueDeclaration.type)
+                if (type is WasmRefNullType) {
+                    // TODO: Add these casts as IR-2-IR lowering instead
+                    generateTypeRTT(valueDeclaration.type)
+                    body.buildRefCast()
+                }
+            }
+        }
     }
 
     override fun visitSetValue(expression: IrSetValue) {
@@ -296,7 +310,7 @@ class BodyGenerator(val context: WasmFunctionCodegenContext) : IrElementVisitorV
 
         // Return types of imported functions cannot have concrete struct/array references.
         // Non-primitive return types are represented as eqref which need to be casted back to expected type on call site.
-        if (function.getWasmImportAnnotation() != null || function.getJsFunAnnotation() != null || function.isExternal) {
+        if (function.getWasmImportAnnotation() != null || function.getJsFunAnnotation() != null || function.isExternal || function.isExported(backendContext)) {
             val resT = context.transformResultType(function.returnType)
             if (resT is WasmRefNullType) {
                 generateTypeRTT(function.returnType)
diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/DeclarationGenerator.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/DeclarationGenerator.kt
index b3b4834ce38..2c4d1864e75 100644
--- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/DeclarationGenerator.kt
+++ b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/DeclarationGenerator.kt
@@ -71,30 +71,30 @@ class DeclarationGenerator(val context: WasmModuleCodegenContext) : IrElementVis
         val watName = declaration.fqNameWhenAvailable.toString()
         val irParameters = declaration.getEffectiveValueParameters()
         // TODO: Exported types should be transformed in a separate lowering by creating shim functions for each export.
+        val isExported = declaration.isExported(context.backendContext)
         val resultType =
             when {
-                declaration.isExported(context.backendContext) -> context.transformExportedResultType(declaration.returnType)
+                isExported -> context.transformExportedResultType(declaration.returnType)
                 // Unit_getInstance returns true Unit reference instead of "void"
                 declaration == unitGetInstanceFunction -> context.transformType(declaration.returnType)
                 else -> context.transformResultType(declaration.returnType)
             }
+
+        val isExportedOrImported = isExported || importedName != null
         val wasmFunctionType =
             WasmFunctionType(
                 name = watName,
                 parameterTypes = irParameters.map {
                     val t = context.transformValueParameterType(it)
-                    if (importedName != null && t is WasmRefNullType) {
-                        WasmEqRef
+                    if (isExportedOrImported && t is WasmRefNullType) {
+                        WasmRefNullType(WasmHeapType.Simple.Data)
                     } else {
                         t
                     }
                 },
                 resultTypes = listOfNotNull(
-                    run {
-                        if (importedName != null && resultType is WasmRefNullType)
-                            WasmEqRef
-                        else
-                            resultType
+                    resultType.let {
+                        if (isExportedOrImported && it is WasmRefNullType) WasmRefNullType(WasmHeapType.Simple.Data) else it
                     }
                 )
             )
@@ -377,4 +377,4 @@ fun IrFunction.getEffectiveValueParameters(): List<IrValueParameter> {
 }
 
 fun IrFunction.isExported(context: WasmBackendContext): Boolean =
-    visibility == DescriptorVisibilities.PUBLIC && fqNameWhenAvailable in context.additionalExportedDeclarations
\ No newline at end of file
+    fqNameWhenAvailable in context.additionalExportedDeclarations || isJsExport()
\ No newline at end of file