[Wasm] Support restricted version of js("code") (KT-56955)

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

@@ -187,6 +187,12 @@ private val wasmStringSwitchOptimizerLowering = makeWasmModulePhase(
     description = "Replace when with constant string cases to binary search by string hashcodes"
 )
 
+private val jsCodeCallsLowering = makeWasmModulePhase(
+    ::JsCodeCallsLowering,
+    name = "JsCodeCallsLowering",
+    description = "Lower calls to js('code') into @JsFun",
+)
+
 private val complexExternalDeclarationsToTopLevelFunctionsLowering = makeWasmModulePhase(
     ::ComplexExternalDeclarationsToTopLevelFunctionsLowering,
     name = "ComplexExternalDeclarationsToTopLevelFunctionsLowering",
@@ -592,6 +598,7 @@ val wasmPhases = SameTypeNamedCompilerPhase(
     name = "IrModuleLowering",
     description = "IR module lowering",
     lower = validateIrBeforeLowering then
+            jsCodeCallsLowering then
             generateTests then
             excludeDeclarationsFromCodegenPhase then
             expectDeclarationsRemovingPhase then

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

@@ -39,6 +39,8 @@ class WasmSymbols(
         context.module.getPackage(FqName("kotlin.enums"))
     private val wasmInternalPackage: PackageViewDescriptor =
         context.module.getPackage(FqName("kotlin.wasm.internal"))
+    private val kotlinJsPackage: PackageViewDescriptor =
+        context.module.getPackage(FqName("kotlin.js"))
     private val collectionsPackage: PackageViewDescriptor =
         context.module.getPackage(StandardNames.COLLECTIONS_PACKAGE_FQ_NAME)
     private val builtInsPackage: PackageViewDescriptor =
@@ -329,6 +331,8 @@ class WasmSymbols(
     private val jsFunClass = getIrClass(FqName("kotlin.JsFun"))
     val jsFunConstructor by lazy { jsFunClass.constructors.single() }
 
+    val jsCode = getFunction("js", kotlinJsPackage)
+
     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/ir2wasm/DeclarationGenerator.kt

@@ -76,7 +76,7 @@ class DeclarationGenerator(
         }
 
         val wasmImportModule = declaration.getWasmImportDescriptor()
-        val jsCode = if (declaration.isExternal) declaration.getJsFunAnnotation() else null
+        val jsCode = declaration.getJsFunAnnotation()
         val importedName = when {
             wasmImportModule != null -> {
                 check(declaration.isExternal) { "Non-external fun with @WasmImport ${declaration.fqNameWhenAvailable}"}

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

@@ -354,15 +354,7 @@ class ComplexExternalDeclarationsToTopLevelFunctionsLowering(val context: WasmBa
         resultType: IrType,
         jsCode: String,
     ): IrSimpleFunction {
-        val res = context.irFactory.buildFun {
-            name = Name.identifier(originalName.asStringStripSpecialMarkers() + suffix)
-            returnType = resultType
-            isExternal = true
-        }
-        val builder = context.createIrBuilder(res.symbol)
-        res.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
-            it.putValueArgument(0, builder.irString(jsCode))
-        }
+        val res = createExternalJsFunction(context, originalName, suffix, resultType, jsCode)
         res.parent = currentFile
         addedDeclarations += res
         return res
@@ -388,6 +380,25 @@ class ComplexExternalDeclarationsToTopLevelFunctionsLowering(val context: WasmBa
     }
 }
 
+fun createExternalJsFunction(
+    context: WasmBackendContext,
+    originalName: Name,
+    suffix: String,
+    resultType: IrType,
+    jsCode: String,
+): IrSimpleFunction {
+    val res = context.irFactory.buildFun {
+        name = Name.identifier(originalName.asStringStripSpecialMarkers() + suffix)
+        returnType = resultType
+        isExternal = true
+    }
+    val builder = context.createIrBuilder(res.symbol)
+    res.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
+        it.putValueArgument(0, builder.irString(jsCode))
+    }
+    return res
+}
+
 /**
  * Redirect usages of complex declarations to top-level functions
  */

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

@@ -0,0 +1,88 @@
+/*
+ * 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.backend.wasm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.visitors.*
+
+/**
+ * Lower calls to `js(code)` into `@JsFun(code) external` functions.
+ */
+class JsCodeCallsLowering(val context: WasmBackendContext) : FileLoweringPass {
+    override fun lower(irFile: IrFile) {
+        irFile.transformDeclarationsFlat { declaration ->
+            when (declaration) {
+                is IrSimpleFunction -> transformFunction(declaration)
+                is IrProperty -> transformProperty(declaration)
+                else -> null
+            }
+        }
+    }
+
+    private fun transformFunction(function: IrSimpleFunction): List<IrDeclaration>? {
+        val body = function.body ?: return null
+        check(body is IrBlockBody)  // Should be lowered to block body
+        val statement = body.statements.singleOrNull() ?: return null
+
+        val isSingleExpressionJsCode: Boolean
+        val jsCode: String
+        when (statement) {
+            is IrReturn -> {
+                jsCode = statement.value.getJsCode() ?: return null
+                isSingleExpressionJsCode = true
+            }
+            is IrCall -> {
+                jsCode = statement.getJsCode() ?: return null
+                isSingleExpressionJsCode = false
+            }
+            else -> return null
+        }
+        val valueParameters = function.valueParameters
+        val jsFunCode = buildString {
+            append('(')
+            append(valueParameters.joinToString { it.name.identifier })
+            append(") => ")
+            if (!isSingleExpressionJsCode) append("{ ")
+            append(jsCode)
+            if (!isSingleExpressionJsCode) append(" }")
+        }
+        val builder = context.createIrBuilder(function.symbol)
+        function.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
+            it.putValueArgument(0, builder.irString(jsFunCode))
+        }
+        function.body = null
+        return null
+    }
+
+    private fun transformProperty(property: IrProperty): List<IrDeclaration>? {
+        val field = property.backingField ?: return null
+        val initializer = field.initializer ?: return null
+        val jsCode = initializer.expression.getJsCode() ?: return null
+        val externalFun = createExternalJsFunction(
+            context,
+            property.name,
+            "_js_code",
+            field.type,
+            jsCode = "() => ($jsCode)",
+        )
+        val builder = context.createIrBuilder(field.symbol)
+        initializer.expression = builder.irCall(externalFun)
+        return listOf(property, externalFun)
+    }
+
+    private fun IrExpression.getJsCode(): String? {
+        val call = this as? IrCall ?: return null
+        if (call.symbol != context.wasmSymbols.jsCode) return null
+        @Suppress("UNCHECKED_CAST")
+        return (call.getValueArgument(0) as IrConst<String>).value
+    }
+}

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.isBuiltInWasmRefType
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExported
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExternalType
+import org.jetbrains.kotlin.backend.wasm.utils.getJsFunAnnotation
 import org.jetbrains.kotlin.backend.wasm.utils.getWasmImportDescriptor
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
@@ -49,7 +50,7 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
         if (declaration.isFakeOverride) return null
         if (declaration !is IrSimpleFunction) return null
         val isExported = declaration.isExported()
-        val isExternal = declaration.isExternal
+        val isExternal = declaration.isExternal || declaration.getJsFunAnnotation() != null
         if (declaration.isPropertyAccessor) return null
         if (declaration.parent !is IrPackageFragment) return null
         if (!isExported && !isExternal) return null

compiler/testData/codegen/boxWasmJsInterop/jsCode.kt

@@ -0,0 +1,11 @@
+fun foo(x: Int, y: Int, z: String): Int = js("x + y + Number(z)")
+
+val x: String = js("typeof 10")
+
+fun box(): String {
+    val res = foo(10, 20, z = "30")
+    if (res != 60) return "Wrong foo: $res"
+    if (x != "number") return "Wrong x: $x"
+
+    return "OK"
+}

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/fir/FirJsCodegenWasmJsInteropTestGenerated.java

@@ -49,6 +49,12 @@ public class FirJsCodegenWasmJsInteropTestGenerated extends AbstractFirJsCodegen
         runTest("compiler/testData/codegen/boxWasmJsInterop/functionTypes.kt");
     }
 
+    @Test
+    @TestMetadata("jsCode.kt")
+    public void testJsCode() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/jsCode.kt");
+    }
+
     @Test
     @TestMetadata("jsExport.kt")
     public void testJsExport() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrCodegenWasmJsInteropJsTestGenerated.java

@@ -49,6 +49,12 @@ public class IrCodegenWasmJsInteropJsTestGenerated extends AbstractIrCodegenWasm
         runTest("compiler/testData/codegen/boxWasmJsInterop/functionTypes.kt");
     }
 
