Loosen tail-call optimization check for functions returning Unit

Do not check, that all Unit predecessors are POPs. This is safe for the same reason, as it is safe to allow some of ARETURN sources not be suspension point results. To elaborate, before Unit, the stack is empty. This is because if there are multiple paths to Unit and at least one of them comes from POP after suspension point (we are interested in this case only - otherwise, the call is not tail-call), in path from said POP the stack is empty, since after suspension point the stack contains only one element. Thus, the stack in other paths leading to Unit has to be empty, otherwise, merge operation is not possible and ASM will report error during analysis. Since the stack is empty in all paths, we can hoist Unit and following ARETURN to predecessors, effectively turning path from suspension point to tail-call.
2026-03-10 08:31:29 +00:00 · 2021-08-19 05:51:09 +02:00
parent d33b70af1a
commit b5fa129540
8 changed files with 66 additions and 52 deletions
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt
+++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt
@@ -39,7 +39,6 @@ internal class MethodNodeExaminer(
    private val popsBeforeSafeUnitInstances = mutableSetOf<AbstractInsnNode>()
    private val areturnsAfterSafeUnitInstances = mutableSetOf<AbstractInsnNode>()
    private val meaningfulSuccessorsCache = hashMapOf<AbstractInsnNode, List<AbstractInsnNode>>()
-    private val meaningfulPredecessorsCache = hashMapOf<AbstractInsnNode, List<AbstractInsnNode>>()

    init {
        if (!disableTailCallOptimizationForFunctionReturningUnit) {
@@ -52,10 +51,8 @@ internal class MethodNodeExaminer(
            for (pop in popsBeforeUnitInstances) {
                val units = pop.meaningfulSuccessors()
                val allUnitsAreSafe = units.all { unit ->
-                    // check no other predecessor exists
-                    unit.meaningfulPredecessors().all { it in popsBeforeUnitInstances } &&
-                            // check they have only returns among successors
-                            unit.meaningfulSuccessors().all { it.opcode == Opcodes.ARETURN }
+                    // check they have only returns among successors
+                    unit.meaningfulSuccessors().all { it.opcode == Opcodes.ARETURN }
                }
                if (!allUnitsAreSafe) continue
                // save them all to the properties
@@ -74,35 +71,22 @@ internal class MethodNodeExaminer(
    private fun AbstractInsnNode.isAreturnAfterSafeUnitInstance(): Boolean = this in areturnsAfterSafeUnitInstances

    private fun AbstractInsnNode.meaningfulSuccessors(): List<AbstractInsnNode> = meaningfulSuccessorsCache.getOrPut(this) {
-        meaningfulSuccessorsOrPredecessors(true)
-    }
-
-    private fun AbstractInsnNode.meaningfulPredecessors(): List<AbstractInsnNode> = meaningfulPredecessorsCache.getOrPut(this) {
-        meaningfulSuccessorsOrPredecessors(false)
-    }
-
-    private fun AbstractInsnNode.meaningfulSuccessorsOrPredecessors(isSuccessors: Boolean): List<AbstractInsnNode> {
        fun AbstractInsnNode.isMeaningful() = isMeaningful && opcode != Opcodes.NOP && opcode != Opcodes.GOTO && this !is LineNumberNode

-        fun AbstractInsnNode.getIndices() =
-            if (isSuccessors) controlFlowGraph.getSuccessorsIndices(this)
-            else controlFlowGraph.getPredecessorsIndices(this)
-
        val visited = mutableSetOf<AbstractInsnNode>()
        fun dfs(insn: AbstractInsnNode) {
            if (insn in visited) return
            visited += insn
            if (!insn.isMeaningful()) {
-                for (succIndex in insn.getIndices()) {
+                for (succIndex in controlFlowGraph.getSuccessorsIndices(insn)) {
                    dfs(methodNode.instructions[succIndex])
                }
            }
        }
-
-        for (succIndex in getIndices()) {
+        for (succIndex in controlFlowGraph.getSuccessorsIndices(this)) {
            dfs(methodNode.instructions[succIndex])
        }
-        return visited.filter { it.isMeaningful() }
+        visited.filter { it.isMeaningful() }
    }

    fun replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks() {
--- 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
@@ -12144,6 +12144,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/functionReference.kt");
                }

+                @Test
+                @TestMetadata("inline.kt")
+                public void testInline() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/inline.kt");
+                }
+
                @Test
                @TestMetadata("override.kt")
                public void testOverride() throws Exception {
--- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline.txt
+++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline.txt
@@ -14,7 +14,6 @@ public final class CrossinlineKt$box$1$filter$$inlined$source$1$1 {
 }

@kotlin.Metadata
-@kotlin.coroutines.jvm.internal.DebugMetadata
 public final class CrossinlineKt$box$1$filter$$inlined$source$1$lambda$1$1 {
    enclosing method CrossinlineKt$box$1$filter$$inlined$source$1$lambda$1.send(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
    field label: int
@@ -96,19 +95,6 @@ public final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$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
-public final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2$1 {
-    enclosing method CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2.send(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
-    field label: int
-    synthetic field result: java.lang.Object
-    synthetic final field this$0: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2
-    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2
-    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2$1
-    public method <init>(p0: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2, p1: 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 CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2 {
    // source: 'crossinline.kt'
@@ -116,7 +102,6 @@ public final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2 {
    synthetic final field $this_source$inlined: Sink
    synthetic final field this$0: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1
    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2
-    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2$1
    public method <init>(p0: Sink, p1: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1): void
    public method close(@org.jetbrains.annotations.Nullable p0: java.lang.Throwable): void
    public @org.jetbrains.annotations.Nullable method send(p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
@@ -202,7 +187,6 @@ public final class CrossinlineKt$filter$$inlined$source$1$1 {
 }

@kotlin.Metadata
-@kotlin.coroutines.jvm.internal.DebugMetadata
 public final class CrossinlineKt$filter$$inlined$source$1$lambda$1$1 {
    enclosing method CrossinlineKt$filter$$inlined$source$1$lambda$1.send(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
    field label: int
--- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline_ir.txt
+++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline_ir.txt
@@ -12,26 +12,12 @@ public final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$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
-public final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2$1 {
-    enclosing method CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2.send(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
-    field label: int
-    synthetic field result: java.lang.Object
-    synthetic final field this$0: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2
-    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2
-    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2$1
-    public method <init>(p0: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2, p1: 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 CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2 {
    // source: 'crossinline.kt'
    enclosing method CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1.consume(LSink;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
    synthetic final field $this_source$inlined: Sink
    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2
-    inner (anonymous) class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$2$1
    public method <init>(p0: Sink): void
    public method close(@org.jetbrains.annotations.Nullable p0: java.lang.Throwable): void
    public @org.jetbrains.annotations.Nullable method send(p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
@@ -130,7 +116,6 @@ public final class CrossinlineKt$filter$$inlined$source$1 {
 }

@kotlin.Metadata
-@kotlin.coroutines.jvm.internal.DebugMetadata
 public final class CrossinlineKt$filter$lambda-3$$inlined$consumeEach$1$1 {
    enclosing method CrossinlineKt$filter$lambda-3$$inlined$consumeEach$1.send(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
    field label: int
--- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/inline.kt
+++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/inline.kt
@@ -0,0 +1,38 @@
+// TARGET_BACKEND: JVM
+// FULL_JDK
+// WITH_RUNTIME
+// WITH_COROUTINES
+// CHECK_TAIL_CALL_OPTIMIZATION
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+private var x = 100
+
+private inline fun inlineFun(
+    potentiallySuspendLambda: () -> Unit
+) {
+    if (x == 99) return
+    potentiallySuspendLambda()
+}
+
+suspend fun myFunWithTailCall() = inlineFun(
+    potentiallySuspendLambda = { suspendFun() }
+)
+
+suspend fun suspendFun(): Unit = 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"
+}
--- 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
@@ -12042,6 +12042,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/functionReference.kt");
                }

+                @Test
+                @TestMetadata("inline.kt")
+                public void testInline() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/inline.kt");
+                }
+
                @Test
                @TestMetadata("override.kt")
                public void testOverride() throws Exception {
--- 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
@@ -12144,6 +12144,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/functionReference.kt");
                }

+                @Test
+                @TestMetadata("inline.kt")
+                public void testInline() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/inline.kt");
+                }
+
                @Test
                @TestMetadata("override.kt")
                public void testOverride() throws Exception {
--- a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java
+++ b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java
@@ -9683,6 +9683,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/functionReference.kt");
                }

+                @TestMetadata("inline.kt")
+                public void testInline() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/inline.kt");
+                }
+
                @TestMetadata("override.kt")
                public void testOverride() throws Exception {
                    runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/override.kt");