JVM_IR: apply TailCallOptimizationLowering to all suspend functions

Even if a function is known to be tail call because it's a compiler
generated bridge, the tail return might still need to be added in case
of Unit return type.

compiler/fir/fir2ir/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java

@@ -7823,6 +7823,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/coroutineReturn.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt");
+            }
+
             @TestMetadata("suspendNonLocalReturn.kt")
             public void testSuspendNonLocalReturn_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendNonLocalReturn.kt", "kotlin.coroutines");

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/TailCallOptimizationLowering.kt

@@ -9,10 +9,9 @@ import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.codegen.hasContinuation
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
@@ -34,8 +33,8 @@ internal val tailCallOptimizationPhase = makeIrFilePhase(
 private class TailCallOptimizationLowering(private val context: JvmBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
         irFile.transformChildren(object : IrElementTransformer<TailCallOptimizationData?> {
-            override fun visitFunction(declaration: IrFunction, data: TailCallOptimizationData?) =
+            override fun visitSimpleFunction(declaration: IrSimpleFunction, data: TailCallOptimizationData?) =
-                super.visitFunction(declaration, TailCallOptimizationData(declaration))
+                super.visitSimpleFunction(declaration, if (declaration.isSuspend) TailCallOptimizationData(declaration) else null)
 
             override fun visitCall(expression: IrCall, data: TailCallOptimizationData?): IrExpression {
                 val transformed = super.visitCall(expression, data) as IrExpression
@@ -52,33 +51,31 @@ private class TailCallOptimizationLowering(private val context: JvmBackendContex
     )
 }
 
-private class TailCallOptimizationData(val function: IrFunction) {
+private class TailCallOptimizationData(val function: IrSimpleFunction) {
     val returnsUnit = function.returnType.isUnit()
     val tailCalls = mutableSetOf<IrCall>()
 
     // Collect all tail calls, including those nested in `when`s, which are not arguments to `return`s.
-    private fun findCallsOnTailPositionWithoutImmediateReturn(statement: IrStatement, immediateReturn: Boolean = false) {
+    private fun IrStatement.findCallsOnTailPositionWithoutImmediateReturn(immediateReturn: Boolean = false) {
         when {
-            statement is IrCall && statement.isSuspend && !immediateReturn && (returnsUnit || statement.type == function.returnType) ->
+            this is IrCall && isSuspend && !immediateReturn && (returnsUnit || type == function.returnType) ->
-                tailCalls += statement
+                tailCalls += this
-            statement is IrBlock ->
+            this is IrBlock ->
-                statement.statements.findTailCall(returnsUnit)?.let(::findCallsOnTailPositionWithoutImmediateReturn)
+                statements.findTailCall(returnsUnit)?.findCallsOnTailPositionWithoutImmediateReturn()
-            statement is IrWhen ->
+            this is IrWhen ->
-                statement.branches.forEach { findCallsOnTailPositionWithoutImmediateReturn(it.result) }
+                branches.forEach { it.result.findCallsOnTailPositionWithoutImmediateReturn() }
-            statement is IrReturn ->
+            this is IrReturn ->
-                findCallsOnTailPositionWithoutImmediateReturn(statement.value, immediateReturn = true)
+                value.findCallsOnTailPositionWithoutImmediateReturn(immediateReturn = true)
-            statement is IrTypeOperatorCall && statement.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT ->
+            this is IrTypeOperatorCall && operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT ->
-                findCallsOnTailPositionWithoutImmediateReturn(statement.argument)
+                argument.findCallsOnTailPositionWithoutImmediateReturn()
             // TODO: Support binary logical operations and elvis, though. KT-23826 and KT-23825
         }
     }
 
     init {
-        if (function.hasContinuation()) {
+        when (val body = function.body) {
-            when (val body = function.body) {
+            is IrBlockBody -> body.statements.findTailCall(returnsUnit)?.findCallsOnTailPositionWithoutImmediateReturn()
-                is IrBlockBody -> body.statements.findTailCall(returnsUnit)?.let(::findCallsOnTailPositionWithoutImmediateReturn)
+            is IrExpressionBody -> body.expression.findCallsOnTailPositionWithoutImmediateReturn()
-                is IrExpressionBody -> findCallsOnTailPositionWithoutImmediateReturn(body.expression)
-            }
         }
     }
 }

compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt

@@ -0,0 +1,34 @@
+// IGNORE_BACKEND_FIR: JVM_IR
+// WITH_RUNTIME
+// WITH_COROUTINES
+
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+suspend fun pause(): Unit = suspendCoroutineUninterceptedOrReturn { x ->
+    x.resume(Unit)
+    COROUTINE_SUSPENDED
+}
+
+interface I {
+    suspend fun f()
+}
+
+var unpaused = false
+
+class A : I {
+    override suspend fun f() {
+        pause()
+        unpaused = true
+    }
+}
+
+class B(val x: I) : I by x
+
+fun box(): String {
+    suspend {
+        B(A()).f()
+    }.startCoroutine(EmptyContinuation)
+    return if (unpaused) "OK" else "fail: ignored suspension"
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -8943,6 +8943,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/coroutineReturn.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt");
+            }
+
             @TestMetadata("suspendNonLocalReturn.kt")
             public void testSuspendNonLocalReturn_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendNonLocalReturn.kt", "kotlin.coroutines.experimental");

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

@@ -8943,6 +8943,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/coroutineReturn.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt");
+            }
+
             @TestMetadata("suspendNonLocalReturn.kt")
             public void testSuspendNonLocalReturn_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendNonLocalReturn.kt", "kotlin.coroutines.experimental");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -7823,6 +7823,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/coroutineReturn.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt");
+            }
+
             @TestMetadata("suspendNonLocalReturn.kt")
             public void testSuspendNonLocalReturn_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendNonLocalReturn.kt", "kotlin.coroutines");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -6583,6 +6583,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/coroutineReturn.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt");
+            }
+
             @TestMetadata("suspendNonLocalReturn.kt")
             public void testSuspendNonLocalReturn_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendNonLocalReturn.kt", "kotlin.coroutines");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -6583,6 +6583,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/coroutineReturn.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/unitTypeReturn/interfaceDelegation.kt");
+            }
+
             @TestMetadata("suspendNonLocalReturn.kt")
             public void testSuspendNonLocalReturn_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendNonLocalReturn.kt", "kotlin.coroutines");