Fix infinite loop in suspend function in case of return from finally [#KT-21961]

* add test

compiler/testData/codegen/box/coroutines/controlFlow/returnFromFinally.kt

@@ -7,6 +7,8 @@ import helpers.*
 import COROUTINES_PACKAGE.*
 import COROUTINES_PACKAGE.intrinsics.*
 
+
+
 class Controller {
     var result = ""
 
@@ -15,9 +17,18 @@ class Controller {
         c.resume(value)
         COROUTINE_SUSPENDED
     }
+
+    var count = 0
+    fun expect(i: Int) {
+        if (++count != i) throw Exception("EXPECTED $i")
+    }
+
+    fun <T> log(value: T) {
+        result += "log($value);"
+    }
 }
 
-fun builder(c: suspend Controller.() -> String): String {
+fun builder(c: suspend Controller.() -> Int): String {
     val controller = Controller()
     c.startCoroutine(controller, handleResultContinuation {
         controller.result += "return($it);"
@@ -25,45 +36,22 @@ fun builder(c: suspend Controller.() -> String): String {
     return controller.result
 }
 
-fun builderUnit(c: suspend Controller.() -> Unit): String {
-    val controller = Controller()
-    c.startCoroutine(controller, handleResultContinuation {
-        controller.result += "return;"
-    })
-    return controller.result
-}
-
-fun <T> id(value: T) = value
-
 fun box(): String {
-    var value = builder {
+    val res = builder {
+        expect(1)
+        log(1)
         try {
-            if (id(23) == 23) {
+            expect(2)
-                return@builder suspendAndLog("OK")
+            suspendAndLog(2)
-            }
+        } finally {
+            expect(3)
+            log(3)
+            return@builder 4
         }
-        finally {
+        log("FAIL")
-            result += "finally;"
+        -1
-        }
-        result += "afterFinally;"
-        "shouldNotReach"
     }
-    if (value != "suspend(OK);finally;return(OK);") return "fail1: $value"
-
-    value = builderUnit {
-        try {
-            if (id(23) == 23) {
-                suspendAndLog("OK")
-                return@builderUnit
-            }
-        }
-        finally {
-            result += "finally;"
-        }
-        result += "afterFinally;"
-        "shouldNotReach"
-    }
-    if (value != "suspend(OK);finally;return;") return "fail2: $value"
 
+    if (res != "log(1);suspend(2);log(3);return(4);") return "FAIL: $res"
     return "OK"
 }

compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt

@@ -0,0 +1,69 @@
+// IGNORE_BACKEND: JS_IR
+// IGNORE_BACKEND: JVM_IR
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+class Controller {
+    var result = ""
+
+    suspend fun <T> suspendAndLog(value: T): T = suspendCoroutineUninterceptedOrReturn { c ->
+        result += "suspend($value);"
+        c.resume(value)
+        COROUTINE_SUSPENDED
+    }
+}
+
+fun builder(c: suspend Controller.() -> String): String {
+    val controller = Controller()
+    c.startCoroutine(controller, handleResultContinuation {
+        controller.result += "return($it);"
+    })
+    return controller.result
+}
+
+fun builderUnit(c: suspend Controller.() -> Unit): String {
+    val controller = Controller()
+    c.startCoroutine(controller, handleResultContinuation {
+        controller.result += "return;"
+    })
+    return controller.result
+}
+
+fun <T> id(value: T) = value
+
+fun box(): String {
+    var value = builder {
+        try {
+            if (id(23) == 23) {
+                return@builder suspendAndLog("OK")
+            }
+        }
+        finally {
+            result += "finally;"
+        }
+        result += "afterFinally;"
+        "shouldNotReach"
+    }
+    if (value != "suspend(OK);finally;return(OK);") return "fail1: $value"
+
+    value = builderUnit {
+        try {
+            if (id(23) == 23) {
+                suspendAndLog("OK")
+                return@builderUnit
+            }
+        }
+        finally {
+            result += "finally;"
+        }
+        result += "afterFinally;"
+        "shouldNotReach"
+    }
+    if (value != "suspend(OK);finally;return;") return "fail2: $value"
+
+    return "OK"
+}

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

@@ -6170,6 +6170,16 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnFromFinally.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("switchLikeWhen.kt")
             public void testSwitchLikeWhen_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt", "kotlin.coroutines.experimental");

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

@@ -6170,6 +6170,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnFromFinally.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("switchLikeWhen.kt")
             public void testSwitchLikeWhen_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt", "kotlin.coroutines.experimental");

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

@@ -6170,6 +6170,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnFromFinally.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("switchLikeWhen.kt")
             public void testSwitchLikeWhen_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt", "kotlin.coroutines.experimental");

js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineBodyTransformer.kt

@@ -290,6 +290,8 @@ class CoroutineBodyTransformer(private val context: CoroutineTransformationConte
             }
         }
 
+        tryStack.removeAt(tryStack.lastIndex)
+
         // Handle finally node
         if (finallyNode != null) {
             currentBlock = finallyBlock
@@ -298,8 +300,6 @@ class CoroutineBodyTransformer(private val context: CoroutineTransformationConte
             hasFinallyBlocks = true
         }
 
-        tryStack.removeAt(tryStack.lastIndex)
-
         currentBlock = successor
     }
 

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

@@ -5430,6 +5430,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnFromFinally.kt", "kotlin.coroutines.experimental");
             }
 
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines.experimental");
+            }
+
             @TestMetadata("switchLikeWhen.kt")
             public void testSwitchLikeWhen_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt", "kotlin.coroutines.experimental");

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

@@ -5895,6 +5895,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnFromFinally.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("returnWithFinally.kt")
+            public void testReturnWithFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/returnWithFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("switchLikeWhen.kt")
             public void testSwitchLikeWhen_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt", "kotlin.coroutines.experimental");