[JS IR] Don't parse/print JS code in js() call twice

#KT-51327 Fixed

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.ir.backend.js
 
+import org.jetbrains.kotlin.backend.common.compilationException
 import org.jetbrains.kotlin.backend.common.ir.Ir
 import org.jetbrains.kotlin.backend.common.ir.Symbols
 import org.jetbrains.kotlin.builtins.PrimitiveType
@@ -22,24 +23,23 @@ import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.backend.js.lower.JsInnerClassesSupport
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.JsPolyfills
-import org.jetbrains.kotlin.ir.backend.js.utils.JsInlineClassesUtils
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.parseJsCode
-import org.jetbrains.kotlin.ir.backend.js.utils.MinimizedNameGenerator
+import org.jetbrains.kotlin.ir.backend.js.utils.*
-import org.jetbrains.kotlin.ir.backend.js.utils.OperatorNames
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.*
-import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
-import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
-import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.DescriptorlessExternalPackageFragmentSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrDynamicTypeImpl
 import org.jetbrains.kotlin.ir.util.SymbolTable
+import org.jetbrains.kotlin.ir.util.getAnnotation
 import org.jetbrains.kotlin.ir.util.kotlinPackageFqn
 import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.js.backend.ast.JsExpressionStatement
+import org.jetbrains.kotlin.js.backend.ast.JsFunction
 import org.jetbrains.kotlin.js.config.ErrorTolerancePolicy
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.js.config.RuntimeDiagnostic
@@ -50,6 +50,7 @@ import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.isNullable
 import org.jetbrains.kotlin.util.collectionUtils.filterIsInstanceMapNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.cast
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 enum class IcCompatibleIr2Js(val isCompatible: Boolean, val incrementalCacheEnabled: Boolean) {
     DISABLED(false, false),
@@ -379,4 +380,26 @@ class JsIrBackendContext(
         /*TODO*/
         print(message)
     }
+
+    private val outlinedJsCodeFunctions = mutableMapOf<IrFunctionSymbol, JsFunction>()
+
+    fun addOutlinedJsCode(symbol: IrSimpleFunctionSymbol, outlinedJsCode: JsFunction) {
+        outlinedJsCodeFunctions[symbol] = outlinedJsCode
+    }
+
+    fun getJsCodeForFunction(symbol: IrFunctionSymbol): JsFunction? {
+        val jsFunction = outlinedJsCodeFunctions[symbol]
+        if (jsFunction != null) return jsFunction
+        val jsFunAnnotation = symbol.owner.getAnnotation(JsAnnotations.jsFunFqn) ?: return null
+        val jsCode = jsFunAnnotation.getSingleConstStringArgument()
+        // FIXME: Provide debug info
+        val statements = parseJsCode(jsCode) ?: compilationException("Could not parse JS code", jsFunAnnotation)
+        val parsedJsFunction = statements.singleOrNull()
+            ?.safeAs<JsExpressionStatement>()
+            ?.expression
+            ?.safeAs<JsFunction>()
+            ?: compilationException("Provided JS code is not a js function", jsFunAnnotation)
+        outlinedJsCodeFunctions[symbol] = parsedJsFunction
+        return parsedJsFunction
+    }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/codegen/IrToJs.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.LoweredIr
 import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity.*
 import org.jetbrains.kotlin.ir.backend.js.export.*
+import org.jetbrains.kotlin.ir.backend.js.lower.JsCodeOutliningLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.StaticMembersLowering
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrFileToJsTransformer
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.processClassModels
@@ -211,6 +212,7 @@ class IrToJs(
         for (fragment in unit.packageFragments) {
             for (declaration in fragment.declarations) {
                 if (declaration is IrDeclarationWithName) {
+                    if (declaration.origin == JsCodeOutliningLowering.OUTLINED_JS_CODE_ORIGIN) continue
                     export(declaration)
                 }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/JsCodeOutliningLowering.kt

@@ -16,22 +16,20 @@ import org.jetbrains.kotlin.ir.backend.js.utils.emptyScope
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.builders.irCall
-import org.jetbrains.kotlin.ir.builders.irComposite
 import org.jetbrains.kotlin.ir.builders.irGet
-import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrContainerExpression
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
-import org.jetbrains.kotlin.ir.util.constructors
+import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 // Outlines `kotlin.js.js(code: String)` calls where JS code references Kotlin locals.
 // Makes locals usages explicit.
@@ -44,6 +42,10 @@ class JsCodeOutliningLowering(val backendContext: JsIrBackendContext) : BodyLowe
         val replacer = JsCodeOutlineTransformer(backendContext, container)
         irBody.transformChildrenVoid(replacer)
     }
+
+    companion object {
+        val OUTLINED_JS_CODE_ORIGIN = object : IrDeclarationOriginImpl("OUTLINED_JS_CODE") {}
+    }
 }
 
 private fun IrElement.containsCallsTo(symbol: IrFunctionSymbol): Boolean {
@@ -144,18 +146,17 @@ private class JsCodeOutlineTransformer(
 
         // Building outlined IR function skeleton
         val outlinedFunction = backendContext.irFactory.buildFun {
-            name = Name.identifier("outlinedJsCode$")
+            name = Name.identifier(container.safeAs<IrDeclarationWithName>()?.name?.asString()?.let { "$it\$outlinedJsCode\$" }
+                                       ?: "outlinedJsCode\$")
             returnType = backendContext.dynamicType
             isExternal = true
-            origin = JsIrBackendContext.callableClosureOrigin
+            origin = JsCodeOutliningLowering.OUTLINED_JS_CODE_ORIGIN
         }
         // We don't need this function's body. Using empty block body stub, because some code might expect all functions to have bodies.
         outlinedFunction.body = backendContext.irFactory.createBlockBody(UNDEFINED_OFFSET, UNDEFINED_OFFSET)
-        outlinedFunction.parent = when (container) {
+        outlinedFunction.parent = container.file
-            is IrDeclarationParent -> container
+        container.file.declarations.add(outlinedFunction)
-            else -> container.parent
+        kotlinLocalsUsedInJs.values.forEach { local ->
-        }
-        kotlinLocalsUsedInJs.forEach { local ->
             outlinedFunction.addValueParameter {
                 name = local.name
                 type = local.type
@@ -175,29 +176,19 @@ private class JsCodeOutlineTransformer(
                 newStatements += JsReturn(JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)))
             }
         }
-        val newFun = JsFunction(emptyScope, JsBlock(newStatements), "")
+        val newFun = JsFunction(emptyScope, JsBlock(newStatements), "Outlined js() call")
-        kotlinLocalsUsedInJs.forEach { irParameter ->
+        kotlinLocalsUsedInJs.keys.forEach { jsName ->
-            newFun.parameters.add(JsParameter(JsName(irParameter.name.identifier, false)))
+            newFun.parameters.add(JsParameter(jsName))
         }
 
-        with(backendContext.createIrBuilder(container.symbol)) {
+        backendContext.addOutlinedJsCode(outlinedFunction.symbol, newFun)
-            // Add @JsFun("function (used_local1, used_local2, ..) { ... }") annotation to outlined function
-            val jsFunCtor = backendContext.intrinsics.jsFunAnnotationSymbol.constructors.single()
-            val jsFunCall =
-                irCall(jsFunCtor).apply {
-                    putValueArgument(0, irString(newFun.toString()))
-                }
-            outlinedFunction.annotations = listOf(jsFunCall)
 
-            val outlinedFunctionCall = irCall(outlinedFunction).apply {
+        return with(backendContext.createIrBuilder(container.symbol)) {
-                kotlinLocalsUsedInJs.forEachIndexed { index, local ->
+            irCall(outlinedFunction).apply {
+                kotlinLocalsUsedInJs.values.forEachIndexed { index, local ->
                     putValueArgument(index, irGet(local))
                 }
             }
-            return irComposite {
-                +outlinedFunction
-                +outlinedFunctionCall
-            }
         }
     }
 }
@@ -249,9 +240,9 @@ private class KotlinLocalsUsageCollector(
 ) : RecursiveJsVisitor() {
     private val functionStack = mutableListOf<JsFunction?>(null)
     private val processedNames = mutableSetOf<String>()
-    private val kotlinLocalsUsedInJs = mutableListOf<IrValueDeclaration>()
+    private val kotlinLocalsUsedInJs = mutableMapOf<JsName, IrValueDeclaration>()
 
-    val usedLocals: List<IrValueDeclaration>
+    val usedLocals: Map<JsName, IrValueDeclaration>
         get() = kotlinLocalsUsedInJs
 
     override fun visitFunction(x: JsFunction) {
@@ -267,7 +258,9 @@ private class KotlinLocalsUsageCollector(
         // We will also collect shadowed usages, but it is OK since the same shadowing will be present in generated JS code.
         // Keeping track of processed names to avoid registering them multiple times
         if (processedNames.add(name.ident) && !name.isDeclaredInsideJsCode()) {
-            kotlinLocalsUsedInJs.addIfNotNull(findValueDeclarationWithName(name.ident))
+            findValueDeclarationWithName(name.ident)?.let {
+                kotlinLocalsUsedInJs[name] = it
+            }
         }
     }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/FunctionWithJsFuncAnnotationInliner.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 import org.jetbrains.kotlin.backend.common.compilationException
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.backend.js.utils.JsGenerationContext
-import org.jetbrains.kotlin.ir.backend.js.utils.getJsFunAnnotation
 import org.jetbrains.kotlin.js.backend.ast.*
 
 class FunctionWithJsFuncAnnotationInliner(private val jsFuncCall: IrCall, private val context: JsGenerationContext) {
@@ -24,13 +23,9 @@ class FunctionWithJsFuncAnnotationInliner(private val jsFuncCall: IrCall, privat
             }
     }
 
-    private fun getJsFunctionImplementation(): JsFunction {
+    private fun getJsFunctionImplementation(): JsFunction =
-        val code = jsFuncCall.symbol.owner.getJsFunAnnotation() ?: compilationException("JsFun annotation is expected", jsFuncCall)
+        context.staticContext.backendContext.getJsCodeForFunction(jsFuncCall.symbol)?.deepCopy()
-        val statements = parseJsCode(code) ?: compilationException("Cannot compute js code", jsFuncCall)
+            ?: compilationException("JS function not found", jsFuncCall)
-        return statements.singleOrNull()
-            ?.let { it as? JsExpressionStatement }
-            ?.let { it.expression as? JsFunction } ?: compilationException("Provided js code is not a js function", jsFuncCall.symbol.owner)
-    }
 
     private fun collectReplacementsForCall(): Map<JsName, JsExpression> {
         val translatedArguments = Array(jsFuncCall.valueArgumentsCount) {
@@ -72,4 +67,4 @@ private class SimpleJsCodeInliner(private val replacements: Map<JsName, JsExpres
             else -> declareNewTemporaryFor(expression)
         }
     }
-}
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrDeclarationToJsTransformer.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 
+import org.jetbrains.kotlin.ir.backend.js.lower.JsCodeOutliningLowering
 import org.jetbrains.kotlin.ir.backend.js.utils.JsGenerationContext
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.js.backend.ast.*
@@ -14,6 +15,7 @@ class IrDeclarationToJsTransformer : BaseIrElementToJsNodeTransformer<JsStatemen
 
     override fun visitSimpleFunction(declaration: IrSimpleFunction, context: JsGenerationContext): JsStatement {
         require(!declaration.isExpect)
+        if (declaration.origin == JsCodeOutliningLowering.OUTLINED_JS_CODE_ORIGIN) return JsEmpty
         return declaration.accept(IrFunctionToJsTransformer(), context).makeStmt()
     }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsExpressionTransformer.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.backend.common.compilationException
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.ir.backend.js.utils.JsGenerationContext
 import org.jetbrains.kotlin.ir.backend.js.utils.Namer
-import org.jetbrains.kotlin.ir.backend.js.utils.emptyScope
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
@@ -214,7 +213,7 @@ class IrElementToJsExpressionTransformer : BaseIrElementToJsNodeTransformer<JsEx
     }
 
     override fun visitCall(expression: IrCall, context: JsGenerationContext): JsExpression {
-        if (context.checkIfJsCode(expression.symbol) || context.checkIfAnnotatedWithJsFunc(expression.symbol)) {
+        if (context.checkIfJsCode(expression.symbol) || context.checkIfHasAssociatedJsCode(expression.symbol)) {
             return JsCallTransformer(expression, context).generateExpression()
         }
         return translateCall(expression, context, this).withSource(expression, context)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt

@@ -17,7 +17,6 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.types.classifierOrNull
-import org.jetbrains.kotlin.ir.types.classOrNull
 import org.jetbrains.kotlin.ir.types.isAny
 import org.jetbrains.kotlin.ir.util.constructedClassType
 import org.jetbrains.kotlin.ir.util.file
@@ -165,7 +164,7 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer<JsSta
         if (expression.symbol.isUnitInstanceFunction(data)) {
             return JsEmpty
         }
-        if (data.checkIfJsCode(expression.symbol) || data.checkIfAnnotatedWithJsFunc(expression.symbol)) {
+        if (data.checkIfJsCode(expression.symbol) || data.checkIfHasAssociatedJsCode(expression.symbol)) {
             return JsCallTransformer(expression, data).generateStatement()
         }
         return translateCall(expression, data, IrElementToJsExpressionTransformer()).withSource(expression, data).makeStmt()

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsCallTransformer.kt

@@ -74,10 +74,10 @@ class JsCallTransformer(private val jsOrJsFuncCall: IrCall, private val context:
                     ?: compilationException("Cannot compute js code", jsOrJsFuncCall)
             }
 
-            context.checkIfAnnotatedWithJsFunc(jsOrJsFuncCall.symbol) ->
+            context.checkIfHasAssociatedJsCode(jsOrJsFuncCall.symbol) ->
                 FunctionWithJsFuncAnnotationInliner(jsOrJsFuncCall, context).generateResultStatement()
 
             else -> compilationException("`js` function call or function with @JsFunc annotation expected", jsOrJsFuncCall)
         }
     }
-}
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsPolyfills.kt

@@ -34,6 +34,6 @@ class JsPolyfills {
     private fun Sequence<IrDeclaration>.asImplementationList(): List<JsStatement> =
         map { it.getJsPolyfill()!! }
             .distinct()
-            .flatMap { parseJsCode(it).orEmpty() }
+            .flatMap { parseJsCode(it).orEmpty() } // FIXME!!!
             .toList()
-}
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/jsAstUtils.kt

@@ -55,10 +55,7 @@ fun jsAssignment(left: JsExpression, right: JsExpression) = JsBinaryOperation(Js
 fun prototypeOf(classNameRef: JsExpression) = JsNameRef(Namer.PROTOTYPE_NAME, classNameRef)
 
 fun translateFunction(declaration: IrFunction, name: JsName?, context: JsGenerationContext): JsFunction {
-    val jsFun = declaration.getJsFunAnnotation()
+    context.staticContext.backendContext.getJsCodeForFunction(declaration.symbol)?.let { function ->
-    if (jsFun != null) {
-        // JsFun internal annotation must have string containing valid function expression
-        val function = (parseJsCode(jsFun)!!.single() as JsExpressionStatement).expression as JsFunction
         function.name = name
         return function
     }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/jsCode.kt

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.js.parser.parseExpressionOrStatement
 // Returns null if constant expression could not be parsed
 fun translateJsCodeIntoStatementList(code: IrExpression, context: JsIrBackendContext): List<JsStatement>? {
     // TODO: support proper symbol linkage and label clash resolution
-
+    // FIXME: Provide debug info
     return parseJsCode(foldString(code, context) ?: return null)
 }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/AnnotationUtils.kt

@@ -57,9 +57,6 @@ fun IrAnnotationContainer.getJsPolyfill(): String? =
 fun IrAnnotationContainer.hasJsPolyfill(): Boolean =
     hasAnnotation(JsAnnotations.JsPolyfillFqn)
 
-fun IrAnnotationContainer.getJsFunAnnotation(): String? =
-    getAnnotation(JsAnnotations.jsFunFqn)?.getSingleConstStringArgument()
-
 fun IrAnnotationContainer.isJsExport(): Boolean =
     hasAnnotation(JsAnnotations.jsExportFqn)
 
@@ -72,8 +69,6 @@ fun IrAnnotationContainer.isJsNativeSetter(): Boolean = hasAnnotation(JsAnnotati
 
 fun IrAnnotationContainer.isJsNativeInvoke(): Boolean = hasAnnotation(JsAnnotations.jsNativeInvoke)
 
-fun IrAnnotationContainer.isAnnotatedWithJsFun(): Boolean = hasAnnotation(JsAnnotations.jsFunFqn)
-
 private fun IrOverridableDeclaration<*>.dfsOverridableJsNameOrNull(): String? {
     for (overriddenSymbol in overriddenSymbols) {
         val symbolOwner = overriddenSymbol.owner

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/JsGenerationContext.kt

@@ -78,5 +78,5 @@ class JsGenerationContext(
 
     fun checkIfJsCode(symbol: IrFunctionSymbol): Boolean = symbol == staticContext.backendContext.intrinsics.jsCode
 
-    fun checkIfAnnotatedWithJsFunc(symbol: IrFunctionSymbol): Boolean = symbol.owner.isAnnotatedWithJsFun()
+    fun checkIfHasAssociatedJsCode(symbol: IrFunctionSymbol): Boolean = staticContext.backendContext.getJsCodeForFunction(symbol) != null
 }

js/js.translator/testData/box/inline/jsCode.kt

@@ -1,8 +1,7 @@
 // EXPECTED_REACHABLE_NODES: 1283
 package foo
 
-// In the IR backend the injected JS code is outlined.
+// CHECK_CONTAINS_NO_CALLS: test
-// CHECK_CONTAINS_NO_CALLS: test TARGET_BACKENDS=JS
 // CHECK_NOT_CALLED_IN_SCOPE: function=sum scope=test
 
 internal inline fun sum(x: Int, y: Int): Int = js("x + y")
@@ -16,4 +15,4 @@ fun box(): String {
     assertEquals(8, test(1, 3))
 
     return "OK"
-}
+}

js/js.translator/testData/box/inline/jsCodeVarDeclared.kt

@@ -1,8 +1,7 @@
 // EXPECTED_REACHABLE_NODES: 1283
 package foo
 
-// In the IR backend the injected JS code is outlined.
+// CHECK_CONTAINS_NO_CALLS: test
-// CHECK_CONTAINS_NO_CALLS: test TARGET_BACKENDS=JS
 // CHECK_NOT_CALLED_IN_SCOPE: function=sum scope=test
 
 internal inline fun sum(x: Int, y: Int): Int = js("var a = x; a + y")
@@ -22,4 +21,4 @@ fun box(): String {
     assertEquals(8, test(1, 3))
 
     return "OK"
-}
+}