diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailRecursionCallsCollector.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailRecursionCallsCollector.kt index 12480482677..a16702fa198 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailRecursionCallsCollector.kt +++ b/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, 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 { +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() - val visitor = object : IrElementVisitor { - - override fun visitElement(element: IrElement, data: ElementKind) { - val childKind = ElementKind.NOT_SURE // Not sure by default. - element.acceptChildren(this, childKind) + var someCallsAreInOtherFunctions = false + val visitor = object : IrElementVisitor { + override fun visitElement(element: IrElement, data: VisitorState) { + element.acceptChildren(this, VisitorState(isTailExpression = false, data.inOtherFunction)) } - 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) { - val valueKind = if (expression.returnTargetSymbol == irFunction.symbol) { - ElementKind.TAIL_STATEMENT - } else { - ElementKind.NOT_SURE - } - expression.value.accept(this, valueKind) + override fun visitReturn(expression: IrReturn, data: VisitorState) { + expression.value.accept(this, VisitorState(expression.returnTargetSymbol == irFunction.symbol, data.inOtherFunction)) } - 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) { - expression.acceptChildren(this, ElementKind.NOT_SURE) + override fun visitCall(expression: IrCall, data: VisitorState) { + expression.acceptChildren(this, VisitorState(isTailExpression = false, data.inOtherFunction)) - // Is it a tail call? - if (data != ElementKind.TAIL_STATEMENT) { - return - } - - // Is it a recursive call? - if (expression.symbol != irFunction.symbol) { + // TODO: the frontend generates diagnostics on calls that are not optimized. This may or may not + // match what the backend does here. It'd be great to validate that the two are in agreement. + if (!data.isTailExpression || 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) { - // Dispatch receiver type is singleton and hence it can't be changed and the call must be tailrec. - result.add(expression) + + val hasSameDispatchReceiver = + irFunction.dispatchReceiverParameter?.type?.classOrNull?.owner?.kind?.isSingleton == true || + 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 } - - expression.dispatchReceiver?.let { - 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 - } + if (data.inOtherFunction) { + someCallsAreInOtherFunctions = true } - result.add(expression) } - override fun visitFunctionReference(expression: IrFunctionReference, data: ElementKind) { - expression.acceptChildren(this, ElementKind.NOT_SURE) + override fun visitFunctionReference(expression: IrFunctionReference, data: VisitorState) { + 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) - return result -} - -/** - * 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 + irFunction.body?.accept(visitor, VisitorState(isTailExpression = true, inOtherFunction = false)) + return TailCalls(result, someCallsAreInOtherFunctions) } diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/TailrecLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/TailrecLowering.kt index 8461839f575..4f9c96d8880 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/TailrecLowering.kt +++ b/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: - val parameterToVariable = parameters.associateWith { - createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true) + // `return recursiveCall(...)` is rewritten into assignments to parameters followed by a jump to the start. + // While we may be able to write to the parameters directly, the recursive call may be inside an inline lambda, + // 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(loop) + +irBreak(this@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, - private val parameterToVariable: Map, + private val parameterToVariable: Map, private val tailRecursionCalls: Set, -) : VariableRemapper(parameterToNew) { +) : VariableRemapper(parameterToVariable) { val parameters = irFunction.explicitParameters diff --git a/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmGeneratorExtensionsImpl.kt b/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmGeneratorExtensionsImpl.kt index 4aa8f320bf1..03df050ef70 100644 --- a/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmGeneratorExtensionsImpl.kt +++ b/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 } diff --git a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt index 1587d1e962d..50a73316703 100644 --- a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt +++ b/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")) - } - } - return declareParameter(receiverParameterDescriptor, ktElement, irOwnerElement) + // TODO: this can produce `$this$null` - expected? + Name.identifier("\$this\$$label") + } else null + } 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 ) } diff --git a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/GeneratorExtensions.kt b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/GeneratorExtensions.kt index c1d3544e554..970eb19ac54 100644 --- a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/GeneratorExtensions.kt +++ b/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? = null open fun remapDebuggerFieldPropertyDescriptor(propertyDescriptor: PropertyDescriptor): PropertyDescriptor = propertyDescriptor + + open val parametersAreAssignable: Boolean + get() = false }