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
2022-01-24 23:12:08 +01:00
parent e05f189082
commit da80ac008b
10 changed files with 104 additions and 42 deletions
@@ -103,6 +103,7 @@ class CoroutineTransformerMethodVisitor(
            )
            if (examiner.allSuspensionPointsAreTailCalls(suspensionPoints)) {
                examiner.replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks()
+                examiner.addCoroutineSuspendedChecksBeforeSafeCheckcasts()
                dropSuspensionMarkers(methodNode)
                dropUnboxInlineClassMarkers(methodNode, suspensionPoints)
                return
@@ -37,6 +37,10 @@ internal class MethodNodeExaminer(
    private val areturnsAfterSafeUnitInstances = mutableSetOf<AbstractInsnNode>()
    private val meaningfulSuccessorsCache = hashMapOf<AbstractInsnNode, List<AbstractInsnNode>>()

+    // 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<AbstractInsnNode>()
+
    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<SuspensionPoint>): 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<SuspensionPoint>

    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)
@@ -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 {
@@ -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"
+}
@@ -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")
@@ -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 <init>(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
@@ -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 <init>(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
@@ -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 {
@@ -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 {
@@ -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");