JVM: write back to tailrec function parameters if possible

This makes the debugger update the values when entering a recursive
call, and also uses fewer variable slots.

This may also work on the JS backend, but Native would probably require
something else.

^KT-47203 Fixed

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

@@ -28,6 +28,8 @@ import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.util.usesDefaultArguments
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 
+data class TailCalls(val ir: Set<IrCall>, val fromManyFunctions: Boolean)
+
 /**
  * Collects calls to be treated as tail recursion.
  * The checks are partially based on the frontend implementation
@@ -37,87 +39,78 @@ 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, followFunctionReference: (IrFunctionReference) -> Boolean): Set<IrCall> {
+fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (IrFunctionReference) -> Boolean): TailCalls {
     if ((irFunction as? IrSimpleFunction)?.isTailrec != true) {
-        return emptySet()
+        return TailCalls(emptySet(), false)
     }
 
+    class VisitorState(val isTailExpression: Boolean, val inOtherFunction: Boolean)
+
     val isUnitReturn = irFunction.returnType.isUnit()
-
     val result = mutableSetOf<IrCall>()
-    val visitor = object : IrElementVisitor<Unit, ElementKind> {
+    var someCallsAreInOtherFunctions = false
-
+    val visitor = object : IrElementVisitor<Unit, VisitorState> {
-        override fun visitElement(element: IrElement, data: ElementKind) {
+        override fun visitElement(element: IrElement, data: VisitorState) {
-            val childKind = ElementKind.NOT_SURE // Not sure by default.
+            element.acceptChildren(this, VisitorState(isTailExpression = false, data.inOtherFunction))
-            element.acceptChildren(this, childKind)
         }
 
-        override fun visitFunction(declaration: IrFunction, data: ElementKind) {
+        override fun visitFunction(declaration: IrFunction, data: VisitorState) {
             // Ignore local functions.
         }
 
-        override fun visitClass(declaration: IrClass, data: ElementKind) {
+        override fun visitClass(declaration: IrClass, data: VisitorState) {
             // Ignore local classes.
         }
 
-        override fun visitTry(aTry: IrTry, data: ElementKind) {
+        override fun visitTry(aTry: IrTry, data: VisitorState) {
             // We do not support tail calls in try-catch-finally, for simplicity of the mental model
             // very few cases there would be real tail-calls, and it's often not so easy for the user to see why
         }
 
-        override fun visitReturn(expression: IrReturn, data: ElementKind) {
+        override fun visitReturn(expression: IrReturn, data: VisitorState) {
-            val valueKind = if (expression.returnTargetSymbol == irFunction.symbol) {
+            expression.value.accept(this, VisitorState(expression.returnTargetSymbol == irFunction.symbol, data.inOtherFunction))
-                ElementKind.TAIL_STATEMENT
-            } else {
-                ElementKind.NOT_SURE
-            }
-            expression.value.accept(this, valueKind)
         }
 
-        override fun visitExpressionBody(body: IrExpressionBody, data: ElementKind) =
+        override fun visitExpressionBody(body: IrExpressionBody, data: VisitorState) =
             body.acceptChildren(this, data)
 
-        override fun visitBlockBody(body: IrBlockBody, data: ElementKind) =
+        override fun visitBlockBody(body: IrBlockBody, data: VisitorState) =
             visitStatementContainer(body, data)
 
-        override fun visitContainerExpression(expression: IrContainerExpression, data: ElementKind) =
+        override fun visitContainerExpression(expression: IrContainerExpression, data: VisitorState) =
             visitStatementContainer(expression, data)
 
-        private fun visitStatementContainer(expression: IrStatementContainer, data: ElementKind) {
+        private fun visitStatementContainer(expression: IrStatementContainer, data: VisitorState) {
             expression.statements.forEachIndexed { index, irStatement ->
-                val statementKind = when {
+                val isTailStatement = if (index == expression.statements.lastIndex) {
                     // The last statement defines the result of the container expression, so it has the same kind.
-                    index == expression.statements.lastIndex -> data
+                    data.isTailExpression
+                } else {
                     // In a Unit-returning function, any statement directly followed by a `return` is a tail statement.
                     isUnitReturn && expression.statements[index + 1].let {
                         it is IrReturn && it.returnTargetSymbol == irFunction.symbol && it.value.isUnitRead()
-                    } -> ElementKind.TAIL_STATEMENT
+                    }
-                    else -> ElementKind.NOT_SURE
                 }
-                irStatement.accept(this, statementKind)
+                irStatement.accept(this, VisitorState(isTailStatement, data.inOtherFunction))
             }
         }
 
         private fun IrExpression.isUnitRead(): Boolean =
             this is IrGetObjectValue && symbol.isClassWithFqName(StandardNames.FqNames.unit)
 
-        override fun visitWhen(expression: IrWhen, data: ElementKind) {
+        override fun visitWhen(expression: IrWhen, data: VisitorState) {
             expression.branches.forEach {
-                it.condition.accept(this, ElementKind.NOT_SURE)
+                it.condition.accept(this, VisitorState(isTailExpression = false, data.inOtherFunction))
                 it.result.accept(this, data)
             }
         }
 
-        override fun visitCall(expression: IrCall, data: ElementKind) {
+        override fun visitCall(expression: IrCall, data: VisitorState) {
-            expression.acceptChildren(this, ElementKind.NOT_SURE)
+            expression.acceptChildren(this, VisitorState(isTailExpression = false, data.inOtherFunction))
 
-            // Is it a tail call?
+            // TODO: the frontend generates diagnostics on calls that are not optimized. This may or may not
-            if (data != ElementKind.TAIL_STATEMENT) {
+            //   match what the backend does here. It'd be great to validate that the two are in agreement.
-                return
+            if (!data.isTailExpression || expression.symbol != irFunction.symbol) {
-            }
-
-            // Is it a recursive call?
-            if (expression.symbol != irFunction.symbol) {
                 return
             }
             // TODO: check type arguments
@@ -126,31 +119,30 @@ fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (
                 // Overridden functions using default arguments at tail call are not included: KT-4285
                 return
             }
-            val dispatchReceiverType = irFunction.dispatchReceiverParameter?.type
+
-            if (dispatchReceiverType?.classOrNull?.owner?.kind?.isSingleton == true) {
+            val hasSameDispatchReceiver =
-                // Dispatch receiver type is singleton and hence it can't be changed and the call must be tailrec.
+                irFunction.dispatchReceiverParameter?.type?.classOrNull?.owner?.kind?.isSingleton == true ||
-                result.add(expression)
+                        expression.dispatchReceiver?.let { it is IrGetValue && it.symbol.owner == irFunction.dispatchReceiverParameter } != false
+            if (!hasSameDispatchReceiver) {
+                // A tail call is not allowed to change dispatch receiver
+                //   class C {
+                //       fun foo(other: C) {
+                //           other.foo(this) // not a tail call
+                //       }
+                //   }
+                // TODO: KT-15341 - if the tailrec function is neither `override` nor `open`, this is fine actually?
+                //   Probably requires editing the frontend too.
                 return
             }
 
-
+            if (data.inOtherFunction) {
-            expression.dispatchReceiver?.let {
+                someCallsAreInOtherFunctions = true
-                if (it !is IrGetValue || it.symbol.owner != irFunction.dispatchReceiverParameter) {
-                    // A tail call is not allowed to change dispatch receiver
-                    //   class C {
-                    //       fun foo(other: C) {
-                    //           other.foo(this) // not a tail call
-                    //       }
-                    //   }
-                    return
-                }
             }
-
             result.add(expression)
         }
 
-        override fun visitFunctionReference(expression: IrFunctionReference, data: ElementKind) {
+        override fun visitFunctionReference(expression: IrFunctionReference, data: VisitorState) {
-            expression.acceptChildren(this, ElementKind.NOT_SURE)
+            expression.acceptChildren(this, VisitorState(isTailExpression = false, data.inOtherFunction))
             // This should match inline lambdas:
             //   tailrec fun foo() {
             //     run { return foo() } // non-local return from `foo`, so this *is* a tail call
@@ -160,27 +152,11 @@ fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (
             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)
+                expression.symbol.owner.body?.accept(this, VisitorState(isTailExpression = false, inOtherFunction = true))
             }
         }
     }
 
-    irFunction.body?.accept(visitor, ElementKind.TAIL_STATEMENT)
+    irFunction.body?.accept(visitor, VisitorState(isTailExpression = true, inOtherFunction = false))
-    return result
+    return TailCalls(result, someCallsAreInOtherFunctions)
-}
-
-/**
- * The kind of IR element used to detect tail calls.
- */
-private enum class ElementKind {
-    /**
-     * This element is the last statement to be executed before the return from the function.
-     * If the return type is not `Unit`, the result of this statement defines the result of the entire function.
-     */
-    TAIL_STATEMENT,
-
-    /**
-     * Not sure if the element meets the requirements to be [TAIL_STATEMENT].
-     */
-    NOT_SURE
 }

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

