[IR] Refactor optimization of eliminating removing extra boxes

Signed-off-by: Evgeniy.Zhelenskiy <Evgeniy.Zhelenskiy@jetbrains.com>

#KT-1179

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/BridgeLowering.kt

@@ -697,7 +697,7 @@ internal class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass
                             irCastIfNeeded(irGet(bridgeParameter), targetParameterType),
                             getOptimizedPublicAccess(target, targetRemappedParameter.rootMfvcNode.mfvc)
                         ) { error("Not applicable") }
-                        val newArguments = instance.makeFlattenedGetterExpressions(this)
+                        val newArguments = instance.makeFlattenedGetterExpressions(this, registerPossibleExtraBoxCreation = {})
                         for (newArgument in newArguments) {
                             irCall.putArgument(targetExplicitParameters[targetIndex++], newArgument)
                         }
@@ -712,7 +712,7 @@ internal class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass
                                 irGet(bridgeExplicitParameters[bridgeIndex++])
                             }
                             val boxCall = bridgeRemappedParameter.rootMfvcNode.makeBoxedExpression(
-                                this, bridgeRemappedParameter.typeArguments, valueArguments
+                                this, bridgeRemappedParameter.typeArguments, valueArguments, registerPossibleExtraBoxCreation = {}
                             )
                             irCall.putArgument(targetParameter, irCastIfNeeded(boxCall, targetParameter.type.upperBound))
                         }

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/InheritedDefaultMethodsOnClassesLowering.kt

@@ -140,7 +140,7 @@ private class InheritedDefaultMethodsOnClassesLowering(val context: JvmBackendCo
                                             irGet(sourceFullValueParameterList[flattenedIndex++])
                                         }
                                         val boxedExpression = remappedParameter.rootMfvcNode.makeBoxedExpression(
-                                            this@irBlockBody, remappedParameter.typeArguments, valueArguments
+                                            this@irBlockBody, remappedParameter.typeArguments, valueArguments, registerPossibleExtraBoxCreation = {}
                                         )
                                         putValueArgument(i, boxedExpression)
                                     }

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/JvmMultiFieldValueClassLowering.kt

@@ -61,7 +61,14 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
         }
     }
 
