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

#KT-51327 Fixed
2022-07-27 22:36:45 +02:00
parent a90c4d5dd5
commit a152bc2313
15 changed files with 75 additions and 72 deletions
@@ -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.utils.MinimizedNameGenerator
-import org.jetbrains.kotlin.ir.backend.js.utils.OperatorNames
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.parseJsCode
+import org.jetbrains.kotlin.ir.backend.js.utils.*
 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.IrClassifierSymbol
-import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
-import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.*
 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
+    }
 }
@@ -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)
                }

@@ -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) {
-            is IrDeclarationParent -> container
-            else -> container.parent
-        }
-        kotlinLocalsUsedInJs.forEach { local ->
+        outlinedFunction.parent = container.file
+        container.file.declarations.add(outlinedFunction)
+        kotlinLocalsUsedInJs.values.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), "")
-        kotlinLocalsUsedInJs.forEach { irParameter ->
-            newFun.parameters.add(JsParameter(JsName(irParameter.name.identifier, false)))
+        val newFun = JsFunction(emptyScope, JsBlock(newStatements), "Outlined js() call")
+        kotlinLocalsUsedInJs.keys.forEach { jsName ->
+            newFun.parameters.add(JsParameter(jsName))
        }

-        with(backendContext.createIrBuilder(container.symbol)) {
-            // 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)
+        backendContext.addOutlinedJsCode(outlinedFunction.symbol, newFun)

-            val outlinedFunctionCall = irCall(outlinedFunction).apply {
-                kotlinLocalsUsedInJs.forEachIndexed { index, local ->
+        return with(backendContext.createIrBuilder(container.symbol)) {
+            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
+            }
        }
    }

@@ -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 {
-        val code = jsFuncCall.symbol.owner.getJsFunAnnotation() ?: compilationException("JsFun annotation is expected", jsFuncCall)
-        val statements = parseJsCode(code) ?: compilationException("Cannot compute js code", 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 getJsFunctionImplementation(): JsFunction =
+        context.staticContext.backendContext.getJsCodeForFunction(jsFuncCall.symbol)?.deepCopy()
+            ?: compilationException("JS function not found", jsFuncCall)

    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)
        }
    }
-}
+}
@@ -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()
    }

@@ -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)
@@ -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()
@@ -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)
        }
    }
-}
+}
@@ -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()
-}
+}
@@ -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()
-    if (jsFun != null) {
-        // JsFun internal annotation must have string containing valid function expression
-        val function = (parseJsCode(jsFun)!!.single() as JsExpressionStatement).expression as JsFunction
+    context.staticContext.backendContext.getJsCodeForFunction(declaration.symbol)?.let { function ->
        function.name = name
        return function
    }
@@ -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)
 }

@@ -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
@@ -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
 }
@@ -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 TARGET_BACKENDS=JS
+// CHECK_CONTAINS_NO_CALLS: test
 // 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"
-}
+}
@@ -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 TARGET_BACKENDS=JS
+// CHECK_CONTAINS_NO_CALLS: test
 // 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"
-}
+}