+    @Test
+    @TestMetadata("jsCode.kt")
+    public void testJsCode() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/jsCode.kt");
+    }
+
     @Test
     @TestMetadata("jsExport.kt")
     public void testJsExport() throws Exception {

libraries/stdlib/wasm/src/kotlin/js/core.kt

@@ -34,4 +34,37 @@ import kotlin.wasm.internal.ExcludedFromCodegen
  * ```
  */
 @ExcludedFromCodegen
-public external val definedExternally: Nothing
+public external val definedExternally: Nothing
+
+/**
+ * This function allows you to incorporate JavaScript [code] into Kotlin/Wasm codebase.
+ * It is used to implement top-level functions and initialize top-level properties.
+ *
+ * It is important to note, that calls to [js] function should be the only expression
+ * in a function body or a property initializer.
+ *
+ * [code] parameter should be a compile-time constant.
+ *
+ * When used in an expression context, [code] should contain a single JavaScript expression. For example:
+ *
+ * ``` kotlin
+ * val version: String = js("process.version")
+ * fun newEmptyJsArray(): JsValue = js("[]")
+ * ```
+ *
+ * When used in a function body, [code] is expected to be a list of JavaScript statements. For example:
+ *
+ * ``` kotlin
+ * fun log(message1: String, message2: String) {
+ *     js("""
+ *     console.log(message1);
+ *     console.log(message2);
+ *     """)
+ * }
+ * ```
+ *
+ * You can use parameters of calling function in JavaScript [code].
+ * However, other Kotlin declarations are not visible inside the [code] block.
+ */
+@ExcludedFromCodegen
+public external fun js(code: String): Nothing

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

@@ -65,6 +65,11 @@ public class IrCodegenWasmJsInteropWasmTestGenerated extends AbstractIrCodegenWa
         runTest("compiler/testData/codegen/boxWasmJsInterop/imperativeWrapperUninitialised.kt");
     }
 
+    @TestMetadata("jsCode.kt")
+    public void testJsCode() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/jsCode.kt");
+    }
+
     @TestMetadata("jsExport.kt")
     public void testJsExport() throws Exception {
         runTest("compiler/testData/codegen/boxWasmJsInterop/jsExport.kt");