[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()
 

compiler/testData/codegen/boxWasmJsInterop/wasmImport.kt

@@ -0,0 +1,66 @@
+// TARGET_BACKEND: WASM
+
+// FILE: 1.mjs
+
+export function add(x, y) { return x + y; }
+
+export function giveMeFive(x) {
+    if (x !== 5)
+        throw "I expected 5";
+}
+
+// FILE: 2.mjs
+
+function sub(x, y) { return x - y; };
+
+export { sub };
+export { sub as "(˹˻𔗎𝅳𝅵𝅷𝅹⁽₍❨❪⟮﴾︵﹙（｟󠀨❲❴⟦⟨⟪⟬⦇⦉⦕⸢⸤︗︷︹︻︽︿﹁﹃﹇﹛﹝［｛｢󠁛󠁻«‘“‹❮" }
+export { sub as "~!@#\$%^&*()_+\`-={}|[]\\\\:\\\";'<>?,./" }
+export { sub as "" }
+export { sub as "\n  \r  \t" }
+export default sub;
+
+// FILE: wasmImport.kt
+import kotlin.wasm.WasmImport
+
+@WasmImport("./1.mjs", "add")
+external fun addImportRenamed(x: Int, y: Int): Int
+
+@WasmImport("./1.mjs")
+external fun giveMeFive(x: Int): Unit  // Test unit return type
+
+@WasmImport("./1.mjs")
+external fun add(x: Int, y: Int): Int
+
+@WasmImport("./2.mjs")
+external fun sub(x: Float, y: Float): Float
+
+@WasmImport("./2.mjs", "(˹˻𔗎𝅳𝅵𝅷𝅹⁽₍❨❪⟮﴾︵﹙（｟󠀨❲❴⟦⟨⟪⟬⦇⦉⦕⸢⸤︗︷︹︻︽︿﹁﹃﹇﹛﹝［｛｢󠁛󠁻«‘“‹❮")
+external fun sub2(x: Float, y: Float): Float
+
+@WasmImport("./2.mjs", "~!@#\$%^&*()_+`-={}|[]\\\\:\\\";'<>?,./")
+external fun sub3(x: Float, y: Float): Float
+
+@WasmImport("./2.mjs", "")
+external fun sub4(x: Float, y: Float): Float
+
+@WasmImport("./2.mjs", "\n  \r  \t")
+external fun sub5(x: Float, y: Float): Float
+
+@WasmImport("./2.mjs", "default")
+external fun sub6(x: Float, y: Float): Float
+
+fun box(): String {
+    if (addImportRenamed(5, 6) != 11) return "Fail1"
+    if (add(5, 6) != 11) return "Fail1"
+    giveMeFive(5)
+
+    if (sub(5f, 6f) != -1f) return "Fail2"
+    if (sub2(5f, 6f) != -1f) return "Fail3"
+    if (sub3(5f, 6f) != -1f) return "Fail4"
+    if (sub4(5f, 6f) != -1f) return "Fail5"
+    if (sub5(5f, 6f) != -1f) return "Fail6"
+    if (sub6(5f, 6f) != -1f) return "Fail7"
+
+    return "OK"
+}

js/js.tests/test/org/jetbrains/kotlin/js/testOld/BasicWasmBoxTest.kt

@@ -67,6 +67,7 @@ abstract class BasicWasmBoxTest(
             val kotlinFiles = mutableListOf<String>()
             val jsFilesBefore = mutableListOf<String>()
             val jsFilesAfter = mutableListOf<String>()
+            val mjsFiles = mutableListOf<String>()
 
             inputFiles.forEach {
                 val name = it.fileName
@@ -79,6 +80,9 @@ abstract class BasicWasmBoxTest(
 
                     name.endsWith(".js") ->
                         jsFilesBefore += name
+
+                    name.endsWith(".mjs") ->
+                        mjsFiles += name
                 }
             }
 
@@ -189,6 +193,9 @@ abstract class BasicWasmBoxTest(
                     println(" ------ $name Test file://$path/test.mjs")
                     val projectName = "kotlin"
                     println(" ------ $name HTML http://0.0.0.0:63342/$projectName/${dir.path}/index.html")
+                    for (mjsPath: String in mjsFiles) {
+                        println(" ------ $name External ESM file://$path/${File(mjsPath).name}")
+                    }
 
                     File(dir, "index.html").writeText(
                         """
@@ -216,6 +223,12 @@ abstract class BasicWasmBoxTest(
 
                 writeCompilationResult(res, dir, "index", sourceMapFileName = null)
                 File(dir, "test.mjs").writeText(testJs)
+
+                for (mjsPath: String in mjsFiles) {
+                    val mjsFile = File(mjsPath)
+                    File(dir, mjsFile.name).writeText(mjsFile.readText())
+                }
+
                 ExternalTool(System.getProperty("javascript.engine.path.V8"))
                     .run(
                         "--experimental-wasm-gc",

js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenWasmJsInteropWasmTestGenerated.java

@@ -84,4 +84,9 @@ public class IrCodegenWasmJsInteropWasmTestGenerated extends AbstractIrCodegenWa
     public void testTypes() throws Exception {
         runTest("compiler/testData/codegen/boxWasmJsInterop/types.kt");
     }
+
+    @TestMetadata("wasmImport.kt")
+    public void testWasmImport() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/wasmImport.kt");
+    }
 }

libraries/stdlib/wasm/src/kotlin/wasm/Annotations.kt

@@ -0,0 +1,21 @@
+/*
+ * 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 kotlin.wasm
+
+/**
+ * Imports a function from the given [module] with the given optional [name].
+ * The declaration name will be used if the [name] argument is not provided.
+ *
+ * Can only be used on top-level external functions.
+ *
+ * In JavaScript environment,
+ * function will be imported from ES module without type adapters.
+ */
+@Target(AnnotationTarget.FUNCTION)
+public annotation class WasmImport(
+    val module: String,
+    val name: String = ""
+)