[IR] Move JvmUpgradeCallableReferences before JvmTailrecLowering

#KT-74383

Author: zhelenskiy

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailRecursionCallsCollector.kt

@@ -39,7 +39,11 @@ data class TailCalls(val ir: Set<IrCall>, val fromManyFunctions: Boolean)
  * It is also not guaranteed that each returned call is detected as tail recursion by the frontend.
  * However any returned call can be correctly optimized as tail recursion.
  */
-fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (IrFunctionReference) -> Boolean): TailCalls {
+fun collectTailRecursionCalls(
+    irFunction: IrFunction,
+    followFunctionReference: (IrFunctionReference) -> Boolean,
+    followRichFunctionReference: (IrRichFunctionReference) -> Boolean
+): TailCalls {
     if ((irFunction as? IrSimpleFunction)?.isTailrec != true) {
         return TailCalls(emptySet(), false)
     }
@@ -155,6 +159,15 @@ fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (
                 expression.symbol.owner.body?.accept(this, VisitorState(isTailExpression = false, inOtherFunction = true))
             }
         }
+
+        override fun visitRichFunctionReference(expression: IrRichFunctionReference, data: VisitorState) {
+            expression.acceptChildren(this, VisitorState(isTailExpression = false, data.inOtherFunction))
+            if (followRichFunctionReference(expression)) {
+                // If control reaches end of lambda, it will *not* end the current function by default,
+                // so the lambda's body itself is not a tail statement.
+                expression.invokeFunction.body?.accept(this, VisitorState(isTailExpression = false, inOtherFunction = true))
+            }
+        }
     }
 
     irFunction.body?.accept(visitor, VisitorState(isTailExpression = true, inOtherFunction = false))

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/TailrecLowering.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.defaultValueForType
 import org.jetbrains.kotlin.ir.util.getArgumentsWithIr
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.ir.util.transformInPlace
 import org.jetbrains.kotlin.ir.visitors.IrVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
@@ -69,12 +70,14 @@ open class TailrecLowering(val context: LoweringContext) : BodyLoweringPass {
 
     open fun followFunctionReference(reference: IrFunctionReference): Boolean = false
 
+    open fun followRichFunctionReference(reference: IrRichFunctionReference): Boolean = false
+
     open fun nullConst(startOffset: Int, endOffset: Int, type: IrType): IrExpression =
         IrConstImpl.defaultValueForType(startOffset, endOffset, type)
 }
 
 private fun TailrecLowering.lowerTailRecursionCalls(irFunction: IrFunction) {
-    val (tailRecursionCalls, someCallsAreFromOtherFunctions) = collectTailRecursionCalls(irFunction, ::followFunctionReference)
+    val (tailRecursionCalls, someCallsAreFromOtherFunctions) = collectTailRecursionCalls(irFunction, ::followFunctionReference, ::followRichFunctionReference)
     if (tailRecursionCalls.isEmpty()) {
         return
     }
@@ -152,6 +155,15 @@ private class BodyTransformer(
         return super.visitFunctionReference(expression)
     }
 
+    override fun visitRichFunctionReference(expression: IrRichFunctionReference): IrExpression {
+        return if (lowering.followRichFunctionReference(expression)) {
+            super.visitRichFunctionReference(expression)
+        } else {
+            expression.boundValues.transformInPlace(this, null)
+            expression
+        }
+    }
+
     private fun IrBuilderWithScope.genTailCall(expression: IrCall) = this.irBlock(expression) {
         // Get all specified arguments:
         val parameterToArgument = expression.getArgumentsWithIr().associateTo(mutableMapOf()) { (parameter, argument) ->

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

@@ -67,14 +67,14 @@ private val jvmFilePhases = createFilePhases(
     ::JvmSingleAbstractMethodLowering,
     ::JvmMultiFieldValueClassLowering,
     ::JvmInlineClassLowering,
+    ::JvmUpgradeCallableReferences,
     ::JvmTailrecLowering,
 
     ::MappedEnumWhenLowering,
 
     ::AssertionLowering,
     ::JvmReturnableBlockLowering,
     ::SingletonReferencesLowering,
-    ::JvmUpgradeCallableReferences,
     ::JvmSharedVariablesLowering,
 
     ::JvmInventNamesForLocalFunctions,

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.jvm.ir.defaultValue
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
+import org.jetbrains.kotlin.ir.expressions.IrRichFunctionReference
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.types.IrType
 
@@ -21,6 +22,9 @@ internal class JvmTailrecLowering(context: JvmBackendContext) : TailrecLowering(
     override fun followFunctionReference(reference: IrFunctionReference): Boolean =
         reference.origin == IrStatementOrigin.INLINE_LAMBDA
 
+    override fun followRichFunctionReference(reference: IrRichFunctionReference): Boolean =
+        reference.origin == IrStatementOrigin.INLINE_LAMBDA
+
     override fun nullConst(startOffset: Int, endOffset: Int, type: IrType): IrExpression =
         type.defaultValue(startOffset, endOffset, context as JvmBackendContext)
 }