JVM: streamline handling of tail-call suspend functions

Just see if every suspend call is followed only by safe instructions or
returns, then insert a suspension point check if there isn't a direct
return.

The first part of this is equivalent to the old implementation, just
refactored. The second part generates strictly more checks; see, for
example, the fixed test, in which the previous implementation failed to
insert a check before a CHECKCAST.

^KT-51818 Fixed

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

@@ -96,15 +96,8 @@ class CoroutineTransformerMethodVisitor(
                 addCompletionParameterToLVT(methodNode)
             }
 
-            val examiner = MethodNodeExaminer(
+            if (methodNode.allSuspensionPointsAreTailCalls(suspensionPoints, !disableTailCallOptimizationForFunctionReturningUnit)) {
-                containingClassInternalName,
+                methodNode.addCoroutineSuspendedChecks(suspensionPoints)
-                methodNode,
-                suspensionPoints,
-                disableTailCallOptimizationForFunctionReturningUnit
-            )
-            if (examiner.allSuspensionPointsAreTailCalls(suspensionPoints)) {
-                examiner.replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks()
-                examiner.addCoroutineSuspendedChecksBeforeSafeCheckcasts()
                 dropSuspensionMarkers(methodNode)
                 dropUnboxInlineClassMarkers(methodNode, suspensionPoints)
                 return

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

@@ -8,220 +8,90 @@ package org.jetbrains.kotlin.codegen.coroutines
 import org.jetbrains.kotlin.codegen.inline.isInlineMarker
 import org.jetbrains.kotlin.codegen.optimization.boxing.isUnitInstance
 import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph
-import org.jetbrains.kotlin.codegen.optimization.common.asSequence
 import org.jetbrains.kotlin.codegen.optimization.common.isMeaningful
 import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
-import org.jetbrains.kotlin.utils.sure
+import org.jetbrains.kotlin.utils.addToStdlib.popLast
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.*
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicInterpreter
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
-import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
 
-internal class MethodNodeExaminer(
+internal fun MethodNode.allSuspensionPointsAreTailCalls(suspensionPoints: List<SuspensionPoint>, optimizeReturnUnit: Boolean): Boolean {
-    containingClassInternalName: String,
+    val frames = MethodTransformer.analyze("fake", this, TcoInterpreter(suspensionPoints))
-    val methodNode: MethodNode,
+    val controlFlowGraph = ControlFlowGraph.build(this)
-    suspensionPoints: List<SuspensionPoint>,
-    disableTailCallOptimizationForFunctionReturningUnit: Boolean
-) {
-    private val frames: Array<Frame<BasicValue>?> =
-        MethodTransformer.analyze(containingClassInternalName, methodNode, TcoInterpreter(suspensionPoints))
-    private val controlFlowGraph = ControlFlowGraph.build(methodNode)
 
-    private val safeUnitInstances = mutableSetOf<AbstractInsnNode>()
+    fun AbstractInsnNode.isSafe(): Boolean =
-    private val popsBeforeSafeUnitInstances = mutableSetOf<AbstractInsnNode>()
+        !isMeaningful || opcode in SAFE_OPCODES || isInvisibleInDebugVarInsn(this@allSuspensionPointsAreTailCalls) || isInlineMarker(this)
-    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
+    fun AbstractInsnNode.transitiveSuccessorsAreSafeOrReturns(): Boolean {
-    // In this case, we can add check for COROUTINE_SUSPENDED, the same as we did for functions, returning Unit.
+        val visited = mutableSetOf(this)
-    private val safeCheckcasts = mutableSetOf<AbstractInsnNode>()
+        val stack = mutableListOf(this)
-
+        while (stack.isNotEmpty()) {
-    init {
+            val insn = stack.popLast()
-        if (!disableTailCallOptimizationForFunctionReturningUnit) {
+            // In Unit-returning functions, the last statement is followed by POP + GETSTATIC Unit.INSTANCE
-            // retrieve all POP insns
+            // if it is itself not Unit-returning.
-            val pops = methodNode.instructions.asSequence().filter { it.opcode == Opcodes.POP }
+            if (insn.opcode == Opcodes.ARETURN || (optimizeReturnUnit && insn.isPopBeforeReturnUnit)) {
-            // for each of them check that all successors are PUSH Unit
+                if (frames[instructions.indexOf(insn)]?.top() !is FromSuspensionPointValue?) {
-            val popsBeforeUnitInstances = pops.map { it to it.meaningfulSuccessors() }
+                    return false
-                .filter { (_, succs) -> succs.all { it.isUnitInstance() } }
-                .map { it.first }.toList()
-            for (pop in popsBeforeUnitInstances) {
-                val units = pop.meaningfulSuccessors()
-                val allUnitsAreSafe = units.all { unit ->
-                    // check they have only returns among successors
-                    unit.meaningfulSuccessors().all { it.opcode == Opcodes.ARETURN }
                 }
-                if (!allUnitsAreSafe) continue
+            } else if (insn !== this && !insn.isSafe()) {
-                // save them all to the properties
+                return false
-                popsBeforeSafeUnitInstances += pop
+            } else {
-                safeUnitInstances += units
+                for (nextIndex in controlFlowGraph.getSuccessorsIndices(insn)) {
-                units.flatMapTo(areturnsAfterSafeUnitInstances) { it.meaningfulSuccessors() }
+                    val nextInsn = instructions.get(nextIndex)
-            }
+                    if (visited.add(nextInsn)) {
-        }
+                        stack.add(nextInsn)
-
-        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
-    // it is part of POP, PUSH Unit, ARETURN sequence.
-    private fun AbstractInsnNode.isSafeUnitInstance(): Boolean = this in safeUnitInstances
-
-    private fun AbstractInsnNode.isPopBeforeSafeUnitInstance(): Boolean = this in popsBeforeSafeUnitInstances
-    private fun AbstractInsnNode.isAreturnAfterSafeUnitInstance(): Boolean = this in areturnsAfterSafeUnitInstances
-
-    private fun AbstractInsnNode.meaningfulSuccessors(): List<AbstractInsnNode> = meaningfulSuccessorsCache.getOrPut(this) {
-        fun AbstractInsnNode.isMeaningful() = isMeaningful && opcode != Opcodes.NOP && opcode != Opcodes.GOTO && this !is LineNumberNode
-
-        val visited = mutableSetOf<AbstractInsnNode>()
-        fun dfs(insn: AbstractInsnNode) {
-            if (insn in visited) return
-            visited += insn
-            if (!insn.isMeaningful()) {
-                for (succIndex in controlFlowGraph.getSuccessorsIndices(insn)) {
-                    dfs(methodNode.instructions[succIndex])
                 }
             }
         }
-        for (succIndex in controlFlowGraph.getSuccessorsIndices(this)) {
+        return true
-            dfs(methodNode.instructions[succIndex])
-        }
-        visited.filter { it.isMeaningful() }
     }
 
-    fun replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks() {
+    return suspensionPoints.all { suspensionPoint ->
-        val isReferenceMap =
+        val index = instructions.indexOf(suspensionPoint.suspensionCallBegin)
-            popsBeforeSafeUnitInstances.associateWith { (frames[methodNode.instructions.indexOf(it)]?.top()?.isReference == true) }
+        tryCatchBlocks.all { index < instructions.indexOf(it.start) || instructions.indexOf(it.end) <= index } &&
-
+                suspensionPoint.suspensionCallEnd.transitiveSuccessorsAreSafeOrReturns()
-        for (pop in popsBeforeSafeUnitInstances) {
-            if (isReferenceMap[pop] == true) {
-                val label = Label()
-                methodNode.instructions.insertBefore(pop, withInstructionAdapter {
-                    dup()
-                    loadCoroutineSuspendedMarker()
-                    ifacmpne(label)
-                    areturn(AsmTypes.OBJECT_TYPE)
-                    mark(label)
-                })
-            }
-        }
-    }
-
-    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()
-
-        val instructions = methodNode.instructions
-        return suspensionPoints.all { suspensionPoint ->
-            val beginIndex = instructions.indexOf(suspensionPoint.suspensionCallBegin)
-            val endIndex = instructions.indexOf(suspensionPoint.suspensionCallEnd)
-
-            val insideTryBlock = methodNode.tryCatchBlocks.any { block ->
-                val tryBlockStartIndex = instructions.indexOf(block.start)
-                val tryBlockEndIndex = instructions.indexOf(block.end)
-
-                beginIndex in tryBlockStartIndex until tryBlockEndIndex
-            }
-            if (insideTryBlock) return@all false
-
-            safelyReachableReturns[endIndex + 1]?.all { returnIndex ->
-                frames[returnIndex]?.top().sure {
-                    "There must be some value on stack to return"
-                } is FromSuspensionPointValue
-            } ?: false
-        }
-    }
-
-    /**
-     * Let's call an instruction safe if its execution is always invisible: stack modifications, branching, variable insns (invisible in debug)
-     *
-     * For some instruction `insn` define the result as following:
-     * - if there is a path leading to the non-safe instruction then result is `null`
-     * - Otherwise result contains all the reachable ARETURN indices
-     *
-     * @return indices of safely reachable returns for each instruction in the method node
-     */
-    private fun findSafelyReachableReturns(): Array<Set<Int>?> {
-        val insns = methodNode.instructions.toArray()
-        val reachableReturnsIndices = Array(insns.size) init@{ index ->
-            val insn = insns[index]
-
-            if (insn.opcode == Opcodes.ARETURN && !insn.isAreturnAfterSafeUnitInstance()) {
-                return@init setOf(index)
-            }
-
-            // Since POP, PUSH Unit, ARETURN behaves like normal return in terms of tail-call optimization, set return index to POP
-            if (insn.isPopBeforeSafeUnitInstance()) {
-                return@init setOf(index)
-            }
-
-            if (!insn.isMeaningful || insn.opcode in SAFE_OPCODES || insn.isInvisibleInDebugVarInsn(methodNode) || isInlineMarker(insn)
-                || insn.isSafeUnitInstance() || insn.isAreturnAfterSafeUnitInstance()
-            ) {
-                setOf()
-            } else null
-        }
-
-        var changed: Boolean
-        do {
-            changed = false
-            for (index in insns.indices.reversed()) {
-                if (insns[index].opcode == Opcodes.ARETURN) continue
-
-                val newResult =
-                    controlFlowGraph
-                        .getSuccessorsIndices(index).plus(index)
-                        .map(reachableReturnsIndices::get)
-                        .fold<Set<Int>?, Set<Int>?>(mutableSetOf()) { acc, successorsResult ->
-                            if (acc != null && successorsResult != null) acc + successorsResult else null
-                        }
-
-                if (newResult != reachableReturnsIndices[index]) {
-                    reachableReturnsIndices[index] = newResult
-                    changed = true
-                }
-            }
-        } while (changed)
-
-        return reachableReturnsIndices
     }
 }
 
+internal fun MethodNode.addCoroutineSuspendedChecks(suspensionPoints: List<SuspensionPoint>) {
+    for (suspensionPoint in suspensionPoints) {
+        if (suspensionPoint.suspensionCallEnd.nextMeaningful?.opcode == Opcodes.ARETURN) {
+            // `if (x == COROUTINE_SUSPENDED) return x else return x` == `return x`
+            continue
+        }
+        instructions.insert(suspensionPoint.suspensionCallEnd, withInstructionAdapter {
+            val label = Label()
+            dup()
+            loadCoroutineSuspendedMarker()
+            ifacmpne(label)
+            areturn(AsmTypes.OBJECT_TYPE)
+            mark(label)
+        })
+    }
+}
+
+private tailrec fun AbstractInsnNode?.skipUntilMeaningful(): AbstractInsnNode? = when {
+    this == null -> null
+    opcode == Opcodes.NOP || !isMeaningful -> next.skipUntilMeaningful()
+    opcode == Opcodes.GOTO -> (this as JumpInsnNode).label.skipUntilMeaningful()
+    else -> this
+}
+
+private val AbstractInsnNode.nextMeaningful: AbstractInsnNode?
+    get() = next.skipUntilMeaningful()
+
+private val AbstractInsnNode.isReturnUnit: Boolean
+    get() = isUnitInstance() && nextMeaningful?.opcode == Opcodes.ARETURN
+
+private val AbstractInsnNode.isPopBeforeReturnUnit: Boolean
+    get() = opcode == Opcodes.POP && nextMeaningful?.isReturnUnit == true
+
 private fun AbstractInsnNode?.isInvisibleInDebugVarInsn(methodNode: MethodNode): Boolean {
     val insns = methodNode.instructions
     val index = insns.indexOf(this)
@@ -230,8 +100,15 @@ private fun AbstractInsnNode?.isInvisibleInDebugVarInsn(methodNode: MethodNode):
     })
 }
 
-private val SAFE_OPCODES =
+private val SAFE_OPCODES = buildSet {
-    ((Opcodes.DUP..Opcodes.DUP2_X2) + Opcodes.NOP + Opcodes.POP + Opcodes.POP2 + (Opcodes.IFEQ..Opcodes.GOTO)).toSet() + Opcodes.CHECKCAST
+    add(Opcodes.NOP)
+    addAll(Opcodes.POP..Opcodes.SWAP) // POP*, DUP*, SWAP
+    addAll(Opcodes.IFEQ..Opcodes.GOTO) // IF*, GOTO
+    // CHECKCAST is technically not safe (can throw), but should be unless the type system is holey.
+    // Treating it as safe permits optimizing functions where a non-Any returning suspend function
+    // call is in a tail position (in bytecode they all return Object, so a cast is sometimes inserted).
+    add(Opcodes.CHECKCAST)
+}
 
 private object FromSuspensionPointValue : BasicValue(AsmTypes.OBJECT_TYPE) {
     override fun equals(other: Any?): Boolean = other is FromSuspensionPointValue

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -12770,6 +12770,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt");
             }
 
+            @Test
+            @TestMetadata("checkcast2.kt")
+            public void testCheckcast2() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt");
+            }
+
             @Test
             @TestMetadata("crossinline.kt")
             public void testCrossinline() throws Exception {

compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt

@@ -0,0 +1,29 @@
+// TARGET_BACKEND: JVM
+// FULL_JDK
+// WITH_STDLIB
+// WITH_COROUTINES
+// CHECK_TAIL_CALL_OPTIMIZATION
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+suspend fun suspendFun(x: String): String = suspendCoroutineUninterceptedOrReturn {
+    TailCallOptimizationChecker.saveStackTrace(it)
+    COROUTINE_SUSPENDED
+}
+
+suspend fun myFunWithTailCall(x: String) {
+    x.let { suspendFun(it) }
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    builder {
+        myFunWithTailCall("...")
+    }
+    TailCallOptimizationChecker.checkNoStateMachineIn("myFunWithTailCall")
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -12644,6 +12644,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt");
             }
 
+            @Test
+            @TestMetadata("checkcast2.kt")
+            public void testCheckcast2() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt");
+            }
+
             @Test
             @TestMetadata("crossinline.kt")
             public void testCrossinline() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -12770,6 +12770,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt");
             }
 
+            @Test
+            @TestMetadata("checkcast2.kt")
+            public void testCheckcast2() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt");
+            }
+
             @Test
             @TestMetadata("crossinline.kt")
             public void testCrossinline() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -10203,6 +10203,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt");
             }
 
+            @TestMetadata("checkcast2.kt")
+            public void testCheckcast2() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt");
+            }
+
             @TestMetadata("crossinline.kt")
             public void testCrossinline() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline.kt");