[JS IR] Lower lambdas into in-line anonymous functions when possible

Previously we always generated factories for contextful lambdas: ```kt fun foo(a: Int) = { a } ``` ```js function foo(a_38) { return foo$lambda(a_38); } // factory! function foo$lambda($a) { return function () { return $a; }; } ``` After this patch, the generated code for `foo` is more concise: ```js function foo(a) { return function() { return a; }; } ```
2022-04-04 18:11:52 +03:00
parent 5b61e60f2f
commit c10af22b27
5 changed files with 350 additions and 26 deletions
@@ -10,12 +10,14 @@ import org.jetbrains.kotlin.backend.common.compilationException

 import org.jetbrains.kotlin.backend.common.lower.LoweredStatementOrigins
 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
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.JsStatementOrigins
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.backend.js.utils.Namer
+import org.jetbrains.kotlin.ir.backend.js.utils.isDispatchReceiver
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
@@ -25,17 +27,111 @@ import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.classifierOrNull
 import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.cast
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs

 class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLoweringPass {

+    /**
+     * A factory that creates [IrFunctionExpression] for lambdas being constructed.
+     *
+     * Basically, replaces a constructor call of a lambda class with a JS function expression,
+     * remapping the captured values.
+     */
+    private inner class FunctionExpressionFactory(
+        private val lambdaDeclaration: IrSimpleFunction,
+        private val constructor: IrConstructor,
+        private val lambdaInfo: LambdaInfo,
+    ) {
+        fun createFunctionExpression(ctorCall: IrConstructorCall): IrExpression {
+            val superClass = lambdaInfo.superInvokeFun.parentAsClass
+            val lambdaType = lambdaInfo.lambdaClass.superTypes.single { it.classifierOrNull === superClass.symbol }
+
+            val lambdaContextMapping = remapCapturedFields(constructor) { ctorParameter ->
+                when (val valueArgument = ctorCall.getValueArgument(ctorParameter.owner.index)!!) {
+                    is IrGetValue -> valueArgument.symbol
+                    // IrTypeOperatorCall can be passed to a lambda constructor if it was generated by the inlineLambdaBody method.
+                    is IrTypeOperatorCall -> valueArgument.argument.safeAs<IrGetValue>()?.symbol
+                    else -> null
+                }
+            }
+
+            val outerReceiverMapping = remapCapturedFields(constructor) { ctorParameter ->
+                val valueArgument = ctorCall.getValueArgument(ctorParameter.owner.index)!!
+                valueArgument.safeAs<IrGetField>()
+            }
+
+            assert(lambdaContextMapping.size + outerReceiverMapping.size == ctorCall.valueArgumentsCount)
+
+            fun Iterable<IrValueSymbol>.findDispatchReceiver() = find { it.owner.isDispatchReceiver }
+
+            val capturedDispatchReceiver = lambdaContextMapping.values.findDispatchReceiver()
+                ?: outerReceiverMapping.values.mapNotNull { it.receiver?.safeAs<IrGetValue>()?.symbol }.findDispatchReceiver()
+
+            lambdaDeclaration.body = inlineLambdaBody(
+                lambdaDeclaration,
+                lambdaInfo.invokeFun,
+                lambdaInfo.createOldToNewInvokeParametersMapping(lambdaDeclaration),
+                lambdaContextMapping,
+                outerReceiverMapping,
+            )
+            val functionExpression = IrFunctionExpressionImpl(
+                ctorCall.startOffset,
+                ctorCall.endOffset,
+                lambdaType,
+                lambdaDeclaration,
+                JsStatementOrigins.CALLABLE_REFERENCE_CREATE
+            )
+
+            // TODO: If we generate arrow functions instead of anonymous functions, there's no need for jsBind
+            // TODO: Do we need to set proper offsets?
+            if (capturedDispatchReceiver != null)
+                return IrCallImpl(
+                    UNDEFINED_OFFSET,
+                    UNDEFINED_OFFSET,
+                    context.irBuiltIns.anyType,
+                    context.intrinsics.jsBind,
+                    valueArgumentsCount = 2,
+                    typeArgumentsCount = 0,
+                    origin = JsStatementOrigins.BIND_CALL,
+                ).apply {
+                    putValueArgument(0, IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, capturedDispatchReceiver))
+                    putValueArgument(1, functionExpression)
+                }
+
+            return functionExpression
+        }
+    }
+
    override fun lower(irFile: IrFile) {
-        val ctorToFactoryMap = mutableMapOf<IrConstructorSymbol, IrSimpleFunctionSymbol>()
+
+        // Regular contextless lambdas are always transformed to function references
        val ctorToFreeFunctionMap = mutableMapOf<IrConstructorSymbol, IrSimpleFunctionSymbol>()
-        irFile.transform(CallableReferenceClassTransformer(ctorToFactoryMap, ctorToFreeFunctionMap), null)
+
+        // Regular lambdas with captured variables are transformed to function expressions whenever possible.
+        // However, we don't do that if the lambda captures a variable declared in a loop, at least when variable
+        // declarations are lowered into 'var' statements in JS. See the CapturedVariablesDeclaredInLoops class.
+        // We also don't do that if there is more than one constructor call for a single lambda.
+        val ctorToFunctionExpressionMap = mutableMapOf<IrConstructorSymbol, FunctionExpressionFactory>()
+
+        // Suspend lambdas are transformed to factory calls
+        val ctorToFactoryMap = mutableMapOf<IrConstructorSymbol, IrSimpleFunctionSymbol>()
+
+        val closureUsageAnalyser = ClosureUsageAnalyser()
+
+        irFile.acceptChildrenVoid(closureUsageAnalyser) // TODO: Only do this if lambda inlining is enabled with a feature flag!
+
+        val callableReferenceClassTransformer = CallableReferenceClassTransformer(
+            ctorToFactoryMap,
+            ctorToFreeFunctionMap,
+            ctorToFunctionExpressionMap,
+            closureUsageAnalyser
+        )
+
+        irFile.transform(callableReferenceClassTransformer, null)
+
        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
            override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
                expression.transformChildrenVoid()
@@ -45,6 +141,13 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
                    return replaceLambdaConstructorCallWithReferenceToLiftedLambda(expression, liftedLambda)
                }

