[IR] Move JvmUpgradeCallableReferences before JvmMultiFieldValueClassLowering

#KT-74383

Author: zhelenskiy

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

@@ -65,8 +65,8 @@ private val jvmFilePhases = createFilePhases(
     ::ForLoopsLowering,
     ::CollectionStubMethodLowering,
     ::JvmSingleAbstractMethodLowering,
-    ::JvmMultiFieldValueClassLowering,
     ::JvmUpgradeCallableReferences,
+    ::JvmMultiFieldValueClassLowering,
     ::JvmInlineClassLowering,
     ::JvmTailrecLowering,
 

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

@@ -760,41 +760,43 @@ internal class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : Jvm
                 "Expected ${replacement.render()}, got ${declarations?.map { it.render() }}"
             }
         }
-        return makeNewLambda(originalFunction, ref, makeBody = { wrapper ->
-            variablesToAdd[replacement]?.let {
-                variablesToAdd[wrapper] = it
-                variablesToAdd.remove(replacement)
-            }
-            if (replacement in possibleExtraBoxUsageGenerated) {
-                possibleExtraBoxUsageGenerated.add(wrapper)
-                possibleExtraBoxUsageGenerated.remove(replacement)
-            }
-            with(context.createJvmIrBuilder(wrapper.symbol)) {
-                irExprBody(irBlock {
-                    val newArguments: List<IrValueDeclaration> = wrapper.parameters.flatMap { parameter ->
-                        if (!parameter.type.needsMfvcFlattening()) {
-                            listOf(parameter)
-                        } else {
-                            // Old parameter value will be only used to set parameter of the lowered function,
-                            // thus it is useless to show it in debugger
-                            parameter.origin = JvmLoweredDeclarationOrigin.TEMPORARY_MULTI_FIELD_VALUE_CLASS_PARAMETER
-                            val rootNode = replacements.getRootMfvcNode(parameter.type.erasedUpperBound)
-                            rootNode.createInstanceFromBox(this, irGet(parameter), AccessType.ChooseEffective, ::variablesSaver)
-                                .makeFlattenedGetterExpressions(this, irCurrentClass, ::registerPossibleExtraBoxUsage)
-                                .mapIndexed { index, expression ->
-                                    savableStandaloneVariableWithSetter(
-                                        expression = expression,
-                                        name = "${parameter.name.asString()}-${rootNode.leaves[index].fullFieldName}",
-                                        origin = JvmLoweredDeclarationOrigin.MULTI_FIELD_VALUE_CLASS_REPRESENTATION_VARIABLE,
-                                        saveVariable = ::variablesSaver,
-                                    )
-                                }
-                        }
+        return makeNewLambdaBlock(originalFunction, ref, makeBody = { wrapper -> makeLambdaBody(wrapper, replacement) })
+    }
+
+    private fun makeLambdaBody(wrapper: IrSimpleFunction, replacement: IrFunction): IrBody {
+        variablesToAdd[replacement]?.let {
+            variablesToAdd[wrapper] = it
+            variablesToAdd.remove(replacement)
+        }
+        if (replacement in possibleExtraBoxUsageGenerated) {
+            possibleExtraBoxUsageGenerated.add(wrapper)
+            possibleExtraBoxUsageGenerated.remove(replacement)
+        }
+        return with(context.createJvmIrBuilder(wrapper.symbol)) {
+            irExprBody(irBlock {
+                val newArguments: List<IrValueDeclaration> = wrapper.parameters.flatMap { parameter ->
+                    if (!parameter.type.needsMfvcFlattening()) {
+                        listOf(parameter)
+                    } else {
+                        // Old parameter value will be only used to set parameter of the lowered function,
+                        // thus it is useless to show it in debugger
+                        parameter.origin = JvmLoweredDeclarationOrigin.TEMPORARY_MULTI_FIELD_VALUE_CLASS_PARAMETER
+                        val rootNode = replacements.getRootMfvcNode(parameter.type.erasedUpperBound)
+                        rootNode.createInstanceFromBox(this, irGet(parameter), AccessType.ChooseEffective, ::variablesSaver)
+                            .makeFlattenedGetterExpressions(this, irCurrentClass, ::registerPossibleExtraBoxUsage)
+                            .mapIndexed { index, expression ->
+                                savableStandaloneVariableWithSetter(
+                                    expression = expression,
+                                    name = "${parameter.name.asString()}-${rootNode.leaves[index].fullFieldName}",
+                                    origin = JvmLoweredDeclarationOrigin.MULTI_FIELD_VALUE_CLASS_REPRESENTATION_VARIABLE,
+                                    saveVariable = ::variablesSaver,
+                                )
+                            }
                     }
-                    +replacement.inline(wrapper.parent, newArguments)
-                })
-            }
-        })
+                }
+                +replacement.inline(wrapper.parent, newArguments)
+            })
+        }
     }
 
     override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
