From da80ac008bf5156f34ed4ef86399812a715f6cb0 Mon Sep 17 00:00:00 2001 From: Ilmir Usmanov Date: Mon, 24 Jan 2022 23:12:08 +0100 Subject: [PATCH] Make CHECKCAST not break TCO So, treat CHECKCASTs as { POP, Unit } sequences. If the CHECKCAST is between suspension point and ARETURN, put check for COROUTINE_SUSPENDED before it and return if the suspension point is suspended. This is safe, since if the function throws CCE, it will be thrown from the last state in state-machine and we cannot reenter the function during resume. So, in case of CHECKCAST throwing CCE the behavior is the same, whether we have state-machine or not. We do not need to disable TCO in some cases, as we do for functions, returning Unit, since in latter case suspend function, returning Unit might appear as returning non-Unit during resumption due to missing { POP, Unit } sequence, which is not executed, since the function is tail-call. However, in case of functions, returning non-Unit there is no such concern, since we do not POP result of suspension point, but rather, return it after CHECKCAST. #KT-50835 Fixed --- .../CoroutineTransformerMethodVisitor.kt | 1 + .../coroutines/TailCallOptimization.kt | 49 ++++++++++++++++--- .../FirBlackBoxCodegenTestGenerated.java | 6 +++ .../tailCallOptimizations/checkcast.kt | 38 ++++++++++++++ .../inlineWithoutStateMachine.kt | 5 -- .../inlineWithoutStateMachine.txt | 15 ------ .../inlineWithoutStateMachine_ir.txt | 15 ------ .../codegen/BlackBoxCodegenTestGenerated.java | 6 +++ .../IrBlackBoxCodegenTestGenerated.java | 6 +++ .../LightAnalysisModeTestGenerated.java | 5 ++ 10 files changed, 104 insertions(+), 42 deletions(-) create mode 100644 compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt index 4671f555bb0..8f00ce1fa87 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt @@ -103,6 +103,7 @@ class CoroutineTransformerMethodVisitor( ) if (examiner.allSuspensionPointsAreTailCalls(suspensionPoints)) { examiner.replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks() + examiner.addCoroutineSuspendedChecksBeforeSafeCheckcasts() dropSuspensionMarkers(methodNode) dropUnboxInlineClassMarkers(methodNode, suspensionPoints) return diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt index 7b640441ec9..64fddcc7fd1 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt @@ -37,6 +37,10 @@ internal class MethodNodeExaminer( private val areturnsAfterSafeUnitInstances = mutableSetOf() private val meaningfulSuccessorsCache = hashMapOf>() + // CHECKCAST is considered safe if it is right before ARETURN and right after suspension point + // In this case, we can add check for COROUTINE_SUSPENDED, the same as we did for functions, returning Unit. + private val safeCheckcasts = mutableSetOf() + init { if (!disableTailCallOptimizationForFunctionReturningUnit) { // retrieve all POP insns @@ -58,6 +62,28 @@ internal class MethodNodeExaminer( units.flatMapTo(areturnsAfterSafeUnitInstances) { it.meaningfulSuccessors() } } } + + fun AbstractInsnNode.isPartOfCheckcastChainBeforeAreturn(): Boolean { + for (succ in meaningfulSuccessors()) { + when (succ.opcode) { + Opcodes.CHECKCAST -> + if (!succ.isPartOfCheckcastChainBeforeAreturn()) return false + + Opcodes.ARETURN -> { + // do nothing + } + else -> return false + } + } + return true + } + + val checkcasts = methodNode.instructions.filter { it.opcode == Opcodes.CHECKCAST } + for (checkcast in checkcasts) { + if (!checkcast.isPartOfCheckcastChainBeforeAreturn()) continue + if (frames[methodNode.instructions.indexOf(checkcast)]?.top() !is FromSuspensionPointValue) continue + safeCheckcasts += checkcast + } } // GETSTATIC kotlin/Unit.INSTANCE is considered safe iff @@ -104,6 +130,19 @@ internal class MethodNodeExaminer( } } + fun addCoroutineSuspendedChecksBeforeSafeCheckcasts() { + for (checkcast in safeCheckcasts) { + val label = Label() + methodNode.instructions.insertBefore(checkcast, withInstructionAdapter { + dup() + loadCoroutineSuspendedMarker() + ifacmpne(label) + areturn(AsmTypes.OBJECT_TYPE) + mark(label) + }) + } + } + fun allSuspensionPointsAreTailCalls(suspensionPoints: List): Boolean { val safelyReachableReturns = findSafelyReachableReturns() @@ -152,7 +191,7 @@ internal class MethodNodeExaminer( } if (!insn.isMeaningful || insn.opcode in SAFE_OPCODES || insn.isInvisibleInDebugVarInsn(methodNode) || isInlineMarker(insn) - || insn.isSafeUnitInstance() || insn.isAreturnAfterSafeUnitInstance() || insn.isCheckcastObject() + || insn.isSafeUnitInstance() || insn.isAreturnAfterSafeUnitInstance() ) { setOf() } else null @@ -184,9 +223,6 @@ internal class MethodNodeExaminer( } } -private fun AbstractInsnNode.isCheckcastObject(): Boolean = - opcode == Opcodes.CHECKCAST && (this as TypeInsnNode).desc == AsmTypes.OBJECT_TYPE.internalName - private fun AbstractInsnNode?.isInvisibleInDebugVarInsn(methodNode: MethodNode): Boolean { val insns = methodNode.instructions val index = insns.indexOf(this) @@ -196,7 +232,7 @@ private fun AbstractInsnNode?.isInvisibleInDebugVarInsn(methodNode: MethodNode): } private val SAFE_OPCODES = - ((Opcodes.DUP..Opcodes.DUP2_X2) + Opcodes.NOP + Opcodes.POP + Opcodes.POP2 + (Opcodes.IFEQ..Opcodes.GOTO)).toSet() + ((Opcodes.DUP..Opcodes.DUP2_X2) + Opcodes.NOP + Opcodes.POP + Opcodes.POP2 + (Opcodes.IFEQ..Opcodes.GOTO)).toSet() + Opcodes.CHECKCAST private object FromSuspensionPointValue : BasicValue(AsmTypes.OBJECT_TYPE) { override fun equals(other: Any?): Boolean = other is FromSuspensionPointValue @@ -226,8 +262,7 @@ private class TcoInterpreter(private val suspensionPoints: List override fun unaryOperation(insn: AbstractInsnNode, value: BasicValue?): BasicValue? { // Assume, that CHECKCAST Object does not break tail-call optimization - // TODO: Investigate, whether any CHECKCAST is safe in terms of tail-call optimization - if (value is FromSuspensionPointValue && insn.isCheckcastObject()) { + if (value is FromSuspensionPointValue && insn.opcode == Opcodes.CHECKCAST) { return value } return super.unaryOperation(insn, value).convert(insn) diff --git a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java index bcc7d4e29ec..8787a4f9fa1 100644 --- a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java +++ b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java @@ -12650,6 +12650,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/also.kt"); } + @Test + @TestMetadata("checkcast.kt") + public void testCheckcast() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); + } + @Test @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt new file mode 100644 index 00000000000..3c3da2fa97a --- /dev/null +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt @@ -0,0 +1,38 @@ +// TARGET_BACKEND: JVM +// FULL_JDK +// WITH_STDLIB +// WITH_COROUTINES +// CHECK_TAIL_CALL_OPTIMIZATION +import helpers.* +import kotlin.coroutines.* +import kotlin.coroutines.intrinsics.* + +private var x = 100 + +private inline fun inlineFun( + potentiallySuspendLambda: () -> String +): String { + if (x == 99) return "" + return potentiallySuspendLambda() +} + +suspend fun myFunWithTailCall() = inlineFun( + potentiallySuspendLambda = { suspendFun() } +) + +suspend fun suspendFun(): String = suspendCoroutineUninterceptedOrReturn { x -> + TailCallOptimizationChecker.saveStackTrace(x) + COROUTINE_SUSPENDED +} + +fun builder(c: suspend () -> Unit) { + c.startCoroutine(EmptyContinuation) +} + +fun box(): String { + builder { + myFunWithTailCall() + } + TailCallOptimizationChecker.checkNoStateMachineIn("myFunWithTailCall") + return "OK" +} diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt index 1c216cbdb83..9ab77b139a2 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt @@ -13,13 +13,8 @@ inline suspend fun suspendThere(v: String): String = suspendCoroutineUnintercept COROUTINE_SUSPENDED } -// TODO: Somehow we still generate continuations for tail-call function, but we don't use them. suspend fun suspendHere(): String = suspendThere("O") -// There is a kind of redundant state machine generated for complexSuspend: -// it's basically has the only suspend call just before return, but there is -// a redundant CHECKCAST String in the run's lambda, so we have to insert the state machine. -// TODO: Think of avoiding such redundant casts suspend fun complexSuspend(): String { return run { suspendThere("K") diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.txt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.txt index 616932fe0a5..20bcbdfcb65 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.txt +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.txt @@ -14,25 +14,10 @@ final class InlineWithoutStateMachineKt$box$1 { public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object } -@kotlin.Metadata -@kotlin.coroutines.jvm.internal.DebugMetadata -final class InlineWithoutStateMachineKt$complexSuspend$1 { - // source: 'inlineWithoutStateMachine.kt' - enclosing method InlineWithoutStateMachineKt.complexSuspend(Lkotlin/coroutines/Continuation;)Ljava/lang/Object; - field L$0: java.lang.Object - field L$1: java.lang.Object - field label: int - synthetic field result: java.lang.Object - inner (anonymous) class InlineWithoutStateMachineKt$complexSuspend$1 - method (p0: kotlin.coroutines.Continuation): void - public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object -} - @kotlin.Metadata public final class InlineWithoutStateMachineKt { // source: 'inlineWithoutStateMachine.kt' inner (anonymous) class InlineWithoutStateMachineKt$box$1 - inner (anonymous) class InlineWithoutStateMachineKt$complexSuspend$1 public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void public final static @org.jetbrains.annotations.Nullable method complexSuspend(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt index 2f4fc2f0893..610d4ca8533 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt @@ -15,25 +15,10 @@ final class InlineWithoutStateMachineKt$box$1 { public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object } -@kotlin.coroutines.jvm.internal.DebugMetadata -@kotlin.Metadata -final class InlineWithoutStateMachineKt$complexSuspend$1 { - // source: 'inlineWithoutStateMachine.kt' - enclosing method InlineWithoutStateMachineKt.complexSuspend(Lkotlin/coroutines/Continuation;)Ljava/lang/Object; - field L$0: java.lang.Object - field L$1: java.lang.Object - field label: int - synthetic field result: java.lang.Object - inner (anonymous) class InlineWithoutStateMachineKt$complexSuspend$1 - method (p0: kotlin.coroutines.Continuation): void - public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object -} - @kotlin.Metadata public final class InlineWithoutStateMachineKt { // source: 'inlineWithoutStateMachine.kt' inner (anonymous) class InlineWithoutStateMachineKt$box$1 - inner (anonymous) class InlineWithoutStateMachineKt$complexSuspend$1 public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void public final static @org.jetbrains.annotations.Nullable method complexSuspend(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java index 867cda87421..4fc25bb1c85 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java @@ -12530,6 +12530,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest { runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/also.kt"); } + @Test + @TestMetadata("checkcast.kt") + public void testCheckcast() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); + } + @Test @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java index f0b6e1dc6f2..0e16bae6f74 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java @@ -12650,6 +12650,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/also.kt"); } + @Test + @TestMetadata("checkcast.kt") + public void testCheckcast() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); + } + @Test @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { diff --git a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java index 609e9355ac3..c3e6c196551 100644 --- a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java +++ b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java @@ -10099,6 +10099,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/also.kt"); } + @TestMetadata("checkcast.kt") + public void testCheckcast() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); + } + @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline.kt");