+                ctorToFunctionExpressionMap[expression.symbol]?.let { functionExpressionFactory ->
+                    return functionExpressionFactory.createFunctionExpression(expression).apply {
+                        // Make sure to also apply this transformation to the inlined lambda body.
+                        transformChildrenVoid()
+                    }
+                }
+
                ctorToFactoryMap[expression.symbol]?.let { factory ->
                    return replaceLambdaConstructorCallWithFactoryCall(expression, factory)
                }
@@ -90,10 +193,86 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        compilationException("Unreachable", irBody)
    }

+    /**
+     * This class helps to determine if a lambda captures a variable declared in a loop.
+     * This is needed because variable declarations are lowered into JavaScript `var` statements, which
+     * are function scoped. Because anonymous functions in JavaScript capture variables by reference, this
+     * may lead to unintended effects like
+     * [this](https://stackoverflow.com/questions/750486/javascript-closure-inside-loops-simple-practical-example).
+     *
+     * ES6 `let` statements don't have this problem.
+     */
+    private class ClosureUsageAnalyser : IrElementVisitorVoid {
+
+        private val lambdaConstructorCalls: MutableMap<IrConstructorSymbol, MutableList<IrConstructorCall>> = hashMapOf()
+        private val variablesDeclaredInLoops: MutableSet<IrValueDeclaration> = hashSetOf()
+        private var loopNestedness = 0
+
+        override fun visitElement(element: IrElement) {
+            element.acceptChildrenVoid(this)
+        }
+
+        override fun visitLoop(loop: IrLoop) {
+            ++loopNestedness
+            loop.acceptChildrenVoid(this)
+            --loopNestedness
+        }
+
+        override fun visitVariable(declaration: IrVariable) {
+            if (loopNestedness > 0)
+                variablesDeclaredInLoops.add(declaration)
+            declaration.acceptChildrenVoid(this)
+        }
+
+        override fun visitConstructorCall(expression: IrConstructorCall) {
+            if (expression.origin == JsStatementOrigins.CALLABLE_REFERENCE_CREATE)
+                lambdaConstructorCalls
+                    .getOrPut(expression.symbol, ::mutableListOf)
+                    .add(expression)
+            expression.acceptChildrenVoid(this)
+        }
+
+        private fun IrElement.referencesVariablesDeclaredInLoops(): Boolean {
+            var result = false
+            acceptVoid(object : IrElementVisitorVoid {
+
+                override fun visitElement(element: IrElement) {
+                    if (!result)
+                        element.acceptChildrenVoid(this)
+                }
+
+                override fun visitGetValue(expression: IrGetValue) {
+                    if (expression.symbol.owner in variablesDeclaredInLoops)
+                        result = true
+                    else
+                        expression.acceptChildrenVoid(this)
+                }
+            })
+            return result
+        }
+
+        private fun IrConstructorCall.referencesVariablesDeclaredInLoops(): Boolean =
+            (0 until valueArgumentsCount).any { i ->
+                getValueArgument(i)!!.referencesVariablesDeclaredInLoops()
+            }
+
+        fun lambdaCapturesVariablesDeclaredInLoops(lambdaClass: IrClass): Boolean {
+            val primaryConstructor = lambdaClass.primaryConstructor ?: return false
+            val ctorCalls = lambdaConstructorCalls[primaryConstructor.symbol] ?: return false
+            return ctorCalls.any { it.referencesVariablesDeclaredInLoops() }
+        }
+
+        fun getLambdaConstructorCalls(constructorSymbol: IrConstructorSymbol): List<IrConstructorCall> =
+            lambdaConstructorCalls[constructorSymbol] ?: emptyList()
+    }
+
    private inner class CallableReferenceClassTransformer(
        private val ctorToFactoryMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>,
-        private val ctorToFreeFunctionMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>
+        private val ctorToFreeFunctionMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>,
+        private val ctorToFunctionExpressionMap: MutableMap<IrConstructorSymbol, FunctionExpressionFactory>,
+        private val closureUsageAnalyser: ClosureUsageAnalyser
    ) : IrElementTransformerVoid() {
+
        override fun visitFile(declaration: IrFile): IrFile {
            declaration.transformChildrenVoid()
            declaration.transformDeclarationsFlat { it.transformCallableReference() }
@@ -130,23 +309,43 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        private fun replaceWithFactory(lambdaClass: IrClass): List<IrDeclaration> {
            val lambdaInfo = LambdaInfo(lambdaClass)

-            // Optimization:
-            // If the lambda has no context, we lift it, i.e. instead of generating a factory function that creates lambda objects,
-            // we generate a named free function. The usage of the lambda is then replaced with a reference to the free function.
-            // This allows us to avoid allocating a new object each time the lambda is created.
-            return if (lambdaClass.origin == CallableReferenceLowering.Companion.LAMBDA_IMPL && !lambdaInfo.isSuspendLambda && lambdaClass.fields.none()) {
-                liftLambda(ctorToFreeFunctionMap, lambdaInfo)
+            return if (lambdaClass.origin == CallableReferenceLowering.Companion.LAMBDA_IMPL && !lambdaInfo.isSuspendLambda) {
+                if (lambdaClass.fields.none()) {
+                    // Optimization:
+                    // If the lambda has no context, we lift it, i.e. instead of generating an anonymous function,
+                    // we generate a named free function. The usage of the lambda is then replaced with a reference to the free function.
+                    // This allows us to avoid allocating a new object each time the lambda is created.
+                    liftLambda(ctorToFreeFunctionMap, lambdaInfo)
+                } else if (
+                    // If the lambda constructor is called from more than one place, don't inline.
+                    closureUsageAnalyser.getLambdaConstructorCalls(lambdaClass.primaryConstructor!!.symbol).size == 1
+                    // In-line anonymous functions that capture variables declared in loops are dangerous.
+                    // See https://stackoverflow.com/questions/750486/javascript-closure-inside-loops-simple-practical-example
+                    && !closureUsageAnalyser.lambdaCapturesVariablesDeclaredInLoops(lambdaClass)
+                ) {
+                    // If possible, generate anonymous functions in-line instead of factories of anonymous functions.
+                    buildFunctionExpression(ctorToFunctionExpressionMap, lambdaInfo)
+                } else {
+                    buildFactoryFunction(ctorToFactoryMap, lambdaInfo)
+                }
            } else {
                buildFactoryFunction(ctorToFactoryMap, lambdaInfo)
            }.onEach { it.parent = lambdaClass.parent }
        }
    }

+    /**
+     * [lambdaContextMapping] — a mapping from lambda class fields to the values outside the lambda that the lambda captures.
+     *
+     * [outerReceiverMapping] — a mapping from lambda class fields to outer class receivers, in case the lambda is inside an inner class,
+     * and it captures the outer classes.
+     */
    private fun inlineLambdaBody(
        lambdaDeclaration: IrSimpleFunction,
        invokeFun: IrSimpleFunction,
        invokeMapping: Map<IrValueParameterSymbol, IrValueParameterSymbol>,
-        factoryMapping: Map<IrFieldSymbol, IrValueParameterSymbol>
+        lambdaContextMapping: Map<IrFieldSymbol, IrValueSymbol>,
+        outerReceiverMapping: Map<IrFieldSymbol, IrGetField> = emptyMap()
    ): IrBlockBody {
        val body = invokeFun.body
            ?: compilationException(
@@ -162,8 +361,22 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        body.transformChildrenVoid(object : IrElementTransformerVoid() {
            override fun visitGetField(expression: IrGetField): IrExpression {
                expression.transformChildrenVoid()
-                val parameter = factoryMapping[expression.symbol] ?: return expression
-                return expression.getValue(parameter)
+
+                lambdaContextMapping[expression.symbol]?.let {
+                    return expression.getValue(it)
+                }
+
+                outerReceiverMapping[expression.symbol]?.let {
+                    return IrGetFieldImpl(
+                        expression.startOffset,
+                        expression.endOffset,
+                        it.symbol,
+                        it.type,
+                        it.receiver?.deepCopyWithSymbols()
+                    )
+                }
+
+                return expression
            }

            override fun visitGetValue(expression: IrGetValue): IrExpression {
@@ -227,21 +440,30 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        )
    }

-    private fun capturedFieldsToParametersMap(constructor: IrConstructor, factoryFunction: IrSimpleFunction): Map<IrFieldSymbol, IrValueParameterSymbol> {
-        val statements = constructor.body?.let { it.cast<IrBlockBody>().statements }
+    /**
+     * Returns a mapping from a lambda class field to the corresponding captured value.
+     *
+     * [remapVP] accepts a lambda constructor's value parameter symbol, for which it should return the corresponding captured value.
+     */
+    private fun <T> remapCapturedFields(
+        lambdaConstructor: IrConstructor,
+        remapVP: (IrValueParameterSymbol) -> T?
+    ): Map<IrFieldSymbol, T> {
+        val statements = lambdaConstructor.body?.let { it.cast<IrBlockBody>().statements }
            ?: compilationException(
                "Expecting Body for function ref constructor",
-                constructor
+                lambdaConstructor
            )
-
-        val fieldSetters = statements.filterIsInstance<IrSetField>()
+        return statements
+            .asSequence()
+            .filterIsInstance<IrSetField>()
            .filter { it.origin == LoweredStatementOrigins.STATEMENT_ORIGIN_INITIALIZER_OF_FIELD_FOR_CAPTURED_VALUE }
-
-        fun remapVP(vp: IrValueParameterSymbol): IrValueParameterSymbol {
-            return factoryFunction.valueParameters[vp.owner.index].symbol
-        }
-
-        return fieldSetters.associate { it.symbol to remapVP(it.value.cast<IrGetValue>().symbol.cast()) }
+            .mapNotNull { irSetField ->
+                remapVP(irSetField.value.cast<IrGetValue>().symbol.cast())?.let {
+                    irSetField.symbol to it
+                }
+            }
+            .toMap()
    }

    private fun extractReferenceReflectionName(getter: IrSimpleFunction): IrExpression {
@@ -309,7 +531,7 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo

            newDeclarations.add(lambdaInfo.lambdaClass)
        } else {
-            val fieldToParameterMapping = capturedFieldsToParametersMap(constructor, factoryFunction)
+            val fieldToParameterMapping = remapCapturedFields(constructor) { factoryFunction.valueParameters[it.owner.index].symbol }
            val oldToNewInvokeParametersMapping = lambdaInfo.createOldToNewInvokeParametersMapping(lambdaDeclaration)
            lambdaDeclaration.body =
                inlineLambdaBody(lambdaDeclaration, lambdaInfo.invokeFun, oldToNewInvokeParametersMapping, fieldToParameterMapping)
@@ -420,6 +642,30 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        return newDeclarations
    }

+    /**
+     * Builds a declaration of the function expression that the lambda will be lowered to.
+     * The declaration doesn't have a body — the body will be generated at a later stage, when we visit lambda constructor calls.
+     */
+    private fun buildFunctionExpression(
+        ctorToFunctionExpressionMap: MutableMap<IrConstructorSymbol, FunctionExpressionFactory>,
+        lambdaInfo: LambdaInfo
+    ): List<IrDeclaration> {
+        val lambdaDeclaration = createLambdaDeclaration(
+            lambdaInfo.invokeFun,
+            lambdaInfo.lambdaClass.name,
+            lambdaInfo.lambdaClass.parent,
+            lambdaInfo.superInvokeFun
+        )
+
+        val constructor = lambdaInfo.lambdaClass.constructors.single()
+
+        ctorToFunctionExpressionMap[constructor.symbol] =
+            FunctionExpressionFactory(lambdaDeclaration, constructor, lambdaInfo)
+
+        // lambdas can contain another lambdas and local classes in so let's not lose them
+        return lambdaInfo.lambdaInnerClasses()
+    }
+
    /**
     * Replaces a contextless lambda class with a free function.
     */
@@ -75,6 +75,12 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer<JsSta
        return expression.accept(IrElementToJsExpressionTransformer(), context).makeStmt()
    }

+    override fun visitFunctionExpression(expression: IrFunctionExpression, context: JsGenerationContext): JsStatement {
+        // If IrFunctionExpression is not used (i. e. the function expression is also a statement),
+        // the generated function cannot be anonymous, so we don't erase its name, unlike in IrElementToJsExpressionTransformer
+        return expression.function.accept(IrFunctionToJsTransformer(), context).makeStmt()
+    }
+
    override fun visitBreak(jump: IrBreak, context: JsGenerationContext): JsStatement {
        return JsBreak(context.getNameForLoop(jump.loop)?.let { JsNameRef(it) }).withSource(jump, context)
    }
@@ -356,6 +356,12 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
            runTest("js/js.translator/testData/box/closure/closureArrayListInstance.kt");
        }

+        @Test
+        @TestMetadata("closureCodeSize.kt")
+        public void testClosureCodeSize() throws Exception {
+            runTest("js/js.translator/testData/box/closure/closureCodeSize.kt");
+        }
+
        @Test
        @TestMetadata("closureFunctionAsArgument.kt")
        public void testClosureFunctionAsArgument() throws Exception {
@@ -356,6 +356,12 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
            runTest("js/js.translator/testData/box/closure/closureArrayListInstance.kt");
        }

+        @Test
+        @TestMetadata("closureCodeSize.kt")
+        public void testClosureCodeSize() throws Exception {
+            runTest("js/js.translator/testData/box/closure/closureCodeSize.kt");
+        }
+
        @Test
        @TestMetadata("closureFunctionAsArgument.kt")
        public void testClosureFunctionAsArgument() throws Exception {
@@ -0,0 +1,60 @@
+// KT-51133
+
+package foo
+
+fun <T, R> myWith(t: T, f: (T) -> R) = f(t)
+
+// CHECK_CONTAINS_NO_CALLS: noCapture except=myWith
+// HAS_NO_CAPTURED_VARS: function=noCapture except=myWith;noCapture$lambda
+fun noCapture() = myWith(42) { it }
+
+// CHECK_CONTAINS_NO_CALLS: captureLocalVariableReadOnly except=myWith IGNORED_BACKENDS=JS
+// HAS_NO_CAPTURED_VARS: function=captureLocalVariableReadOnly except=myWith IGNORED_BACKENDS=JS
+fun captureLocalVariableReadOnly(a: Int) = myWith(1) { a + it }
+
+// CHECK_CONTAINS_NO_CALLS: captureLocalVariableReadWrite except=myWith IGNORED_BACKENDS=JS
+// HAS_NO_CAPTURED_VARS: function=captureLocalVariableReadWrite except=myWith IGNORED_BACKENDS=JS
+fun captureLocalVariableReadWrite(): Int {
+    var a = 41
+    return myWith(1) {
+        a += it
+        a
+    }
+}
+
+val thirteen = 13
+
+// CHECK_CONTAINS_NO_CALLS: captureGlobalVariable except=myWith
+// HAS_NO_CAPTURED_VARS: function=captureGlobalVariable except=myWith;captureGlobalVariable$lambda
+fun captureGlobalVariable() = myWith(2) { thirteen * it }
+
+class A(val i: Int) {
+
+    fun captureClassField() = myWith(100) { it + i }
+
+    inner class B(val j: Int) {
+
+        inner class C(val k: Int) {
+            fun captureInnerClassField() = myWith(200) { it + i + j + k }
+
+            fun unusedLambda() {
+                { i + j + k }
+            }
+        }
+    }
+}
+
+fun unusedLambda(f: () -> Unit) {
+    { f() }
+}
+
+fun box(): String {
+    if (noCapture() != 42) return "fail noCapture()"
+    if (captureLocalVariableReadOnly(41) != 42) return "fail captureLocalVariableReadOnly(41)"
+    if (captureLocalVariableReadWrite() != 42) return "fail captureLocalVariableReadWrite()"
+    if (captureGlobalVariable() != 26) return "fail captureGlobalVariable()"
+    if (A(23).captureClassField() != 123) return "fail A(23).captureClassField()"
+    if (A(300).B(400).C(500).captureInnerClassField() != 1400) return "A(300).B(400).C(500).captureInnerClassField()"
+
+    return "OK"
+}