@@ -814,12 +816,54 @@ internal class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : Jvm
         }.unwrapBlock()
     }
 
+    override fun visitRichFunctionReference(expression: IrRichFunctionReference): IrExpression {
+        val originalFunction = expression.invokeFunction
+        val replacement = originalFunction.getReplacement()
+        if (replacement == null) {
+            expression.invokeFunction = visitFunctionNew(originalFunction) as IrSimpleFunction
+            return expression
+        }
+        val declarations = transformFunctionFlat(originalFunction)
+        require(declarations == listOf(replacement)) {
+            "Expected ${replacement.render()}, got ${declarations?.map { it.render() }}"
+        }
+        val wrapper = makeNewLamdaFunction(originalFunction, makeBody = { wrapper -> makeLambdaBody(wrapper, replacement) })
+        expression.boundValues.transformInPlace(this, null)
+        expression.invokeFunction = wrapper
+        return expression
+    }
+
     private fun IrFunction.getReplacement(): IrFunction? = replacements.getReplacementFunction(this)
         ?: (this as? IrConstructor)?.let { replacements.getReplacementForRegularClassConstructor(it) }
 
-    private fun makeNewLambda(
+    private fun makeNewLambdaBlock(
         originalFunction: IrFunction, expression: IrFunctionReference, makeBody: (wrapper: IrSimpleFunction) -> IrBody
     ): IrContainerExpression {
+        val wrapper = makeNewLamdaFunction(originalFunction, makeBody)
+
+        val newReference = IrFunctionReferenceImpl(
+            startOffset = UNDEFINED_OFFSET,
+            endOffset = UNDEFINED_OFFSET,
+            type = expression.type,
+            symbol = wrapper.symbol,
+            typeArgumentsCount = expression.typeArguments.size,
+            reflectionTarget = expression.reflectionTarget,
+            origin = expression.origin,
+        ).apply {
+            copyTypeArgumentsFrom(expression)
+            arguments.assignFrom(expression.arguments) { it?.transform(this@JvmMultiFieldValueClassLowering, null) }
+            copyAttributes(expression)
+        }
+        return context.createJvmIrBuilder(getCurrentScopeSymbol(), expression).irBlock(origin = IrStatementOrigin.LAMBDA) {
+            +wrapper
+            +newReference
+        }
+    }
+
+    private fun makeNewLamdaFunction(
+        originalFunction: IrFunction,
+        makeBody: (IrSimpleFunction) -> IrBody,
+    ): IrSimpleFunction {
         val currentDeclarationParent = currentDeclarationParent!!
         val wrapper = context.irFactory.buildFun {
             updateFrom(originalFunction)
@@ -847,24 +891,7 @@ internal class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : Jvm
                 postActionAfterTransformingClassDeclaration(this)
             }
         }
-
-        val newReference = IrFunctionReferenceImpl(
-            startOffset = UNDEFINED_OFFSET,
-            endOffset = UNDEFINED_OFFSET,
-            type = expression.type,
-            symbol = wrapper.symbol,
-            typeArgumentsCount = expression.typeArguments.size,
-            reflectionTarget = expression.reflectionTarget,
-            origin = expression.origin,
-        ).apply {
-            copyTypeArgumentsFrom(expression)
-            arguments.assignFrom(expression.arguments) { it?.transform(this@JvmMultiFieldValueClassLowering, null) }
-            copyAttributes(expression)
-        }
-        return context.createJvmIrBuilder(getCurrentScopeSymbol(), expression).irBlock(origin = IrStatementOrigin.LAMBDA) {
-            +wrapper
-            +newReference
-        }
+        return wrapper
     }
 
     override fun visitFunctionAccess(expression: IrFunctionAccessExpression): IrExpression {