From b2315a4a054a8af6366d0929b4d0cb3e39dd71ea Mon Sep 17 00:00:00 2001 From: pyos Date: Mon, 6 Sep 2021 11:25:14 +0200 Subject: [PATCH] IR: optimize `f(); return` in `tailrec fun f(): Unit` It's equivalent to `return f()`. --- .../common/TailRecursionCallsCollector.kt | 55 +++++++++++-------- .../backend/common/lower/TailrecLowering.kt | 22 +++++--- .../functions/tailRecursion/unitBlocks.kt | 10 ++-- 3 files changed, 50 insertions(+), 37 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 cbd95fb2648..29d98ac788c 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 @@ -16,15 +16,15 @@ package org.jetbrains.kotlin.backend.common -import org.jetbrains.kotlin.descriptors.ClassKind import org.jetbrains.kotlin.ir.IrElement import org.jetbrains.kotlin.ir.declarations.IrClass import org.jetbrains.kotlin.ir.declarations.IrFunction import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction import org.jetbrains.kotlin.ir.expressions.* -import org.jetbrains.kotlin.ir.types.IrSimpleType import org.jetbrains.kotlin.ir.types.classOrNull +import org.jetbrains.kotlin.ir.types.IdSignatureValues import org.jetbrains.kotlin.ir.types.isUnit +import org.jetbrains.kotlin.ir.util.parentClassOrNull import org.jetbrains.kotlin.ir.util.usesDefaultArguments import org.jetbrains.kotlin.ir.visitors.IrElementVisitor @@ -42,8 +42,9 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set { return emptySet() } - val result = mutableSetOf() + val isUnitReturn = irFunction.returnType.isUnit() + val result = mutableSetOf() val visitor = object : IrElementVisitor { override fun visitElement(element: IrElement, data: ElementKind) { @@ -73,18 +74,39 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set { expression.value.accept(this, valueKind) } - override fun visitContainerExpression(expression: IrContainerExpression, data: ElementKind) { + override fun visitExpressionBody(body: IrExpressionBody, data: ElementKind) = + body.acceptChildren(this, data) + + override fun visitBlockBody(body: IrBlockBody, data: ElementKind) = + visitStatementContainer(body, data) + + override fun visitContainerExpression(expression: IrContainerExpression, data: ElementKind) = + visitStatementContainer(expression, data) + + private fun visitStatementContainer(expression: IrStatementContainer, data: ElementKind) { expression.statements.forEachIndexed { index, irStatement -> - val statementKind = if (index == expression.statements.lastIndex) { + val statementKind = when { // The last statement defines the result of the container expression, so it has the same kind. - data - } else { - ElementKind.NOT_SURE + index == expression.statements.lastIndex -> data + // 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) } } + private fun IrExpression.isUnitRead(): Boolean = + when (this) { + is IrGetObjectValue -> symbol + // On the JVM, if SingletonReferencesLowering has already finished, a `Unit` reference + // is now an IrGetField to the INSTANCE field. + is IrGetField -> symbol.owner.parentClassOrNull?.symbol + else -> null + }?.signature == IdSignatureValues.unit + override fun visitWhen(expression: IrWhen, data: ElementKind) { expression.branches.forEach { it.condition.accept(this, ElementKind.NOT_SURE) @@ -131,25 +153,10 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set { } result.add(expression) - } - - } - - val body = irFunction.body - if (body !is IrBlockBody) { - return emptySet() // TODO: should an assert be here instead? - } - - body.statements.forEachIndexed { index, irStatement -> - val kind = if (index == body.statements.lastIndex && irFunction.returnType.isUnit()) { - ElementKind.TAIL_STATEMENT - } else { - ElementKind.NOT_SURE - } - irStatement.accept(visitor, kind) } + irFunction.body?.accept(visitor, ElementKind.TAIL_STATEMENT) return result } 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 f9ab5b4e9d3..d3f0740e42e 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 @@ -52,18 +52,25 @@ open class TailrecLowering(val context: BackendContext) : BodyLoweringPass { override fun visitFunction(declaration: IrFunction) { declaration.acceptChildrenVoid(this) - lowerTailRecursionCalls(context, declaration, useProperComputationOrderOfTailrecDefaultParameters()) + lowerTailRecursionCalls(declaration) } }) - lowerTailRecursionCalls(context, container, useProperComputationOrderOfTailrecDefaultParameters()) + lowerTailRecursionCalls(container) } } + private fun lowerTailRecursionCalls(function: IrFunction) = + lowerTailRecursionCalls(context, function, useProperComputationOrderOfTailrecDefaultParameters()) + open fun useProperComputationOrderOfTailrecDefaultParameters() = true } -private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunction, properComputationOrderOfTailrecDefaultParameters: Boolean) { +private fun lowerTailRecursionCalls( + context: BackendContext, + irFunction: IrFunction, + properComputationOrderOfTailrecDefaultParameters: Boolean +) { val tailRecursionCalls = collectTailRecursionCalls(irFunction) if (tailRecursionCalls.isEmpty()) { return @@ -78,8 +85,8 @@ private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunct oldBody.statements.clear() oldBody.statements += builder.irBlockBody { // Define variables containing current values of parameters: - val parameterToVariable = parameters.associate { - it to createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true) + val parameterToVariable = parameters.associateWith { + createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true) } // (these variables are to be updated on any tail call). @@ -89,9 +96,8 @@ private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunct body = irBlock(startOffset, endOffset, resultType = context.irBuiltIns.unitType) { // Read variables containing current values of parameters: - val parameterToNew = parameters.associate { - val variable = parameterToVariable[it]!! - it to createTmpVariable(irGet(variable), nameHint = it.symbol.suggestVariableName()) + val parameterToNew = parameters.associateWith { + createTmpVariable(irGet(parameterToVariable[it]!!), nameHint = it.symbol.suggestVariableName()) } val transformer = BodyTransformer( diff --git a/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/unitBlocks.kt b/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/unitBlocks.kt index 7ae0c8f0e6a..80c0a3bd02d 100644 --- a/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/unitBlocks.kt +++ b/compiler/testData/codegen/box/diagnostics/functions/tailRecursion/unitBlocks.kt @@ -7,14 +7,14 @@ // IGNORE_BACKEND: JS tailrec fun test(x : Int) : Unit { - if (x == 1) { + if (x > 800000) { test(x - 1) - } else if (x == 2) { + } else if (x > 600000) { test(x - 1) return - } else if (x == 3) { - test(x - 1) - if (x == 3) { + } else if (x > 400000) { + test(1) + if (x > 200000) { test(x - 1) } return