JVM_IR: optimize tailrec calls in inline lambdas

^KT-48600 Fixed

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

@@ -36,7 +36,7 @@ import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
  * 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): Set<IrCall> {
+fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (IrFunctionReference) -> Boolean): Set<IrCall> {
     if ((irFunction as? IrSimpleFunction)?.isTailrec != true) {
         return emptySet()
     }
@@ -147,6 +147,21 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set<IrCall> {
 
             result.add(expression)
         }
+
+        override fun visitFunctionReference(expression: IrFunctionReference, data: ElementKind) {
+            expression.acceptChildren(this, ElementKind.NOT_SURE)
+            // This should match inline lambdas:
+            //   tailrec fun foo() {
+            //     run { return foo() } // non-local return from `foo`, so this *is* a tail call
+            //   }
+            // Whether crossinline lambdas are matched is unimportant, as they can't contain any returns
+            // from `foo` anyway.
+            if (followFunctionReference(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.symbol.owner.body?.accept(this, ElementKind.NOT_SURE)
+            }
+        }
     }
 
     irFunction.body?.accept(visitor, ElementKind.TAIL_STATEMENT)

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

@@ -62,17 +62,20 @@ open class TailrecLowering(val context: BackendContext) : BodyLoweringPass {
     }
 
     private fun lowerTailRecursionCalls(function: IrFunction) =
-        lowerTailRecursionCalls(context, function, useProperComputationOrderOfTailrecDefaultParameters())
+        lowerTailRecursionCalls(context, function, useProperComputationOrderOfTailrecDefaultParameters(), ::followFunctionReference)
 
     open fun useProperComputationOrderOfTailrecDefaultParameters() = true
+
+    open fun followFunctionReference(reference: IrFunctionReference): Boolean = false
 }
 
 private fun lowerTailRecursionCalls(
     context: BackendContext,
     irFunction: IrFunction,
-    properComputationOrderOfTailrecDefaultParameters: Boolean
+    properComputationOrderOfTailrecDefaultParameters: Boolean,
+    followFunctionReference: (IrFunctionReference) -> Boolean
 ) {
-    val tailRecursionCalls = collectTailRecursionCalls(irFunction)
+    val tailRecursionCalls = collectTailRecursionCalls(irFunction, followFunctionReference)
     if (tailRecursionCalls.isEmpty()) {
         return
     }
@@ -104,7 +107,8 @@ private fun lowerTailRecursionCalls(
                 val transformer = BodyTransformer(
                     builder, irFunction, loop,
                     parameterToNew, parameterToVariable, tailRecursionCalls,
-                    properComputationOrderOfTailrecDefaultParameters
+                    properComputationOrderOfTailrecDefaultParameters,
+                    followFunctionReference
                 )
 
                 oldBodyStatements.forEach {
@@ -124,7 +128,8 @@ private class BodyTransformer(
     val parameterToNew: Map<IrValueParameter, IrValueDeclaration>,
     val parameterToVariable: Map<IrValueParameter, IrVariable>,
     val tailRecursionCalls: Set<IrCall>,
-    val properComputationOrderOfTailrecDefaultParameters: Boolean
+    val properComputationOrderOfTailrecDefaultParameters: Boolean,
+    val followFunctionReference: (IrFunctionReference) -> Boolean
 ) : VariableRemapper(parameterToNew) {
 
     val parameters = irFunction.explicitParameters
@@ -137,6 +142,13 @@ private class BodyTransformer(
         return builder.at(expression).genTailCall(expression)
     }
 
+    override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+        if (followFunctionReference(expression)) {
+            expression.symbol.owner.body?.transformChildrenVoid(this)
+        }
+        return super.visitFunctionReference(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/lower/JvmTailrecLowering.kt

@@ -7,10 +7,15 @@ package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.lower.TailrecLowering
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredStatementOrigin
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 
 class JvmTailrecLowering(context: JvmBackendContext) : TailrecLowering(context) {
     override fun useProperComputationOrderOfTailrecDefaultParameters(): Boolean =
         context.ir.context.configuration.languageVersionSettings.supportsFeature(LanguageFeature.ProperComputationOrderOfTailrecDefaultParameters)
-}
+
+    override fun followFunctionReference(reference: IrFunctionReference): Boolean =
+        reference.origin == JvmLoweredStatementOrigin.INLINE_LAMBDA
+}

compiler/testData/codegen/box/diagnostics/functions/tailRecursion/recursiveCallInInlineLambda.kt

@@ -1,5 +1,5 @@
 // KT-16549
-// IGNORE_BACKEND: JVM
+// IGNORE_BACKEND: JVM, JS
 // IGNORE_FIR_DIAGNOSTICS_DIFF
 
 class TailInline {
@@ -7,7 +7,7 @@ class TailInline {
         return action()
     }
 
-    private var countDown = 10
+    private var countDown = 100000
 
     tailrec fun test(): Int {
         if (countDown < 5) return countDown