-    private val boxUsageGenerated = mutableSetOf<IrDeclaration>()
+    private val possibleExtraBoxUsageGenerated = mutableSetOf<IrDeclaration>()
+
+    private val irCurrentScope
+        get() = currentScope!!.irElement as IrDeclaration
+
+    private fun registerPossibleExtraBoxUsage() {
+        possibleExtraBoxUsageGenerated.add(irCurrentScope)
+    }
 
     /**
      * The class is used to get replacing expression and MFVC instance if present for the given old value declaration.
@@ -89,15 +96,11 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
         fun IrBuilderWithScope.makeReplacement(expression: IrGetValue): IrExpression? {
             oldValueSymbol2NewValueSymbol[expression.symbol]?.let { return irGet(it.owner) }
             val instance = oldSymbol2MfvcNodeInstance[expression.symbol] ?: return null
-            val res = instance.makeGetterExpression(this)
+            val res = instance.makeGetterExpression(this, ::registerPossibleExtraBoxUsage)
-            boxUsageGenerated.add(irCurrentScope)
             expression2MfvcNodeInstanceAccessor[res] = MfvcNodeInstanceAccessor.Getter(instance)
             return res
         }
 
-        private val irCurrentScope
-            get() = currentScope!!.irElement as IrDeclaration
-
         private fun splitExpressions(expressions: List<IrExpression>): Pair<List<IrExpression>, List<IrExpression>> {
             val repeatable = expressions.takeLastWhile { it.isRepeatableGetter() }
             return expressions.subList(0, expressions.size - repeatable.size) to repeatable
@@ -141,8 +144,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
             val typeArguments = makeTypeArgumentsFromField(expression)
             val instance: ReceiverBasedMfvcNodeInstance =
                 node.createInstanceFromBox(this, typeArguments, expression.receiver, AccessType.AlwaysPrivate, ::variablesSaver)
-            val getterExpression = instance.makeGetterExpression(this)
+            val getterExpression = instance.makeGetterExpression(this, ::registerPossibleExtraBoxUsage)
-            boxUsageGenerated.add(irCurrentScope)
             expression2MfvcNodeInstanceAccessor[getterExpression] = MfvcNodeInstanceAccessor.Getter(instance)
             +getterExpression
             return getterExpression
@@ -177,8 +179,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
             }
             val instance: ReceiverBasedMfvcNodeInstance =
                 node.createInstanceFromBox(this, typeArguments, dispatchReceiver, accessType, ::variablesSaver)
-            val getterExpression = instance.makeGetterExpression(this)
+            val getterExpression = instance.makeGetterExpression(this, ::registerPossibleExtraBoxUsage)
-            boxUsageGenerated.add(irCurrentScope)
             expression2MfvcNodeInstanceAccessor[getterExpression] = MfvcNodeInstanceAccessor.Getter(instance)
             +getterExpression
             return getterExpression
@@ -227,9 +228,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
             val newAccessor = accessor[name] ?: return@handleSavedExpression null
             val expression = when (newAccessor) {
                 is MfvcNodeInstanceAccessor.Setter -> newAccessor.instance.makeSetterExpressions(scope, newAccessor.values)
-                is MfvcNodeInstanceAccessor.Getter -> newAccessor.instance.makeGetterExpression(scope).also {
+                is MfvcNodeInstanceAccessor.Getter -> newAccessor.instance.makeGetterExpression(scope, ::registerPossibleExtraBoxUsage)
-                    boxUsageGenerated.add(irCurrentScope)
-                }
             }
             expression2MfvcNodeInstanceAccessor[expression] = newAccessor
             expression
@@ -239,8 +238,9 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
             scope: IrBuilderWithScope,
             expression: IrExpression,
             handler: IrBlockBuilder.(values: List<IrExpression>) -> IrExpression
-        ): IrExpression? =
+        ): IrExpression? = scope.handleSavedExpression(expression) {
-            scope.handleSavedExpression(expression) { irBlock { +handler(it.instance.makeFlattenedGetterExpressions(this)) }.unwrapBlock() }
+            irBlock { +handler(it.instance.makeFlattenedGetterExpressions(this, ::registerPossibleExtraBoxUsage)) }.unwrapBlock()
+        }
 
         fun registerReplacement(expression: IrExpression, instance: MfvcNodeInstance) {
             expression2MfvcNodeInstanceAccessor[expression] = MfvcNodeInstanceAccessor.Getter(instance)
@@ -291,7 +291,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                         is IrFunction -> replacingDeclaration.body = replacingDeclaration.body?.makeBodyWithAddedVariables(
                             context, variablesToAdd[replacingDeclaration] ?: emptySet(), replacingDeclaration.symbol
                         )?.apply {
-                            if (replacingDeclaration in boxUsageGenerated) {
+                            if (replacingDeclaration in possibleExtraBoxUsageGenerated) {
                                 removeAllExtraBoxes()
                             }
                         }
@@ -299,7 +299,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                         is IrAnonymousInitializer -> replacingDeclaration.body = replacingDeclaration.body.makeBodyWithAddedVariables(
                             context, variablesToAdd[replacingDeclaration.parent] ?: emptySet(), replacingDeclaration.symbol
                         ).apply {
-                            if (replacingDeclaration in boxUsageGenerated) {
+                            if (replacingDeclaration in possibleExtraBoxUsageGenerated) {
                                 removeAllExtraBoxes()
                             }
                         } as IrBlockBody
@@ -554,7 +554,10 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                                             .map { irGet(it) }
                                         val targetParameter = targetExplicitParameters[flattenedTargetIndex++]
                                         val boxedExpression = remappedSourceParameter.rootMfvcNode.makeBoxedExpression(
-                                            this@irBlock, remappedSourceParameter.typeArguments, valueArguments
+                                            this@irBlock,
+                                            remappedSourceParameter.typeArguments,
+                                            valueArguments,
+                                            ::registerPossibleExtraBoxUsage
                                         )
                                         putArgument(targetParameter, boxedExpression)
                                             .also { flattenedSourceIndex += remappedSourceParameter.valueParameters.size }
@@ -569,7 +572,8 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                                     val instance = rootNode.createInstanceFromBox(
                                         this@irBlock, irGet(receiver), getOptimizedPublicAccess(rootNode.mfvc), ::variablesSaver,
                                     )
-                                    val flattenedExpressions = instance.makeFlattenedGetterExpressions(this@irBlock)
+                                    val flattenedExpressions =
+                                        instance.makeFlattenedGetterExpressions(this@irBlock, ::registerPossibleExtraBoxUsage)
                                     for (expression in flattenedExpressions) {
                                         putArgument(targetExplicitParameters[flattenedTargetIndex++], expression)
                                     }
@@ -747,8 +751,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                         this, function.constructedClassType as IrSimpleType, Name.identifier("constructor_tmp"), ::variablesSaver
                     )
                     flattenExpressionTo(expression, instance)
-                    val getterExpression = instance.makeGetterExpression(this)
+                    val getterExpression = instance.makeGetterExpression(this, ::registerPossibleExtraBoxUsage)
-                    boxUsageGenerated.add(currentScope)
                     valueDeclarationsRemapper.registerReplacement(getterExpression, instance)
                     +getterExpression
                 }
@@ -1117,7 +1120,9 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                 }
                 +irCall(rootNode.primaryConstructorImpl).apply {
                     copyTypeArgumentsFrom(expression)
-                    for ((index, leafExpression) in instance.makeFlattenedGetterExpressions(this@flattenExpressionTo).withIndex()) {
+                    val flattenedGetterExpressions =
+                        instance.makeFlattenedGetterExpressions(this@flattenExpressionTo, ::registerPossibleExtraBoxUsage)
+                    for ((index, leafExpression) in flattenedGetterExpressions.withIndex()) {
                         putValueArgument(index, leafExpression)
                     }
                 }
@@ -1137,7 +1142,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
             this, transformedExpression, getOptimizedPublicAccess(rootNode.mfvc), ::variablesSaver,
         )
         require(expressionInstance.size == instance.size) { "Incompatible assignment sizes: ${expressionInstance.size}, ${instance.size}" }
-        instance.addSetterStatements(this, expressionInstance.makeFlattenedGetterExpressions(this))
+        instance.addSetterStatements(this, expressionInstance.makeFlattenedGetterExpressions(this, ::registerPossibleExtraBoxUsage))
     }
 
     private fun getOptimizedPublicAccess(parent: IrClass): AccessType =

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/MfvcNode.kt

@@ -142,7 +142,10 @@ sealed interface MfvcNodeWithSubnodes : MfvcNode {
 }
 
 fun MfvcNodeWithSubnodes.makeBoxedExpression(
-    scope: IrBuilderWithScope, typeArguments: TypeArguments, valueArguments: List<IrExpression>
+    scope: IrBuilderWithScope,
+    typeArguments: TypeArguments,
+    valueArguments: List<IrExpression>,
+    registerPossibleExtraBoxCreation: () -> Unit,
 ): IrExpression = scope.irCall(boxMethod).apply {
     val resultType = type.substitute(typeArguments) as IrSimpleType
     require(resultType.erasedUpperBound == type.erasedUpperBound) { "Substitution of $type led to $resultType" }
@@ -152,6 +155,7 @@ fun MfvcNodeWithSubnodes.makeBoxedExpression(
     for ((index, valueArgument) in valueArguments.withIndex()) {
         putValueArgument(index, valueArgument)
     }
+    registerPossibleExtraBoxCreation()
 }
 
 operator fun MfvcNodeWithSubnodes.get(names: List<Name>): MfvcNode? {

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/MfvcNodeFactory.kt

@@ -234,7 +234,7 @@ fun createIntermediateMfvcNode(
     ) { receiver ->
         val valueArguments = subnodes.flatMap { it.fields!! }
             .map { field -> irGetField(if (field.isStatic) null else irGet(receiver!!), field) }
-        rootNode.makeBoxedExpression(this, typeArguments, valueArguments)
+        rootNode.makeBoxedExpression(this, typeArguments, valueArguments, registerPossibleExtraBoxCreation = {})
     }
 
     val hasPureUnboxMethod = defaultMethodsImplementationSourceNode.isPure() && subnodes.all { it.hasPureUnboxMethod }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/MfvcNodeInstance.kt

@@ -23,8 +23,8 @@ interface MfvcNodeInstance {
     val typeArguments: TypeArguments
     val type: IrSimpleType
 
-    fun makeFlattenedGetterExpressions(scope: IrBlockBuilder): List<IrExpression>
+    fun makeFlattenedGetterExpressions(scope: IrBlockBuilder, registerPossibleExtraBoxCreation: () -> Unit): List<IrExpression>
-    fun makeGetterExpression(scope: IrBuilderWithScope): IrExpression
+    fun makeGetterExpression(scope: IrBuilderWithScope, registerPossibleExtraBoxCreation: () -> Unit): IrExpression
     operator fun get(name: Name): MfvcNodeInstance?
     fun makeStatements(scope: IrBuilderWithScope, values: List<IrExpression>): List<IrStatement>
 }
@@ -57,13 +57,16 @@ class ValueDeclarationMfvcNodeInstance(
 
     override val type: IrSimpleType = makeTypeFromMfvcNodeAndTypeArguments(node, typeArguments)
 
-    override fun makeFlattenedGetterExpressions(scope: IrBlockBuilder): List<IrExpression> =
+    override fun makeFlattenedGetterExpressions(scope: IrBlockBuilder, registerPossibleExtraBoxCreation: () -> Unit): List<IrExpression> =
         makeFlattenedGetterExpressions(scope as IrBuilderWithScope)
+
     private fun makeFlattenedGetterExpressions(scope: IrBuilderWithScope): List<IrExpression> = valueDeclarations.map { scope.irGet(it) }
 
-    override fun makeGetterExpression(scope: IrBuilderWithScope): IrExpression = when (node) {
+    override fun makeGetterExpression(scope: IrBuilderWithScope, registerPossibleExtraBoxCreation: () -> Unit): IrExpression = when (node) {
         is LeafMfvcNode -> makeFlattenedGetterExpressions(scope).single()
-        is MfvcNodeWithSubnodes -> node.makeBoxedExpression(scope, typeArguments, makeFlattenedGetterExpressions(scope))
+        is MfvcNodeWithSubnodes -> node.makeBoxedExpression(
+            scope, typeArguments, makeFlattenedGetterExpressions(scope), registerPossibleExtraBoxCreation
+        )
     }
 
     override fun get(name: Name): ValueDeclarationMfvcNodeInstance? {
@@ -158,19 +161,21 @@ class ReceiverBasedMfvcNodeInstance(
         require(node is RootMfvcNode == (unboxMethod == null)) { "Only root node has node getter" }
     }
 
-    override fun makeFlattenedGetterExpressions(scope: IrBlockBuilder): List<IrExpression> = when (node) {
+    override fun makeFlattenedGetterExpressions(
-        is LeafMfvcNode -> listOf(makeGetterExpression(scope))
+        scope: IrBlockBuilder, registerPossibleExtraBoxCreation: () -> Unit
+    ): List<IrExpression> = when (node) {
+        is LeafMfvcNode -> listOf(makeGetterExpression(scope, registerPossibleExtraBoxCreation))
         is MfvcNodeWithSubnodes -> when {
             node is IntermediateMfvcNode && canUsePrivateAccessFor(node) && fields != null ->
                 fields.map { scope.irGetField(makeReceiverCopy(), it) }
 
             node !is IntermediateMfvcNode || node.hasPureUnboxMethod && (canUsePrivateAccessFor(node) || node.unboxMethod.parentAsClass.isMultiFieldValueClass) ->
-                // We cannot relay on purity of non-mfvc class getter because the incremental compiler will not recompile the code
+                // We cannot rely on purity of non-mfvc class getter because the incremental compiler will not recompile the code
                 // that relayed after the change of purity if it is not local for the class.
-                node.subnodes.flatMap { get(it.name)!!.makeFlattenedGetterExpressions(scope) }
+                node.subnodes.flatMap { get(it.name)!!.makeFlattenedGetterExpressions(scope, registerPossibleExtraBoxCreation) }
 
             else -> {
-                val value = makeGetterExpression(scope)
+                val value = makeGetterExpression(scope, registerPossibleExtraBoxCreation = { /* The box is definitely useful */ })
                 val asVariable = scope.savableStandaloneVariableWithSetter(
                     value,
                     origin = IrDeclarationOrigin.GENERATED_MULTI_FIELD_VALUE_CLASS_PARAMETER,
@@ -180,16 +185,17 @@ class ReceiverBasedMfvcNodeInstance(
                 val root = node.rootNode
                 val variableInstance =
                     root.createInstanceFromBox(scope, typeArguments, scope.irGet(asVariable), accessType, saveVariable)
-                variableInstance.makeFlattenedGetterExpressions(scope)
+                variableInstance.makeFlattenedGetterExpressions(scope, registerPossibleExtraBoxCreation)
             }
         }
     }
 
-    override fun makeGetterExpression(scope: IrBuilderWithScope): IrExpression = with(scope) {
+    override fun makeGetterExpression(scope: IrBuilderWithScope, registerPossibleExtraBoxCreation: () -> Unit): IrExpression = with(scope) {
         when {
             node is LeafMfvcNode && canUsePrivateAccessFor(node) && fields != null -> irGetField(makeReceiverCopy(), fields.single())
-            node is IntermediateMfvcNode && accessType == AccessType.AlwaysPrivate && fields != null ->
+            node is IntermediateMfvcNode && accessType == AccessType.AlwaysPrivate && fields != null -> node.makeBoxedExpression(
-                node.makeBoxedExpression(this, typeArguments, fields.map { irGetField(makeReceiverCopy(), it) })
+                this, typeArguments, fields.map { irGetField(makeReceiverCopy(), it) }, registerPossibleExtraBoxCreation
+            )
 
             unboxMethod != null -> irCall(unboxMethod).apply {
                 val dispatchReceiverParameter = unboxMethod.dispatchReceiverParameter