Make initial continuation able to be resumed with exception

#KT-14719 Fixed

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

@@ -72,7 +72,7 @@ class CoroutineTransformerMethodVisitor(
         }
 
         // Add global exception handler
-        processHandleExceptionCall(methodNode)
+        val isThereGlobalExceptionHandler = processHandleExceptionCall(methodNode)
 
         // Spill stack to variables before suspension points, try/catch blocks
         FixStackWithLabelNormalizationMethodTransformer().transform(classBuilder.thisName, methodNode)
@@ -95,8 +95,19 @@ class CoroutineTransformerMethodVisitor(
         methodNode.instructions.apply {
             val startLabel = LabelNode()
             val defaultLabel = LabelNode()
+            val firstToInsertBefore =
+                    if (isThereGlobalExceptionHandler) {
+                        // Insert after relevant NOP insn (i.e. before next instruction of this NOP)
+                        val globalExceptionHandler = methodNode.tryCatchBlocks.last()
+                        assert(globalExceptionHandler.start is LabelNode && globalExceptionHandler.start.next.opcode == Opcodes.NOP) {
+                            "In a case of global exception handler first insn should be a label"
+                        }
+                        globalExceptionHandler.start.next.next
+                    }
+                    else
+                        first
             // tableswitch(this.label)
-            insertBefore(first,
+            insertBefore(firstToInsertBefore,
                          insnListOf(
                                  VarInsnNode(Opcodes.ALOAD, 0),
                                  FieldInsnNode(
@@ -113,6 +124,7 @@ class CoroutineTransformerMethodVisitor(
                          )
             )
 
+            insert(startLabel, withInstructionAdapter(InstructionAdapter::generateResumeWithExceptionCheck))
 
             insert(last, withInstructionAdapter {
                 visitLabel(defaultLabel.label)
@@ -335,15 +347,7 @@ class CoroutineTransformerMethodVisitor(
             remove(possibleTryCatchBlockStart.previous)
 
             insert(possibleTryCatchBlockStart, withInstructionAdapter {
-                // Check if resumeWithException has been called
+                generateResumeWithExceptionCheck()
-                load(2, AsmTypes.OBJECT_TYPE)
-                dup()
-                val noExceptionLabel = Label()
-                ifnull(noExceptionLabel)
-                athrow()
-
-                mark(noExceptionLabel)
-                pop()
 
                 // Load continuation argument just like suspending function returns it
                 load(1, AsmTypes.OBJECT_TYPE)
@@ -404,9 +408,9 @@ class CoroutineTransformerMethodVisitor(
         return
     }
 
-    private fun processHandleExceptionCall(methodNode: MethodNode) {
+    private fun processHandleExceptionCall(methodNode: MethodNode): Boolean {
         val instructions = methodNode.instructions
-        val marker = instructions.toArray().firstOrNull() { it.isHandleExceptionMarker() } ?: return
+        val marker = instructions.toArray().firstOrNull { it.isHandleExceptionMarker() } ?: return false
 
         assert(instructions.toArray().count { it.isHandleExceptionMarker() } == 1) {
             "Found more than one handleException markers"
@@ -427,6 +431,8 @@ class CoroutineTransformerMethodVisitor(
         instructions.set(exceptionArgument, VarInsnNode(Opcodes.ALOAD, maxVar))
 
         methodNode.tryCatchBlocks.add(TryCatchBlockNode(startLabel, endLabel, endLabel, AsmTypes.JAVA_THROWABLE_TYPE.internalName))
+
+        return true
     }
 
     private fun AbstractInsnNode.isHandleExceptionMarker() =
@@ -436,6 +442,18 @@ class CoroutineTransformerMethodVisitor(
             this is MethodInsnNode && this.owner == COROUTINE_MARKER_OWNER && this.name == HANDLE_EXCEPTION_ARGUMENT_MARKER_NAME
 }
 
+private fun InstructionAdapter.generateResumeWithExceptionCheck() {
+    // Check if resumeWithException has been called
+    load(2, AsmTypes.OBJECT_TYPE)
+    dup()
+    val noExceptionLabel = Label()
+    ifnull(noExceptionLabel)
+    athrow()
+
+    mark(noExceptionLabel)
+    pop()
+}
+
 private fun Type.fieldNameForVar(index: Int) = descriptor.first() + "$" + index
 
 private fun withInstructionAdapter(block: InstructionAdapter.() -> Unit): InsnList {

compiler/testData/codegen/box/coroutines/beginWithException.kt

@@ -0,0 +1,33 @@
+// WITH_RUNTIME
+class Controller {
+    var exception: Throwable? = null
+
+    operator fun handleException(t: Throwable, c: Continuation<Nothing>) {
+        exception = t
+    }
+
+    suspend fun suspendHere(x: Continuation<Any>) {}
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    val controller = Controller()
+    c(controller).resumeWithException(RuntimeException("OK"))
+
+    if (controller.exception?.message != "OK") {
+        throw RuntimeException("Unexpected result: ${controller.exception?.message}")
+    }
+}
+
+fun box(): String {
+    var result = "OK"
+    builder {
+        suspendHere()
+        result = "fail 1"
+    }
+
+    builder {
+        result = "fail 2"
+    }
+
+    return result
+}

compiler/testData/codegen/box/coroutines/beginWithExceptionNoHandleException.kt

@@ -0,0 +1,33 @@
+// WITH_RUNTIME
+class Controller {
+    suspend fun suspendHere(x: Continuation<Any>) {}
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    try {
+        val controller = Controller()
+        c(controller).resumeWithException(RuntimeException("OK"))
+    }
+    catch(e: Exception) {
+        if (e?.message != "OK") {
+            throw RuntimeException("Unexpected result: ${e?.message}")
+        }
+        return
+    }
+
+    throw RuntimeException("Exception must be thrown above")
+}
+
+fun box(): String {
+    var result = "OK"
+    builder {
+        suspendHere()
+        result = "fail 1"
+    }
+
+    builder {
+        result = "fail 2"
+    }
+
+    return result
+}

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

@@ -4417,6 +4417,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/coroutines"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("beginWithException.kt")
+        public void testBeginWithException() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/beginWithException.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("beginWithExceptionNoHandleException.kt")
+        public void testBeginWithExceptionNoHandleException() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/beginWithExceptionNoHandleException.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("coercionToUnit.kt")
         public void testCoercionToUnit() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coercionToUnit.kt");

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

@@ -4417,6 +4417,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/coroutines"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("beginWithException.kt")
+        public void testBeginWithException() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/beginWithException.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("beginWithExceptionNoHandleException.kt")
+        public void testBeginWithExceptionNoHandleException() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/beginWithExceptionNoHandleException.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("coercionToUnit.kt")
         public void testCoercionToUnit() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coercionToUnit.kt");