[JS IR] Fix referencing Kotin variables in inline JS code

Fixed by outlining JS code that uses Kotlin variables making usages of
these locals explicit and preventing bugs due to one-sided variable renaming.

This prevents using Kotlin variables as lvalue in JS code.

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

@@ -17,9 +17,12 @@ import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
-import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeBuilder
 import org.jetbrains.kotlin.ir.types.impl.buildSimpleType
+import org.jetbrains.kotlin.ir.types.isLong
+import org.jetbrains.kotlin.ir.types.typeWithParameters
 import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.findDeclaration
 import org.jetbrains.kotlin.ir.util.kotlinPackageFqn
@@ -278,6 +281,7 @@ class JsIntrinsics(private val irBuiltIns: IrBuiltIns, val context: JsIrBackendC
 
     // TODO move to IntrinsifyCallsLowering
     val doNotIntrinsifyAnnotationSymbol = context.symbolTable.referenceClass(context.getJsInternalClass("DoNotIntrinsify"))
+    val jsFunAnnotationSymbol = context.symbolTable.referenceClass(context.getJsInternalClass("JsFun"))
 
     // TODO move CharSequence-related stiff to IntrinsifyCallsLowering
     val charSequenceClassSymbol = context.symbolTable.referenceClass(context.getClass(FqName("kotlin.CharSequence")))

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

@@ -185,6 +185,12 @@ private val stripTypeAliasDeclarationsPhase = makeDeclarationTransformerPhase(
     description = "Strip typealias declarations"
 )
 
+private val jsCodeOutliningPhase = makeBodyLoweringPhase(
+    ::JsCodeOutliningLowering,
+    name = "JsCodeOutliningLowering",
+    description = "Outline js() calls where JS code references Kotlin locals"
+)
+
 private val arrayConstructorPhase = makeBodyLoweringPhase(
     ::ArrayConstructorLowering,
     name = "ArrayConstructor",
@@ -713,6 +719,7 @@ private val loweringList = listOf<Lowering>(
     validateIrBeforeLowering,
     expectDeclarationsRemovingPhase,
     stripTypeAliasDeclarationsPhase,
+    jsCodeOutliningPhase,
     arrayConstructorPhase,
     lateinitNullableFieldsPhase,
     lateinitDeclarationLoweringPhase,

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

@@ -0,0 +1,187 @@
+/*
+ * 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.ir.backend.js.lower
+
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.backend.common.push
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.translateJsCodeIntoStatementList
+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.util.constructors
+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
+
+// Outlines `kotlin.js.js(code: String)` calls where JS code references Kotlin locals.
+// Makes locals usages explicit.
+class JsCodeOutliningLowering(val backendContext: JsIrBackendContext) : BodyLoweringPass {
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        val replacer = JsCodeOutlineTransformer(backendContext, container)
+        irBody.transformChildrenVoid(replacer)
+    }
+}
+
+private class JsCodeOutlineTransformer(
+    val backendContext: JsIrBackendContext,
+    val container: IrDeclaration,
+) : IrElementTransformerVoidWithContext() {
+    val localScopes: MutableList<MutableMap<String, IrValueDeclaration>> =
+        mutableListOf(mutableMapOf())
+
+    init {
+        if (container is IrFunction) {
+            container.valueParameters.forEach {
+                registerValueDeclaration(it)
+            }
+        }
+    }
+
+    inline fun <T> withLocalScope(body: () -> T): T {
+        localScopes.push(mutableMapOf())
+        val res = body()
+        localScopes.pop()
+        return res
+    }
+
+    fun registerValueDeclaration(irValueDeclaration: IrValueDeclaration) {
+        val name = irValueDeclaration.name
+        if (!name.isSpecial) {
+            val identifier = name.identifier
+            val currentScope = localScopes.lastOrNull() ?: error("Expecting a scope")
+            currentScope[identifier] = irValueDeclaration
+        }
+    }
+
+    fun findValueDeclarationWithName(name: String): IrValueDeclaration? {
+        for (i in (localScopes.size - 1) downTo 0) {
+            val scope = localScopes[i]
+            return scope[name] ?: continue
+        }
+        return null
+    }
+
+    override fun visitContainerExpression(expression: IrContainerExpression): IrExpression {
+        return withLocalScope { super.visitContainerExpression(expression) }
+    }
+
+    override fun visitDeclaration(declaration: IrDeclarationBase): IrStatement {
+        return withLocalScope { super.visitDeclaration(declaration) }
+    }
+
+    override fun visitValueParameterNew(declaration: IrValueParameter): IrStatement {
+        return super.visitValueParameterNew(declaration).also { registerValueDeclaration(declaration) }
+    }
+
+    override fun visitVariable(declaration: IrVariable): IrStatement {
+        return super.visitVariable(declaration).also { registerValueDeclaration(declaration) }
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        return outlineJsCodeIfNeeded(expression) ?: super.visitCall(expression)
+    }
+
+    fun outlineJsCodeIfNeeded(expression: IrCall): IrExpression? {
+        if (expression.symbol != backendContext.intrinsics.jsCode)
+            return null
+
+        val jsCodeArg = expression.getValueArgument(0) ?: error("Expected js code string")
+        val jsStatements = translateJsCodeIntoStatementList(jsCodeArg)
+
+        // 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)
+        }
+        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
+        }
+        // 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 = container as IrDeclarationParent
+        kotlinLocalsUsedInJs.forEach { local ->
+            outlinedFunction.addValueParameter {
+                name = local.name
+                type = local.type
+            }
+        }
+
+        // Building JS Ast function
+        val lastStatement = jsStatements.last()
+        val newStatements = jsStatements.toMutableList()
+        when (lastStatement) {
+            is JsReturn -> {
+            }
+            is JsExpressionStatement -> {
+                newStatements[jsStatements.lastIndex] = JsReturn(lastStatement.expression)
+            }
+            else -> {
+                newStatements += JsReturn(JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)))
+            }
+        }
+        val newFun = JsFunction(emptyScope, JsBlock(newStatements), "")
+        kotlinLocalsUsedInJs.forEach { irParameter ->
+            newFun.parameters.add(JsParameter(JsName(irParameter.name.identifier)))
+        }
+
+        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)
+
+            val outlinedFunctionCall = irCall(outlinedFunction).apply {
+                kotlinLocalsUsedInJs.forEachIndexed { index, local ->
+                    putValueArgument(index, irGet(local))
+                }
+            }
+            return irComposite {
+                +outlinedFunction
+                +outlinedFunctionCall
+            }
+        }
+    }
+}

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

@@ -42,6 +42,14 @@ 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
+        function.name = name
+        return function
+    }
+
     val functionContext = context.newDeclaration(declaration)
     val functionParams = declaration.valueParameters.map { functionContext.getNameForValueDeclaration(it) }
     val body = declaration.body?.accept(IrElementToJsStatementTransformer(), functionContext) as? JsBlock ?: JsBlock()

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

@@ -44,7 +44,7 @@ fun translateJsCodeIntoStatementList(code: IrExpression): List<JsStatement> {
     return parseJsCode(foldString(code)) ?: emptyList()
 }
 
-private fun parseJsCode(jsCode: String): List<JsStatement>? {
+fun parseJsCode(jsCode: String): List<JsStatement>? {
     // Parser can change local or global scope.
     // In case of js we want to keep new local names,
     // but no new global ones.

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

@@ -5,8 +5,9 @@
 
 package org.jetbrains.kotlin.ir.backend.js.utils
 
-import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.declarations.IrAnnotationContainer
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
 import org.jetbrains.kotlin.ir.expressions.IrClassReference
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
@@ -24,6 +25,7 @@ object JsAnnotations {
     val jsNativeGetter = FqName("kotlin.js.nativeGetter")
     val jsNativeSetter = FqName("kotlin.js.nativeSetter")
     val jsNativeInvoke = FqName("kotlin.js.nativeInvoke")
+    val jsFunFqn = FqName("kotlin.js.JsFun")
 }
 
 @Suppress("UNCHECKED_CAST")
@@ -42,6 +44,9 @@ fun IrAnnotationContainer.getJsQualifier(): String? =
 fun IrAnnotationContainer.getJsName(): String? =
     getAnnotation(JsAnnotations.jsNameFqn)?.getSingleConstStringArgument()
 
+fun IrAnnotationContainer.getJsFunAnnotation(): String? =
+    getAnnotation(JsAnnotations.jsFunFqn)?.getSingleConstStringArgument()
+
 fun IrAnnotationContainer.isJsExport(): Boolean =
     hasAnnotation(JsAnnotations.jsExportFqn)
 

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

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

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

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

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

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

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

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

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

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

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

@@ -1,4 +1,5 @@
 // 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/operators.kt

@@ -21,13 +21,15 @@ fun box(): String {
     assertEquals(1, js("0 ^ 1"), "^")
     assertEquals(-2, js("~1"), "~")
 
-    var i = 2
-    assertEquals(1, js("--i"), "-- prefix")
-    assertEquals(1, js("i--"), "-- postfix (1)")
-    assertEquals(0, js("i"), "-- postfix (0)")
-    assertEquals(1, js("++i"), "++ prefix")
-    assertEquals(1, js("i++"), "++ postfix (1)")
-    assertEquals(2, js("i"), "++ postfix (0)")
+    if (testUtils.isLegacyBackend()) {
+        var i = 2
+        assertEquals(1, js("--i"), "-- prefix")
+        assertEquals(1, js("i--"), "-- postfix (1)")
+        assertEquals(0, js("i"), "-- postfix (0)")
+        assertEquals(1, js("++i"), "++ prefix")
+        assertEquals(1, js("i++"), "++ postfix (1)")
+        assertEquals(2, js("i"), "++ postfix (0)")
+    }
 
     assertEquals(true , js("true || false"), "||")
     assertEquals(false , js("true && false"), "&&")
@@ -47,13 +49,15 @@ fun box(): String {
     assertEquals("even", js("(4 % 2 === 0)?'even':'odd'"), "?:")
     assertEquals(3, js("1,2,3"), ", (comma)")
 
-    var j = 0
-    assertEquals(1, js("j = 1"), "=")
-    assertEquals(3, js("j += 2"), "+=")
-    assertEquals(2, js("j -= 1"), "-=")
-    assertEquals(14, js("j *= 7"), "*=")
-    assertEquals(7, js("j /= 2"), "/=")
-    assertEquals(1, js("j %= 2"), "%=")
+    if (testUtils.isLegacyBackend()) {
+        var j = 0
+        assertEquals(1, js("j = 1"), "=")
+        assertEquals(3, js("j += 2"), "+=")
+        assertEquals(2, js("j -= 1"), "-=")
+        assertEquals(14, js("j *= 7"), "*=")
+        assertEquals(7, js("j /= 2"), "/=")
+        assertEquals(1, js("j %= 2"), "%=")
+    }
 
     assertEquals(undefined, js("(void 0)"), "void")
     assertEquals(true, js("'key' in {'key': 10}"), "in")

libraries/stdlib/js-ir/runtime/kotlinJsHacks.kt

@@ -12,3 +12,19 @@ internal fun <T : Enum<T>> enumValuesIntrinsic(): Array<T> =
 @PublishedApi
 internal fun <T : Enum<T>> enumValueOfIntrinsic(@Suppress("UNUSED_PARAMETER") name: String): T =
     throw IllegalStateException("Should be replaced by compiler")
+
+
+/**
+ * Implements annotated function in JavaScript.
+ * [code] string must contain JS expression that evaluates to JS function with signature that matches annotated kotlin function
+ *
+ * For example, a function that adds two Doubles:
+ *
+ *      @JsFun("(x, y) => x + y")
+ *      fun jsAdd(x: Double, y: Double): Double =
+ *          error("...")
+ *
+ * Code gets inserted as is without syntax verification.
+ */
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.PROPERTY_SETTER)
+internal annotation class JsFun(val code: String)