[IR][inliner] Reworked local classes in inline lambdas

#KT-52795 Fixed
This commit is contained in:
Igor Chevdar
2022-07-19 15:05:39 +03:00
committed by Space
parent 8981eb85de
commit 92109a6bef
3 changed files with 118 additions and 43 deletions
@@ -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<IrClass>) {
LocalDeclarationsTransformer(irElement, container, classesToLower).lowerLocalDeclarations()
LocalDeclarationsTransformer(irElement, container, null, classesToLower).lowerLocalDeclarations()
}
fun lower(
irBlock: IrBlock, container: IrDeclaration, closestParent: IrDeclarationParent,
classesToLower: Set<IrClass>, functionsToSkip: Set<IrSimpleFunction>
) {
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<IrClass>?
val irElement: IrElement, val container: IrDeclaration, val closestParent: IrDeclarationParent? = null,
val classesToLower: Set<IrClass>? = null, val functionsToSkip: Set<IrSimpleFunction>? = null
) {
val localFunctions: MutableMap<IrFunction, LocalFunctionContext> = LinkedHashMap()
val localClasses: MutableMap<IrClass, LocalClassContext> = 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)
}
}
@@ -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
@@ -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<IrDeclarationParent> {
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<IrClass>()
expression.acceptChildrenVoid(object : IrElementVisitorVoid {
override fun visitElement(element: IrElement) {
element.acceptChildrenVoid(this)
}
val localFunctions = mutableSetOf<IrFunction>()
val adaptedFunctions = mutableSetOf<IrSimpleFunction>()
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)
}
}