diff --git a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/InteropCallableReferenceLowering.kt b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/InteropCallableReferenceLowering.kt index 1453f739067..68b49692db6 100644 --- a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/InteropCallableReferenceLowering.kt +++ b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/InteropCallableReferenceLowering.kt @@ -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()?.symbol + else -> null + } + } + + val outerReceiverMapping = remapCapturedFields(constructor) { ctorParameter -> + val valueArgument = ctorCall.getValueArgument(ctorParameter.owner.index)!! + valueArgument.safeAs() + } + + assert(lambdaContextMapping.size + outerReceiverMapping.size == ctorCall.valueArgumentsCount) + + fun Iterable.findDispatchReceiver() = find { it.owner.isDispatchReceiver } + + val capturedDispatchReceiver = lambdaContextMapping.values.findDispatchReceiver() + ?: outerReceiverMapping.values.mapNotNull { it.receiver?.safeAs()?.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() + + // Regular contextless lambdas are always transformed to function references val ctorToFreeFunctionMap = mutableMapOf() - 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() + + // Suspend lambdas are transformed to factory calls + val ctorToFactoryMap = mutableMapOf() + + 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> = hashMapOf() + private val variablesDeclaredInLoops: MutableSet = 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 = + lambdaConstructorCalls[constructorSymbol] ?: emptyList() + } + private inner class CallableReferenceClassTransformer( private val ctorToFactoryMap: MutableMap, - private val ctorToFreeFunctionMap: MutableMap + private val ctorToFreeFunctionMap: MutableMap, + private val ctorToFunctionExpressionMap: MutableMap, + 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 { 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, - factoryMapping: Map + lambdaContextMapping: Map, + outerReceiverMapping: Map = 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 { - val statements = constructor.body?.let { it.cast().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 remapCapturedFields( + lambdaConstructor: IrConstructor, + remapVP: (IrValueParameterSymbol) -> T? + ): Map { + val statements = lambdaConstructor.body?.let { it.cast().statements } ?: compilationException( "Expecting Body for function ref constructor", - constructor + lambdaConstructor ) - - val fieldSetters = statements.filterIsInstance() + return statements + .asSequence() + .filterIsInstance() .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().symbol.cast()) } + .mapNotNull { irSetField -> + remapVP(irSetField.value.cast().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, + lambdaInfo: LambdaInfo + ): List { + 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. */ diff --git a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt index 9153c51316b..0a72714a1dc 100644 --- a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt +++ b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt @@ -75,6 +75,12 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer 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" +}