azalea/kotlin-fork
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

JVM: write back to tailrec function parameters if possible

Browse Source 
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
This commit is contained in:
pyos
2022-04-25 13:48:21 +02:00
committed by max-kammerer
parent ca446f008e
commit 239bcea3b9
5 changed files with 94 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailRecursionCallsCollector.kt
+51 -75
View File
@@ -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> {
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<Unit, VisitorState> {
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)
}
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/TailrecLowering.kt
+28 -21
View File
@@ -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<IrValueParameter, IrValueDeclaration>,
private val parameterToVariable: Map<IrValueParameter, IrVariable>,
private val parameterToVariable: Map<IrValueParameter, IrValueDeclaration>,
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
+3
View File
@@ -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
+9 -7
View File
@@ -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
)
}
compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/GeneratorExtensions.kt
+3
View File
@@ -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
}