feat(JsCode): inline usage of js code which captures kotlin variables.

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.ir.backend.js.lower
 
 import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
@@ -132,41 +131,23 @@ private class JsCodeOutlineTransformer(
         if (expression.symbol != backendContext.intrinsics.jsCode)
             return null
 
-        val jsCodeArg = expression.getValueArgument(0)
-            ?: compilationException(
-                "Expected js code string",
-                expression
-            )
+        val jsCodeArg = expression.getValueArgument(0) ?: compilationException("Expected js code string", expression)
         val jsStatements = translateJsCodeIntoStatementList(jsCodeArg, backendContext) ?: return null
 
         // Collect used Kotlin local variables and parameters.
-        val kotlinLocalsUsedInJs = mutableListOf<IrValueDeclaration>()
-        val processedNames = mutableSetOf<String>()
-        jsStatements.forEach { statement ->
-            object : RecursiveJsVisitor() {
-                override fun visitNameRef(nameRef: JsNameRef) {
-                    super.visitNameRef(nameRef)
-                    val name = nameRef.name
-                    // With this approach we should be able to find all usages of Kotlin variables in JS code.
-                    // We will also collect shadowed usages, but it is OK since the same shadowing will be present in generated JS code.
-                    if (name != null && nameRef.qualifier == null) {
-                        // Keeping track of processed names to avoid registering them multiple times
-                        if (processedNames.add(name.ident)) {
-                            kotlinLocalsUsedInJs.addIfNotNull(findValueDeclarationWithName(name.ident))
-                        }
-                    }
-                }
-            }.accept(statement)
-        }
+        val scope = JsScopesCollector().apply { acceptList(jsStatements) }
+        val localsUsageCollector = KotlinLocalsUsageCollector(scope, ::findValueDeclarationWithName).apply { acceptList(jsStatements) }
+        val kotlinLocalsUsedInJs = localsUsageCollector.usedLocals
+
         if (kotlinLocalsUsedInJs.isEmpty())
             return null
 
         // Building outlined IR function skeleton
         val outlinedFunction = backendContext.irFactory.buildFun {
             name = Name.identifier("outlinedJsCode$")
-            visibility = DescriptorVisibilities.LOCAL
             returnType = backendContext.dynamicType
-            origin = OUTLINED_ORIGIN
+            isExternal = true
+            origin = JsIrBackendContext.callableClosureOrigin
         }
         // 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)
@@ -219,8 +200,78 @@ private class JsCodeOutlineTransformer(
             }
         }
     }
+}
 
-    companion object {
-        object OUTLINED_ORIGIN : IrDeclarationOriginImpl("OUTLINED_ORIGIN")
+class JsScopesCollector : RecursiveJsVisitor() {
+    private val functionsStack = mutableListOf(Scope(null))
+    private val functionalScopes = mutableMapOf<JsFunction?, Scope>(null to functionsStack.first())
+
+    private class Scope(val parent: Scope?) {
+        private val variables = hashSetOf<String>()
+
+        fun add(variableName: String) {
+            variables.add(variableName)
+        }
+
+        fun variableWithNameExists(variableName: String): Boolean {
+            return variables.contains(variableName) ||
+                    parent?.variableWithNameExists(variableName) == true
+        }
     }
-}
+
+    override fun visitVars(x: JsVars) {
+        super.visitVars(x)
+        val currentScope = functionsStack.last()
+        x.vars.forEach { currentScope.add(it.name.ident) }
+    }
+
+    override fun visitFunction(x: JsFunction) {
+        val parentScope = functionsStack.last()
+        val newScope = Scope(parentScope).apply {
+            val name = x.name?.ident
+            if (name != null) add(name)
+            x.parameters.forEach { add(it.name.ident) }
+        }
+        functionsStack.push(newScope)
+        functionalScopes[x] = newScope
+        super.visitFunction(x)
+        functionsStack.pop()
+    }
+
+    fun varWithNameExistsInScopeOf(function: JsFunction?, variableName: String): Boolean {
+        return functionalScopes[function]!!.variableWithNameExists(variableName)
+    }
+}
+
+private class KotlinLocalsUsageCollector(
+    private val scopeInfo: JsScopesCollector,
+    private val findValueDeclarationWithName: (String) -> IrValueDeclaration?
+) : RecursiveJsVisitor() {
+    private val functionStack = mutableListOf<JsFunction?>(null)
+    private val processedNames = mutableSetOf<String>()
+    private val kotlinLocalsUsedInJs = mutableListOf<IrValueDeclaration>()
+
+    val usedLocals: List<IrValueDeclaration>
+        get() = kotlinLocalsUsedInJs
+
+    override fun visitFunction(x: JsFunction) {
+        functionStack.push(x)
+        super.visitFunction(x)
+        functionStack.pop()
+    }
+
+    override fun visitNameRef(nameRef: JsNameRef) {
+        super.visitNameRef(nameRef)
+        val name = nameRef.name.takeIf { nameRef.qualifier == null } ?: return
+        // With this approach we should be able to find all usages of Kotlin variables in JS code.
+        // 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))
+        }
+    }
+
+    private fun JsName.isDeclaredInsideJsCode(): Boolean {
+        return scopeInfo.varWithNameExistsInScopeOf(functionStack.peek(), ident)
+    }
+}

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

