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