@@ -75,7 +75,7 @@ open class TailrecLowering(val context: BackendContext) : BodyLoweringPass {
 }
 
 private fun TailrecLowering.lowerTailRecursionCalls(irFunction: IrFunction) {
-    val tailRecursionCalls = collectTailRecursionCalls(irFunction, ::followFunctionReference)
+    val (tailRecursionCalls, someCallsAreFromOtherFunctions) = collectTailRecursionCalls(irFunction, ::followFunctionReference)
     if (tailRecursionCalls.isEmpty()) {
         return
     }
@@ -84,33 +84,41 @@ private fun TailrecLowering.lowerTailRecursionCalls(irFunction: IrFunction) {
     val oldBodyStatements = ArrayList(oldBody.statements)
     val builder = context.createIrBuilder(irFunction.symbol).at(oldBody)
 
-    val parameters = irFunction.explicitParameters
-
     oldBody.statements.clear()
     oldBody.statements += builder.irBlockBody {
-        // Define variables containing current values of parameters:
+        // `return recursiveCall(...)` is rewritten into assignments to parameters followed by a jump to the start.
-        val parameterToVariable = parameters.associateWith {
+        // While we may be able to write to the parameters directly, the recursive call may be inside an inline lambda,
-            createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true)
+        // so the parameters are captured and assigning to them requires temporarily rewriting their types (see
+        // `SharedVariablesLowering`), and that we can't do. So we have to create new `var`s for this purpose.
+        // TODO: an optimization pass will rewrite the types of vars back since the lambdas are guaranteed to be inlined
+        //  in place (otherwise they can't jump to the start of the function at all), so this is all a waste of CPU time.
+        val parameterToVariable = irFunction.explicitParameters.associateWith {
+            if (someCallsAreFromOtherFunctions || !it.isAssignable)
+                createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true)
+            else
+                it
         }
-        // (these variables are to be updated on any tail call).
-
-        +irWhile().apply {
-            val loop = this
-            condition = irTrue()
 
+        +irDoWhile().apply loop@{
             body = irBlock(startOffset, endOffset, resultType = context.irBuiltIns.unitType) {
-                // Read variables containing current values of parameters:
-                val parameterToNew = parameters.associateWith {
-                    createTmpVariable(irGet(parameterToVariable[it]!!), nameHint = it.symbol.suggestVariableName())
-                }
                 val transformer = BodyTransformer(
-                    this@lowerTailRecursionCalls, builder, irFunction, loop, parameterToNew, parameterToVariable, tailRecursionCalls
+                    this@lowerTailRecursionCalls, builder, irFunction, this@loop, parameterToVariable, tailRecursionCalls
                 )
                 oldBodyStatements.forEach {
                     +it.transformStatement(transformer)
                 }
-
+                +irBreak(this@loop)
-                +irBreak(loop)
+            }
+            condition = irBlock {
+                // The problem with creating new `var`s is that they do not show up in the debugger, so stopping inside
+                // a nested call will still display the parameters from the outermost call. To fix this, we need to
+                // write the new values back even though the parameters are now otherwise unused.
+                for ((parameter, variable) in parameterToVariable.entries) {
+                    if (parameter.isAssignable && parameter !== variable) {
+                        +irSet(parameter, irGet(variable))
+                    }
+                }
+                +irTrue()
             }
         }
     }.statements
@@ -124,10 +132,9 @@ private class BodyTransformer(
     private val builder: IrBuilderWithScope,
     irFunction: IrFunction,
     private val loop: IrLoop,
-    parameterToNew: Map<IrValueParameter, IrValueDeclaration>,
+    private val parameterToVariable: Map<IrValueParameter, IrValueDeclaration>,
-    private val parameterToVariable: Map<IrValueParameter, IrVariable>,
     private val tailRecursionCalls: Set<IrCall>,
-) : VariableRemapper(parameterToNew) {
+) : VariableRemapper(parameterToVariable) {
 
     val parameters = irFunction.explicitParameters
 

compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmGeneratorExtensionsImpl.kt

@@ -249,4 +249,7 @@ open class JvmGeneratorExtensionsImpl(
         }
         return null
     }
+
+    override val parametersAreAssignable: Boolean
+        get() = true
 }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt

@@ -371,15 +371,16 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
         ktElement: KtPureElement?,
         irOwnerElement: IrElement
     ): IrValueParameter {
-        if (context.languageVersionSettings.supportsFeature(LanguageFeature.NewCapturedReceiverFieldNamingConvention)) {
+        val name = if (context.languageVersionSettings.supportsFeature(LanguageFeature.NewCapturedReceiverFieldNamingConvention)) {
             if (ktElement is KtFunctionLiteral) {
-                val name = getCallLabelForLambdaArgument(ktElement, this.context.bindingContext)?.let {
+                val label = getCallLabelForLambdaArgument(ktElement, this.context.bindingContext)?.let {
                     it.takeIf(Name::isValidIdentifier) ?: "\$receiver"
                 }
-                return declareParameter(receiverParameterDescriptor, ktElement, irOwnerElement, name = Name.identifier("\$this\$$name"))
+                // TODO: this can produce `$this$null` - expected?
-            }
+                Name.identifier("\$this\$$label")
-        }
+            } else null
-        return declareParameter(receiverParameterDescriptor, ktElement, irOwnerElement)
+        } else null
+        return declareParameter(receiverParameterDescriptor, ktElement, irOwnerElement, name)
     }
 
     private fun getCallLabelForLambdaArgument(declaration: KtFunctionLiteral, bindingContext: BindingContext): String? {
@@ -427,7 +428,8 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
             descriptor, descriptor.type.toIrType(),
             (descriptor as? ValueParameterDescriptor)?.varargElementType?.toIrType(),
             name,
-            index
+            index,
+            isAssignable = (irOwnerElement as? IrSimpleFunction)?.isTailrec == true && context.extensions.parametersAreAssignable
         )
     }
 

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/GeneratorExtensions.kt

@@ -45,4 +45,7 @@ open class GeneratorExtensions : StubGeneratorExtensions() {
     open fun unwrapSyntheticJavaProperty(descriptor: PropertyDescriptor): Pair<FunctionDescriptor, FunctionDescriptor?>? = null
 
     open fun remapDebuggerFieldPropertyDescriptor(propertyDescriptor: PropertyDescriptor): PropertyDescriptor = propertyDescriptor
+
+    open val parametersAreAssignable: Boolean
+        get() = false
 }