From 56c8fdc6c4c133a3d7c9f0dcd2cfce7398004bac Mon Sep 17 00:00:00 2001 From: Ilmir Usmanov Date: Wed, 8 Jan 2020 15:48:28 +0100 Subject: [PATCH] JVM_IR: Perform tail-call optimization on IR It is easier to introduce a new lowering so the codegen will emit code for the old tail-call optimizer to understand. Also, this is more flexible and would allow to optimize cases, which are now feasible with the old optimizer. Note, that because of bytecode inlining, we cannot replace the old one, but we cannot emit code, that is simpler for it to optimize. --- .../kotlin/backend/jvm/JvmBackendContext.kt | 4 +- .../jetbrains/kotlin/backend/jvm/JvmLower.kt | 3 +- .../backend/jvm/codegen/CoroutineCodegen.kt | 12 +- .../backend/jvm/codegen/ExpressionCodegen.kt | 9 +- .../jvm/lower/AddContinuationLowering.kt | 7 +- .../jvm/lower/TailCallOptimizationLowering.kt | 123 ++++++++++++++++++ .../tailSuspendUnitFun.kt | 1 - .../tailCallOptimizations/unreachable.kt | 1 - 8 files changed, 149 insertions(+), 11 deletions(-) create mode 100644 compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/TailCallOptimizationLowering.kt diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt index dc3f661f72d..a39eb9f5604 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt @@ -1,5 +1,5 @@ /* - * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. */ @@ -31,6 +31,7 @@ import org.jetbrains.kotlin.ir.builders.irNull import org.jetbrains.kotlin.ir.declarations.* import org.jetbrains.kotlin.ir.declarations.impl.IrConstructorImpl import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns +import org.jetbrains.kotlin.ir.expressions.IrCall import org.jetbrains.kotlin.ir.expressions.IrExpression import org.jetbrains.kotlin.ir.expressions.IrFunctionReference import org.jetbrains.kotlin.ir.symbols.* @@ -105,6 +106,7 @@ class JvmBackendContext( val continuationClassBuilders = mutableMapOf() val suspendFunctionOriginalToView = mutableMapOf() val suspendFunctionViewToOriginal = mutableMapOf() + val suspendTailCallsWithUnitReplacement = mutableSetOf() val fakeContinuation: IrExpression = createFakeContinuation(this) val staticDefaultStubs = mutableMapOf() diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt index 4239bb77182..df397c41160 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt @@ -1,5 +1,5 @@ /* - * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. */ @@ -295,6 +295,7 @@ private val jvmFilePhases = interfaceDefaultCallsPhase then interfaceObjectCallsPhase then + tailCallOptimizationPhase then addContinuationPhase then innerClassesPhase then diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt index ac1867d0917..46a84bb927d 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt @@ -1,5 +1,5 @@ /* - * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. */ @@ -72,12 +72,16 @@ internal fun generateStateMachineForNamedFunction( internalNameForDispatchReceiver = classCodegen.visitor.thisName, putContinuationParameterToLvt = false, disableTailCallOptimizationForFunctionReturningUnit = irFunction.returnType.isUnit() && - (irFunction as? IrSimpleFunction)?.allOverridden()?.toList()?.let { functions -> - functions.isNotEmpty() && functions.any { !it.returnType.isUnit() } - } == true + irFunction.anyOfOverriddenFunctionsReturnsNonUnit() ) } +internal fun IrFunction.anyOfOverriddenFunctionsReturnsNonUnit(): Boolean { + return (this as? IrSimpleFunction)?.allOverridden()?.toList()?.let { functions -> + functions.isNotEmpty() && functions.any { !it.returnType.isUnit() } + } == true +} + internal fun generateStateMachineForLambda( classCodegen: ClassCodegen, methodVisitor: MethodVisitor, diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt index dcae6697f2b..26fb205f56b 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt @@ -482,7 +482,7 @@ class ExpressionCodegen( irFunction.shouldNotContainSuspendMarkers() -> false // Noinline function are always suspension points !owner.isInline -> true - // The suspend intrinsics are, albeit inline, are also suspension points + // The suspend intrinsics are, albeit inline, also suspension points owner.fqNameForIrSerialization == FqName("kotlin.coroutines.intrinsics.IntrinsicsKt.suspendCoroutineUninterceptedOrReturn") -> true // Inside $$forInline functions crossinline calls are (usually) not suspension points, otherwise, flow will be pessimized (dispatchReceiver as? IrGetField)?.symbol?.owner?.origin == @@ -656,6 +656,13 @@ class ExpressionCodegen( val returnType = if (owner == irFunction) signature.returnType else methodSignatureMapper.mapReturnType(owner) val afterReturnLabel = Label() expression.value.accept(this, data).coerce(returnType, owner.returnType).materialize() + // We replaced COERTION_TO_UNIT with IrReturn during TailCallOptimizationLowering. + // Generate POP GETSTATIC kotlin/Unit.INSTANCE now. + // Otherwise, tail-call optimization will not work. See tailSuspendUnitFun.kt test. + if (expression in context.suspendTailCallsWithUnitReplacement) { + mv.pop() + mv.getstatic("kotlin/Unit", "INSTANCE", "Lkotlin/Unit;") + } generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data) expression.markLineNumber(startOffset = true) if (isNonLocalReturn) { diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt index 950cc246b8d..cffea93b4bb 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt @@ -1,5 +1,5 @@ /* - * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. */ @@ -14,6 +14,8 @@ import org.jetbrains.kotlin.backend.jvm.JvmBackendContext import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator import org.jetbrains.kotlin.backend.jvm.ir.defaultValue +import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrBlock +import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase import org.jetbrains.kotlin.codegen.coroutines.* import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX import org.jetbrains.kotlin.descriptors.Modality @@ -43,7 +45,8 @@ import org.jetbrains.kotlin.util.OperatorNameConventions internal val addContinuationPhase = makeIrFilePhase( ::AddContinuationLowering, "AddContinuation", - "Add continuation classes to suspend functions and transform suspend lambdas into continuations" + "Add continuation classes to suspend functions and transform suspend lambdas into continuations", + prerequisite = setOf(localDeclarationsPhase, tailCallOptimizationPhase) ) private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass { diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/TailCallOptimizationLowering.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/TailCallOptimizationLowering.kt new file mode 100644 index 00000000000..69c809c84c4 --- /dev/null +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/TailCallOptimizationLowering.kt @@ -0,0 +1,123 @@ +/* + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.kotlin.backend.jvm.lower + +import org.jetbrains.kotlin.backend.common.FileLoweringPass +import org.jetbrains.kotlin.backend.common.ir.isSuspend +import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase +import org.jetbrains.kotlin.backend.jvm.JvmBackendContext +import org.jetbrains.kotlin.backend.jvm.codegen.anyOfOverriddenFunctionsReturnsNonUnit +import org.jetbrains.kotlin.backend.jvm.codegen.isKnownToBeTailCall +import org.jetbrains.kotlin.ir.IrElement +import org.jetbrains.kotlin.ir.IrStatement +import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin +import org.jetbrains.kotlin.ir.declarations.IrFile +import org.jetbrains.kotlin.ir.declarations.IrFunction +import org.jetbrains.kotlin.ir.expressions.* +import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl +import org.jetbrains.kotlin.ir.types.isUnit +import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable +import org.jetbrains.kotlin.ir.util.isSuspend +import org.jetbrains.kotlin.ir.util.parentAsClass +import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid +import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid +import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid +import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid +import org.jetbrains.kotlin.name.FqName +import org.jetbrains.kotlin.name.Name + +internal val tailCallOptimizationPhase = makeIrFilePhase( + ::TailCallOptimizationLowering, + "TailCallOptimization", + "Add or move returns to suspension points on tail-call positions" +) + +// TODO: Make this lowering common +private class TailCallOptimizationLowering(private val context: JvmBackendContext) : IrElementVisitorVoid, FileLoweringPass { + override fun lower(irFile: IrFile) { + visitFile(irFile) + } + + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } + + override fun visitFunction(function: IrFunction) { + if (!function.isSuspend || function.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA || function.isKnownToBeTailCall()) return + // Disable tail-call optimization for functions, returning Unit and overriding function, returning non-unit. + if (function.returnType.isUnit() && function.anyOfOverriddenFunctionsReturnsNonUnit()) return + + val tailCalls = mutableSetOf() + + // Find all tail-calls inside suspend function. + // We should add IrReturn before them, so the codegen will generate code, which is understandable by old BE's tail-call optimizer. + // This function collect all possible tail-calls, if they are not prepended by IrReturns, so we can prepend them. + // If the call has IrReturn before it, we do not need to add another one. + // If the call is inside a branch of `when` statement (or `if`, which is lowered to `when`), go through the branches and + // find tail-calls there. + fun findCallsOnTailPositionWithoutImmediateReturn(statement: IrStatement, immediateReturn: Boolean = false) { + when (statement) { + is IrCall -> if (statement.isSuspend && !immediateReturn) { + tailCalls += statement + } + is IrBlock -> findCallsOnTailPositionWithoutImmediateReturn( + statement.statements.findTailCall(function.returnType.isUnit()) ?: return + ) + is IrWhen -> for (branch in statement.branches) { + findCallsOnTailPositionWithoutImmediateReturn(branch.result) + } + is IrReturn -> findCallsOnTailPositionWithoutImmediateReturn(statement.value, immediateReturn = true) + is IrTypeOperatorCall -> if (statement.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT) { + findCallsOnTailPositionWithoutImmediateReturn(statement.argument) + } + else -> { + // Not a call, or something containing a tail-call, break + // TODO: Support binary logical operations and elvis, though. KT-23826 and KT-23825 + } + } + } + + when (val body = function.body) { + null -> return + is IrBlockBody -> findCallsOnTailPositionWithoutImmediateReturn( + body.statements.findTailCall(function.returnType.isUnit()) ?: return + ) + is IrExpressionBody -> findCallsOnTailPositionWithoutImmediateReturn(body.expression) + else -> error("Unexpected $body") + } + + function.transformChildrenVoid(object : IrElementTransformerVoid() { + override fun visitCall(call: IrCall): IrExpression { + if (call !in tailCalls) return call + if (!function.returnType.isUnit() && call.type != function.returnType) return call + // Replace ARETURN with { POP, GETSTATIC kotlin/Unit.INSTANCE, ARETURN } during codegen later. + // Otherwise, additional CHECKCAST will break tail-call optimization + if (function.returnType.isUnit()) { + context.suspendTailCallsWithUnitReplacement += call.attributeOwnerId + } + return IrReturnImpl(call.startOffset, call.endOffset, context.irBuiltIns.nothingType, function.symbol, call) + } + }) + } +} + +// Find tail-call inside a single block. This function is needed, since there can be +// return statement in the middle of the function and thus we cannot just assume, that its last statement is tail-call +private fun List.findTailCall(functionReturnsUnit: Boolean): IrStatement? { + val mayBeReturn = find { it is IrReturn } as? IrReturn + return when (val value = mayBeReturn?.value) { + is IrGetField -> if (functionReturnsUnit && value.isGetFieldOfUnit()) { + // This is simple `return` in the middle of a function + // Tail-call should be just before it + subList(0, indexOf(mayBeReturn)).findTailCall(functionReturnsUnit) + } else mayBeReturn + null -> lastOrNull() + else -> mayBeReturn + } +} + +private fun IrGetField.isGetFieldOfUnit(): Boolean = + type.isUnit() && symbol.owner.name == Name.identifier("INSTANCE") && symbol.owner.parentAsClass.fqNameWhenAvailable == FqName("kotlin.Unit") diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/tailSuspendUnitFun.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/tailSuspendUnitFun.kt index 40c66fd9586..7f6a82e2d70 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/tailSuspendUnitFun.kt +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/tailSuspendUnitFun.kt @@ -1,5 +1,4 @@ // IGNORE_BACKEND_FIR: JVM_IR -// IGNORE_BACKEND: JVM_IR // TARGET_BACKEND: JVM // FULL_JDK // WITH_RUNTIME diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt index abb8c6f18a2..4cde20d7669 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt @@ -1,5 +1,4 @@ // IGNORE_BACKEND_FIR: JVM_IR -// IGNORE_BACKEND: JVM_IR // TARGET_BACKEND: JVM // FULL_JDK // WITH_RUNTIME