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

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

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt

@@ -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)