@@ -0,0 +1,75 @@
+/*
+ * 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 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) {
+    private val function = getJsFunctionImplementation()
+    private val replacements = collectReplacementsForCall()
+
+    fun generateResultStatement(): List<JsStatement> {
+        return function.body.statements
+            .run {
+                SimpleJsCodeInliner(replacements)
+                    .apply { acceptList(this@run) }
+                    .withTemporaryVariablesForExpressions(this)
+            }
+    }
+
+    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 collectReplacementsForCall(): Map<JsName, JsExpression> {
+        val translatedArguments = Array(jsFuncCall.valueArgumentsCount) {
+            jsFuncCall.getValueArgument(it)!!.accept(IrElementToJsExpressionTransformer(), context)
+        }
+        return function.parameters
+            .mapIndexed { i, param -> param.name to translatedArguments[i] }
+            .toMap()
+    }
+}
+
+private class SimpleJsCodeInliner(private val replacements: Map<JsName, JsExpression>): RecursiveJsVisitor() {
+    private val temporaryNamesForExpressions = mutableMapOf<JsName, JsExpression>()
+
+    fun withTemporaryVariablesForExpressions(statements: List<JsStatement>): List<JsStatement> {
+        if (temporaryNamesForExpressions.isEmpty()) {
+            return statements
+        }
+
+        val variableDeclarations = temporaryNamesForExpressions.map { JsVars(JsVars.JsVar(it.key, it.value)) }
+        return variableDeclarations + statements
+    }
+
+    override fun visitNameRef(nameRef: JsNameRef) {
+        super.visitNameRef(nameRef)
+        if (nameRef.qualifier != null) return
+        nameRef.name = nameRef.name?.getReplacement() ?: return
+    }
+
+    private fun JsName.declareNewTemporaryFor(expression: JsExpression): JsName {
+        return JsName(ident, true)
+            .also { temporaryNamesForExpressions[it] = expression }
+    }
+
+    private fun JsName.getReplacement(): JsName? {
+        val expression = replacements[this] ?: return null
+        return when {
+            expression is JsNameRef && expression.qualifier == null -> expression.name!!
+            else -> declareNewTemporaryFor(expression)
+        }
+    }
+}

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

@@ -221,41 +221,8 @@ class IrElementToJsExpressionTransformer : BaseIrElementToJsNodeTransformer<JsEx
     }
 
     override fun visitCall(expression: IrCall, context: JsGenerationContext): JsExpression {
-        if (context.checkIfJsCode(expression.symbol)) {
-            val statements = translateJsCodeIntoStatementList(
-                expression.getValueArgument(0)
-                    ?: compilationException(
-                        "JsCode is expected",
-                        expression
-                    ),
-                context.staticContext.backendContext
-            ) ?: compilationException(
-                "Cannot compute js code",
-                expression
-            )
-
-            if (statements.isEmpty()) return JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)) // TODO: report warning or even error
-
-            val lastStatement = statements.last()
-            if (statements.size == 1) {
-                if (lastStatement is JsExpressionStatement) return lastStatement.expression.withSource(expression, context)
-            }
-
-            val newStatements = statements.toMutableList()
-
-            when (lastStatement) {
-                is JsReturn -> {
-                }
-                is JsExpressionStatement -> {
-                    newStatements[statements.lastIndex] = JsReturn(lastStatement.expression)
-                }
-                // TODO: report warning or even error
-                else -> newStatements += JsReturn(JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)))
-            }
-
-            val syntheticFunction = JsFunction(emptyScope, JsBlock(newStatements), "")
-            return JsInvocation(syntheticFunction).withSource(expression, context)
-
+        if (context.checkIfJsCode(expression.symbol) || context.checkIfAnnotatedWithJsFunc(expression.symbol)) {
+            return JsCallTransformer(expression, context).generateExpression()
         }
         return translateCall(expression, context, this).withSource(expression, context)
             .also {

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

@@ -135,24 +135,8 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer<JsSta
     }
 
     override fun visitCall(expression: IrCall, data: JsGenerationContext): JsStatement {
-        if (data.checkIfJsCode(expression.symbol)) {
-            val statements = translateJsCodeIntoStatementList(
-                expression.getValueArgument(0)
-                    ?: compilationException(
-                        "JsCode is expected",
-                        expression
-                    ),
-                data.staticContext.backendContext
-            ) ?: compilationException(
-                "Cannot compute js code",
-                expression
-            )
-            return when (statements.size) {
-                0 -> JsEmpty
-                1 -> statements.single().withSource(expression, data)
-                // TODO: use transparent block (e.g. JsCompositeBlock)
-                else -> JsBlock(statements)
-            }
+        if (data.checkIfJsCode(expression.symbol) || data.checkIfAnnotatedWithJsFunc(expression.symbol)) {
+            return JsCallTransformer(expression, data).generateStatement()
         }
         return translateCall(expression, data, IrElementToJsExpressionTransformer()).withSource(expression, data).makeStmt()
             .also { data.staticContext.polyfills.visitDeclaration(expression.symbol.owner) }

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

@@ -0,0 +1,83 @@
+/*
+ * 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 org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
+
+import org.jetbrains.kotlin.backend.common.compilationException
+import org.jetbrains.kotlin.ir.backend.js.utils.JsGenerationContext
+import org.jetbrains.kotlin.ir.backend.js.utils.emptyScope
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.js.backend.ast.*
+
+class JsCallTransformer(private val jsOrJsFuncCall: IrCall, private val context: JsGenerationContext) {
+    private val statements = getJsStatements()
+
+    fun generateStatement(): JsStatement {
+        if (statements.isEmpty()) return JsEmpty
+
+        val newStatements = statements.toMutableList().apply {
+            val expression = (last() as? JsReturn)?.expression ?: return@apply
+
+            if (expression is JsPrefixOperation && expression.operator == JsUnaryOperator.VOID) {
+                removeLastOrNull()
+            } else {
+                set(lastIndex, expression.makeStmt())
+            }
+        }
+
+        return when (newStatements.size) {
+            0 -> JsEmpty
+            1 -> newStatements.single().withSource(jsOrJsFuncCall, context)
+            // TODO: use transparent block (e.g. JsCompositeBlock)
+            else -> JsBlock(newStatements)
+        }
+    }
+
+    fun generateExpression(): JsExpression {
+        if (statements.isEmpty()) return JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)) // TODO: report warning or even error
+
+        val lastStatement = statements.last()
+        val lastExpression = when (lastStatement) {
+            is JsReturn -> lastStatement.expression
+            is JsExpressionStatement -> lastStatement.expression
+            else -> null
+        }
+        if (statements.size == 1 && lastExpression != null) {
+            return lastExpression.withSource(jsOrJsFuncCall, context)
+        }
+
+        val newStatements = statements.toMutableList()
+
+        when (lastStatement) {
+            is JsReturn -> {
+            }
+            is JsExpressionStatement -> {
+                newStatements[statements.lastIndex] = JsReturn(lastStatement.expression)
+            }
+            // TODO: report warning or even error
+            else -> newStatements += JsReturn(JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)))
+        }
+
+        val syntheticFunction = JsFunction(emptyScope, JsBlock(newStatements), "")
+        return JsInvocation(syntheticFunction).withSource(jsOrJsFuncCall, context)
+    }
+
+    private fun getJsStatements(): List<JsStatement> {
+        return when {
+            context.checkIfJsCode(jsOrJsFuncCall.symbol) -> {
+                translateJsCodeIntoStatementList(
+                    jsOrJsFuncCall.getValueArgument(0) ?: compilationException("JsCode is expected", jsOrJsFuncCall),
+                    context.staticContext.backendContext
+                )
+                    ?: compilationException("Cannot compute js code", jsOrJsFuncCall)
+            }
+
+            context.checkIfAnnotatedWithJsFunc(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/utils/AnnotationUtils.kt

@@ -61,6 +61,8 @@ fun IrAnnotationContainer.isJsNativeSetter(): Boolean = hasAnnotation(JsAnnotati
 
 fun IrAnnotationContainer.isJsNativeInvoke(): Boolean = hasAnnotation(JsAnnotations.jsNativeInvoke)
 
+fun IrAnnotationContainer.isAnnotatedWithJsFun(): Boolean = hasAnnotation(JsAnnotations.jsFunFqn)
+
 fun IrDeclarationWithName.getJsNameForOverriddenDeclaration(): String? {
     val jsName = getJsName()
 

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

@@ -77,4 +77,6 @@ class JsGenerationContext(
     }
 
     fun checkIfJsCode(symbol: IrFunctionSymbol): Boolean = symbol == staticContext.backendContext.intrinsics.jsCode
+
+    fun checkIfAnnotatedWithJsFunc(symbol: IrFunctionSymbol): Boolean = symbol.owner.isAnnotatedWithJsFun()
 }

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

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JS_IR
 // EXPECTED_REACHABLE_NODES: 1282
 package foo
 

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

@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JS_IR
-// IGNORE_BACKEND: JS_IR_ES6
 // EXPECTED_REACHABLE_NODES: 1282
 package foo
 

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

@@ -1,5 +1,4 @@
 // EXPECTED_REACHABLE_NODES: 1282
-// IGNORE_BACKEND: JS_IR
 package foo
 
 fun box(): String {

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

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JS_IR
 // EXPECTED_REACHABLE_NODES: 1285
 package foo
 

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

@@ -1,5 +1,4 @@
 // EXPECTED_REACHABLE_NODES: 1282
-// IGNORE_BACKEND: JS_IR
 package foo
 
 // CHECK_LABELS_COUNT: function=box name=block count=2

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

@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JS_IR
-// IGNORE_BACKEND: JS_IR_ES6
 // KJS_WITH_FULL_RUNTIME
 // EXPECTED_REACHABLE_NODES: 1687
 external fun p(m: String): String
@@ -26,9 +24,15 @@ fun test4(): String {
     return js("p('test4')")
 }
 
+fun f() = js("p('test5')")
+
 fun test5(): String {
     val p = "wrong5"
-    fun f() = js("p('test5')")
+    // The behavoiur of the classical backend is weird and buggy
+    // From the user side, the local variable `p` is captured
+    // but we have different behaviour because the renaming phase in classical backend
+    // will be invoked after the lambda will be moved up
+    // fun f() = js("p('test5')")
     return f()
 }