From 2afab62dae3a0f95e363a4261fa619962069a15b Mon Sep 17 00:00:00 2001 From: pyos Date: Mon, 6 Sep 2021 13:06:42 +0200 Subject: [PATCH] JVM_IR: optimize tailrec calls in inline lambdas ^KT-48600 Fixed --- .../common/TailRecursionCallsCollector.kt | 17 +++++++++++++- .../backend/common/lower/TailrecLowering.kt | 22 ++++++++++++++----- .../backend/jvm/lower/JvmTailrecLowering.kt | 7 +++++- .../recursiveCallInInlineLambda.kt | 4 ++-- 4 files changed, 41 insertions(+), 9 deletions(-) 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 26c97f32455..55cd7a004bb 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 @@ -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 { +fun collectTailRecursionCalls(irFunction: IrFunction, followFunctionReference: (IrFunctionReference) -> Boolean): Set { if ((irFunction as? IrSimpleFunction)?.isTailrec != true) { return emptySet() } @@ -147,6 +147,21 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set { 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) 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 87d4bd84e07..62497d12b63 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 @@ -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, val parameterToVariable: Map, val tailRecursionCalls: Set, - 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) -> diff --git a/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/JvmTailrecLowering.kt b/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/JvmTailrecLowering.kt index b90799ae761..38c4aa9ab8a 100644 --- a/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/JvmTailrecLowering.kt +++ b/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) -} \ No newline at end of file + + override fun followFunctionReference(reference: IrFunctionReference): Boolean = + reference.origin == JvmLoweredStatementOrigin.INLINE_LAMBDA +} diff --git a/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/recursiveCallInInlineLambda.kt b/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/recursiveCallInInlineLambda.kt index 6c82ccdab20..21c0423839c 100644 --- a/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/recursiveCallInInlineLambda.kt +++ b/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