diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt index dce0a538485..92e884a8614 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt @@ -87,11 +87,18 @@ class LocalDeclarationsLowering( IrDeclarationOriginImpl("FIELD_FOR_CROSSINLINE_CAPTURED_VALUE", isSynthetic = true) override fun lower(irBody: IrBody, container: IrDeclaration) { - LocalDeclarationsTransformer(irBody, container, null).lowerLocalDeclarations() + LocalDeclarationsTransformer(irBody, container).lowerLocalDeclarations() } fun lower(irElement: IrElement, container: IrDeclaration, classesToLower: Set) { - LocalDeclarationsTransformer(irElement, container, classesToLower).lowerLocalDeclarations() + LocalDeclarationsTransformer(irElement, container, null, classesToLower).lowerLocalDeclarations() + } + + fun lower( + irBlock: IrBlock, container: IrDeclaration, closestParent: IrDeclarationParent, + classesToLower: Set, functionsToSkip: Set + ) { + LocalDeclarationsTransformer(irBlock, container, closestParent, classesToLower, functionsToSkip).lowerLocalDeclarations() } internal class ScopeWithCounter(val irElement: IrElement) { @@ -136,10 +143,33 @@ class LocalDeclarationsLowering( } } + sealed class OwnerForLoweredDeclaration(val isLocal: Boolean) { + abstract fun addChild(declaration: IrDeclaration) + abstract fun closestDeclarationParent(): IrDeclarationParent + + // Usually, just move local functions to the nearest class or file. + class DeclarationContainer(private val irDeclarationContainer: IrDeclarationContainer) : OwnerForLoweredDeclaration(false) { + override fun addChild(declaration: IrDeclaration) = irDeclarationContainer.addChild(declaration) + + override fun closestDeclarationParent() = irDeclarationContainer + } + + // But, local functions defined in an inline lambda need to be popped up to the root inline call. + class Block(private val irBlock: IrBlock, private val irDeclarationParent: IrDeclarationParent) : OwnerForLoweredDeclaration(true) { + private val initialStatementsCount = irBlock.statements.size + override fun addChild(declaration: IrDeclaration) { + // Place all children at the block's start but in order they are being added. + irBlock.statements.add(irBlock.statements.size - initialStatementsCount, declaration) + } + + override fun closestDeclarationParent() = irDeclarationParent + } + } + class LocalFunctionContext( override val declaration: IrSimpleFunction, val index: Int, - val ownerForLoweredDeclaration: IrDeclarationContainer + val ownerForLoweredDeclaration: OwnerForLoweredDeclaration ) : LocalContextWithClosureAsParameters() { lateinit var closure: Closure @@ -217,7 +247,8 @@ class LocalDeclarationsLowering( } private inner class LocalDeclarationsTransformer( - val irElement: IrElement, val container: IrDeclaration, val classesToLower: Set? + val irElement: IrElement, val container: IrDeclaration, val closestParent: IrDeclarationParent? = null, + val classesToLower: Set? = null, val functionsToSkip: Set? = null ) { val localFunctions: MutableMap = LinkedHashMap() val localClasses: MutableMap = LinkedHashMap() @@ -618,8 +649,9 @@ class LocalDeclarationsLowering( throw AssertionError("local functions must not have dispatch receiver") } - val memberOwner = localFunctionContext.ownerForLoweredDeclaration - val newName = generateNameForLiftedDeclaration(oldDeclaration, memberOwner) + val owner = localFunctionContext.ownerForLoweredDeclaration + val ownerParent = owner.closestDeclarationParent() + val newName = generateNameForLiftedDeclaration(oldDeclaration, ownerParent) // TODO: consider using fields to access the closure of enclosing class. val (capturedValues, capturedTypeParameters) = localFunctionContext.closure @@ -627,7 +659,7 @@ class LocalDeclarationsLowering( val newDeclaration = context.irFactory.buildFun { updateFrom(oldDeclaration) name = newName - visibility = DescriptorVisibilities.PRIVATE + visibility = if (owner.isLocal) DescriptorVisibilities.LOCAL else DescriptorVisibilities.PRIVATE modality = Modality.FINAL } @@ -646,7 +678,7 @@ class LocalDeclarationsLowering( tp.superTypes = tp.superTypes.map { localFunctionContext.remapType(it) } } - newDeclaration.parent = memberOwner + newDeclaration.parent = ownerParent newDeclaration.returnType = localFunctionContext.remapType(oldDeclaration.returnType) newDeclaration.dispatchReceiverParameter = null newDeclaration.extensionReceiverParameter = oldDeclaration.extensionReceiverParameter?.run { @@ -917,6 +949,7 @@ class LocalDeclarationsLowering( } override fun visitSimpleFunction(declaration: IrSimpleFunction, data: Data) { + if (functionsToSkip?.contains(declaration) == true) return super.visitSimpleFunction(declaration, data.withInline(declaration.isInline)) if (declaration.visibility == DescriptorVisibilities.LOCAL) { @@ -928,8 +961,12 @@ class LocalDeclarationsLowering( if (declaration.name.isSpecial || declaration.name in enclosingScope.usedLocalFunctionNames) enclosingScope.counter++ else -1 - localFunctions[declaration] = - LocalFunctionContext(declaration, index, enclosingScope.irElement as IrDeclarationContainer) + val ownerForLoweredDeclaration = + data.currentClass?.let { OwnerForLoweredDeclaration.DeclarationContainer(it.irElement as IrDeclarationContainer) } + ?: (irElement as? IrBlock)?.let { OwnerForLoweredDeclaration.Block(irElement, closestParent!!) } + ?: OwnerForLoweredDeclaration.DeclarationContainer(enclosingScope.irElement as IrDeclarationContainer) + localFunctions[declaration] = LocalFunctionContext(declaration, index, ownerForLoweredDeclaration) + enclosingScope.usedLocalFunctionNames.add(declaration.name) } } diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt index 8ad536e1ea4..664577b0dbf 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt @@ -37,6 +37,9 @@ import org.jetbrains.kotlin.util.OperatorNameConventions fun IrValueParameter.isInlineParameter(type: IrType = this.type) = index >= 0 && !isNoinline && !type.isNullable() && (type.isFunction() || type.isSuspendFunction()) +fun IrExpression.isAdaptedFunctionReference() = + this is IrBlock && this.origin == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE + interface InlineFunctionResolver { fun getFunctionDeclaration(symbol: IrFunctionSymbol): IrFunction } @@ -415,9 +418,6 @@ class FunctionInlining( && irCall.dispatchReceiver is IrGetValue } - private fun IrExpression.isAdaptedFunctionReference() = - this is IrBlock && this.origin == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE - private inner class ParameterToArgument( val parameter: IrValueParameter, val argumentExpression: IrExpression diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt index 1c18de886a1..ceb797ffd49 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt @@ -11,10 +11,7 @@ import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering import org.jetbrains.kotlin.ir.IrElement import org.jetbrains.kotlin.ir.IrStatement 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.IrExpression -import org.jetbrains.kotlin.ir.expressions.IrGetValue +import org.jetbrains.kotlin.ir.expressions.* import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl import org.jetbrains.kotlin.ir.util.setDeclarationsParent @@ -38,25 +35,74 @@ class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : B } override fun lower(irBody: IrBody, container: IrDeclaration) { - irBody.transformChildrenVoid(object : IrElementTransformerVoidWithContext() { - override fun visitCall(expression: IrCall): IrExpression { - if (!expression.symbol.owner.isInline) - return super.visitCall(expression) + irBody.transformChildren(object : IrElementTransformer { + override fun visitDeclaration(declaration: IrDeclarationBase, data: IrDeclarationParent) = + super.visitDeclaration(declaration, (declaration as? IrDeclarationParent) ?: data) + override fun visitCall(expression: IrCall, data: IrDeclarationParent): IrElement { + val rootCallee = expression.symbol.owner + if (!rootCallee.isInline) + return super.visitCall(expression, data) + + val inlineLambdas = (0 until expression.valueArgumentsCount) + .mapNotNull { index -> + (expression.getValueArgument(index) as? IrFunctionExpression)?.function + ?.takeIf { rootCallee.valueParameters[index].isInlineParameter() } + } val localClasses = mutableSetOf() - expression.acceptChildrenVoid(object : IrElementVisitorVoid { - override fun visitElement(element: IrElement) { - element.acceptChildrenVoid(this) - } + val localFunctions = mutableSetOf() + val adaptedFunctions = mutableSetOf() + val transformer = this + for (lambda in inlineLambdas) { + lambda.acceptChildrenVoid(object : IrElementVisitorVoid { + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } - override fun visitClass(declaration: IrClass) { - localClasses.add(declaration) - } - }) - if (localClasses.isEmpty()) + override fun visitClass(declaration: IrClass) { + localClasses.add(declaration) + } + + override fun visitFunctionExpression(expression: IrFunctionExpression) { + expression.function.acceptChildrenVoid(this) + } + + override fun visitFunction(declaration: IrFunction) { + declaration.transformChildren(transformer, declaration) + + localFunctions.add(declaration) + } + + override fun visitCall(expression: IrCall) { + val callee = expression.symbol.owner + if (!callee.isInline) { + expression.acceptChildrenVoid(this) + return + } + + expression.extensionReceiver?.acceptVoid(this) + expression.dispatchReceiver?.acceptVoid(this) + (0 until expression.valueArgumentsCount).forEach { index -> + val argument = expression.getValueArgument(index) + val parameter = callee.valueParameters[index] + // Skip adapted function references - they will be inlined later. + if (parameter.isInlineParameter() && argument?.isAdaptedFunctionReference() == true) + adaptedFunctions += (argument as IrBlock).statements[0] as IrSimpleFunction + else + argument?.acceptVoid(this) + } + } + }) + } + + if (localClasses.isEmpty() && localFunctions.isEmpty()) return expression - LocalDeclarationsLowering(context).lower(expression, container, localClasses) + val irBlock = IrBlockImpl(expression.startOffset, expression.endOffset, expression.type).apply { + statements += expression + } + LocalDeclarationsLowering(context).lower(irBlock, container, data, localClasses, adaptedFunctions) + irBlock.statements.addAll(0, localClasses) expression.transformChildrenVoid(object : IrElementTransformerVoid() { override fun visitClass(declaration: IrClass): IrStatement { @@ -66,19 +112,11 @@ class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : B ) } }) - localClasses.forEach { - it.setDeclarationsParent( - currentDeclarationParent - ?: (container as? IrDeclarationParent) - ?: container.parent - ) - } - return IrBlockImpl(expression.startOffset, expression.endOffset, expression.type).apply { - statements += localClasses - statements += expression - } + localClasses.forEach { it.setDeclarationsParent(data) } + + return irBlock } - }) + }, container as? IrDeclarationParent ?: container.parent) } }