[Wasm] Add @WasmImport annotation

Imports top-level function from given module

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
 import org.jetbrains.kotlin.backend.common.serialization.linkerissues.checkNoUnboundSymbols
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmCompiledModuleFragment
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmModuleFragmentGenerator
+import org.jetbrains.kotlin.backend.wasm.ir2wasm.toJsStringLiteral
 import org.jetbrains.kotlin.backend.wasm.lower.markExportedDeclarations
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.ir.backend.js.MainModule
@@ -179,12 +180,28 @@ fun WasmCompiledModuleFragment.generateJs(): String {
         return ifNotCached;
     }    """.trimIndent()
 
-    val jsCodeBody = jsFuns.joinToString(",\n") { "\"" + it.importName + "\" : " + it.jsCode }
-    val jsCodeBodyIndented = jsCodeBody.prependIndent("    ")
-    val jsCode =
-        "\nconst js_code = {\n$jsCodeBodyIndented\n};\n"
+    val imports = jsModuleImports
+        .toList()
+        .sorted()
+        .joinToString("\n") {
+            val moduleSpecifier = it.toJsStringLiteral()
+            "  $moduleSpecifier: await import($moduleSpecifier),"
+        }
 
-    return runtime + jsCode
+    val jsCodeBody = jsFuns.joinToString(",\n") {
+        "${it.importName.toJsStringLiteral()} : ${it.jsCode}"
+    }
+    val jsCodeBodyIndented = jsCodeBody.prependIndent("    ")
+    val importObject = """
+const _import_object = { 
+${imports} 
+  js_code: { 
+${jsCodeBodyIndented}
+  }
+};
+"""
+
+    return runtime + importObject
 }
 
 fun generateJsWasmLoader(wasmFilePath: String, externalJs: String): String =
@@ -210,17 +227,17 @@ fun generateJsWasmLoader(wasmFilePath: String, externalJs: String): String =
       const dirpath = path.dirname(filepath);
       const wasmBuffer = fs.readFileSync(path.resolve(dirpath, '$wasmFilePath'));
       const wasmModule = new WebAssembly.Module(wasmBuffer);
-      wasmInstance = new WebAssembly.Instance(wasmModule, { js_code });
+      wasmInstance = new WebAssembly.Instance(wasmModule, _import_object);
     }
     
     if (isD8) {
       const wasmBuffer = read('$wasmFilePath', 'binary');
       const wasmModule = new WebAssembly.Module(wasmBuffer);
-      wasmInstance = new WebAssembly.Instance(wasmModule, { js_code });
+      wasmInstance = new WebAssembly.Instance(wasmModule, _import_object);
     }
     
     if (isBrowser) {
-      wasmInstance = (await WebAssembly.instantiateStreaming(fetch('$wasmFilePath'), { js_code })).instance;
+      wasmInstance = (await WebAssembly.instantiateStreaming(fetch('$wasmFilePath'), _import_object)).instance;
     }
     
     const wasmExports = wasmInstance.exports;

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/DeclarationGenerator.kt

@@ -67,11 +67,28 @@ class DeclarationGenerator(
             return
         }
 
-        val jsCode = declaration.getJsFunAnnotation() ?: if (declaration.isExternal) declaration.name.asString() else null
-        val importedName = jsCode?.let {
-            val jsCodeName = jsCodeName(declaration)
-            context.addJsFun(jsCodeName, it)
-            WasmImportPair("js_code", jsCodeName(declaration))
+        val wasmImportModule = declaration.getWasmImportDescriptor()
+
+        val jsCode = declaration.getJsFunAnnotation()
+            // TODO: Why are we importing declarations by with raw declaration.name.asString() jsCode?
+            ?: if (declaration.isExternal && wasmImportModule == null) declaration.name.asString() else null
+
+        val importedName = when {
+            wasmImportModule != null -> {
+                check(declaration.isExternal) { "Non-external fun with @WasmImport ${declaration.fqNameWhenAvailable}"}
+                context.addJsModuleImport(wasmImportModule.moduleName)
+                wasmImportModule
+            }
+            jsCode != null -> {
+                // TODO: check consistency (currently fails with generated __convertKotlinClosureToJsClosure)
+                // check(declaration.isExternal) { "Non-external fun with @JsFun ${declaration.fqNameWhenAvailable}"}
+                val jsCodeName = jsCodeName(declaration)
+                context.addJsFun(jsCodeName, jsCode)
+                WasmImportPair("js_code", jsCodeName(declaration))
+            }
+            else -> {
+                null
+            }
         }
 
         if (declaration.isFakeOverride)

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/JsHelpers.kt

@@ -5,17 +5,7 @@
 
 package org.jetbrains.kotlin.backend.wasm.ir2wasm
 
-import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.jsAssignment
-import org.jetbrains.kotlin.js.backend.ast.JsArrayLiteral
-import org.jetbrains.kotlin.js.backend.ast.JsBlock
-import org.jetbrains.kotlin.js.backend.ast.JsNameRef
-import org.jetbrains.kotlin.js.backend.ast.JsStringLiteral
+import org.jetbrains.kotlin.js.backend.JsToStringGenerationVisitor
 
-fun generateStringLiteralsSupport(literals: List<String>): String {
-    return JsBlock(
-        jsAssignment(
-            JsNameRef("stringLiterals", "runtime"),
-            JsArrayLiteral(literals.map { JsStringLiteral(it) })
-        ).makeStmt()
-    ).toString()
-}
+fun String.toJsStringLiteral(): CharSequence =
+    JsToStringGenerationVisitor.javaScriptString(this)

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmCompiledModuleFragment.kt

@@ -58,6 +58,7 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
     class JsCodeSnippet(val importName: String, val jsCode: String)
 
     val jsFuns = mutableListOf<JsCodeSnippet>()
+    val jsModuleImports = mutableSetOf<String>()
 
     class FunWithPriority(val function: WasmFunction, val priority: String)
 

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/WasmModuleCodegenContext.kt

@@ -198,5 +198,9 @@ class WasmModuleCodegenContext(
         wasmFragment.jsFuns +=
             WasmCompiledModuleFragment.JsCodeSnippet(importName = importName, jsCode = jsCode)
     }
+
+    fun addJsModuleImport(module: String) {
+        wasmFragment.jsModuleImports += module
+    }
 }
 

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.backend.wasm.utils.getJsFunAnnotation
+import org.jetbrains.kotlin.backend.wasm.utils.getWasmImportDescriptor
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
@@ -175,6 +176,11 @@ class ComplexExternalDeclarationsToTopLevelFunctionsLowering(val context: WasmBa
     }
 
     fun processExternalSimpleFunction(function: IrSimpleFunction) {
+        // Skip JS interop adapters form WasmImport.
+        // It needs to keep original signature to interop with other Wasm modules.
+        if (function.getWasmImportDescriptor() != null)
+            return
+
         val jsFun = function.getJsFunAnnotation()
         // Wrap external functions without @JsFun to lambdas `foo` -> `(a, b) => foo(a, b)`.
         // This way we wouldn't fail if we don't call them.

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/utils/Annotations.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.wasm.utils
 import org.jetbrains.kotlin.ir.backend.js.utils.getSingleConstStringArgument
 import org.jetbrains.kotlin.ir.declarations.IrAnnotationContainer
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.IrClassReference
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.types.makeNullable
@@ -20,6 +21,17 @@ import org.jetbrains.kotlin.wasm.ir.WasmImportPair
 fun IrAnnotationContainer.hasExcludedFromCodegenAnnotation(): Boolean =
     hasAnnotation(FqName("kotlin.wasm.internal.ExcludedFromCodegen"))
 
+fun IrFunction.getWasmImportDescriptor(): WasmImportPair? {
+    val annotation = getAnnotation(FqName("kotlin.wasm.WasmImport"))
+        ?: return null
+
+    @Suppress("UNCHECKED_CAST")
+    return WasmImportPair(
+        (annotation.getValueArgument(0) as IrConst<String>).value,
+        (annotation.getValueArgument(1) as? IrConst<String>)?.value ?: this.name.asString()
+    )
+}
+
 fun IrAnnotationContainer.getWasmOpAnnotation(): String? =
     getAnnotation(FqName("kotlin.wasm.internal.WasmOp"))?.getSingleConstStringArgument()