Support jump out of the constructor call in suspend functions

Jump out from expression (e.g., break or continue expression in call
arguments) requires stack normalization, which inserts POP instructions.
POPping an uninitialized value is similar to ASTORE, except that it
doesn't store a value to a local variable. Such POP instructions should
be removed during postprocessing of the uninitialized stores.

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

@@ -19,6 +19,8 @@ package org.jetbrains.kotlin.codegen.coroutines
 import org.jetbrains.kotlin.codegen.inline.insnText
 import org.jetbrains.kotlin.codegen.optimization.common.*
 import org.jetbrains.kotlin.codegen.optimization.fixStack.peek
+import org.jetbrains.kotlin.codegen.optimization.fixStack.peekWords
+import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.*
@@ -90,19 +92,21 @@ class UninitializedStoresProcessor(private val methodNode: MethodNode) {
                 this::UninitializedNewValueFrame
         ).analyze()
 
+        interpreter.analyzePopInstructions(frames)
+
         for ((index, insn) in methodNode.instructions.toArray().withIndex()) {
             val frame = frames[index] ?: continue
             val uninitializedValue = frame.getUninitializedValueForConstructorCall(insn) ?: continue
 
             val newInsn = uninitializedValue.newInsn
-            val copyUsages: Set<AbstractInsnNode> = interpreter.uninitializedValuesToCopyUsages[newInsn]!!
+            val removableUsages: Set<AbstractInsnNode> = interpreter.uninitializedValuesToRemovableUsages[newInsn]!!
-            assert(copyUsages.isNotEmpty()) { "At least DUP copy operation expected" }
+            assert(removableUsages.isNotEmpty()) { "At least DUP copy operation expected" }
 
             // Value generated by NEW wasn't store to local/field (only DUPed)
-            if (copyUsages.size == 1) continue
+            if (removableUsages.size == 1) continue
 
             methodNode.instructions.run {
-                removeAll(copyUsages)
+                removeAll(removableUsages)
 
                 // Replace 'NEW C' instruction with "manual" initialization of class 'C':
                 //      LDC [typeName for C]
@@ -194,10 +198,10 @@ class UninitializedStoresProcessor(private val methodNode: MethodNode) {
     private fun AbstractInsnNode.isConstructorCall() = this is MethodInsnNode && this.name == "<init>"
 
     private class UninitializedNewValueMarkerInterpreter(private val instructions: InsnList) : OptimizationBasicInterpreter() {
-        val uninitializedValuesToCopyUsages = hashMapOf<AbstractInsnNode, MutableSet<AbstractInsnNode>>()
+        val uninitializedValuesToRemovableUsages = hashMapOf<AbstractInsnNode, MutableSet<AbstractInsnNode>>()
         override fun newOperation(insn: AbstractInsnNode): BasicValue? {
             if (insn.opcode == Opcodes.NEW) {
-                uninitializedValuesToCopyUsages.getOrPut(insn) { mutableSetOf() }
+                uninitializedValuesToRemovableUsages.getOrPut(insn) { mutableSetOf() }
                 return UninitializedNewValue(insn as TypeInsnNode, insn.desc)
             }
             return super.newOperation(insn)
@@ -206,7 +210,7 @@ class UninitializedStoresProcessor(private val methodNode: MethodNode) {
         override fun copyOperation(insn: AbstractInsnNode, value: BasicValue?): BasicValue? {
             if (value is UninitializedNewValue) {
                 checkUninitializedObjectCopy(value.newInsn, insn)
-                uninitializedValuesToCopyUsages[value.newInsn]!!.add(insn)
+                uninitializedValuesToRemovableUsages[value.newInsn]!!.add(insn)
                 return value
             }
             return super.copyOperation(insn, value)
@@ -239,5 +243,33 @@ class UninitializedStoresProcessor(private val methodNode: MethodNode) {
 
         private val AbstractInsnNode.debugText
             get() = "${instructions.indexOf(this)}: $insnText"
+
+        fun analyzePopInstructions(frames: Array<Frame<BasicValue>?>) {
+            val insns = instructions.toArray()
+            for (i in frames.indices) {
+                val frame = frames[i] ?: continue
+                val insn = insns[i]
+                when (insn.opcode) {
+                    Opcodes.POP -> analyzePop(insn, frame)
+                    Opcodes.POP2 -> analyzePop2(insn, frame)
+                }
+            }
+        }
+
+        private fun analyzePop(insn: AbstractInsnNode, frame: Frame<BasicValue>) {
+            val top = frame.top() ?: error("Stack underflow on POP: ${insn.debugText}")
+            if (top is UninitializedNewValue) {
+                uninitializedValuesToRemovableUsages[top.newInsn]!!.add(insn)
+            }
+        }
+
+        private fun analyzePop2(insn: AbstractInsnNode, frame: Frame<BasicValue>) {
+            val top2 = frame.peekWords(2) ?: error("Stack underflow on POP2: ${insn.debugText}")
+            for (value in top2) {
+                if (value is UninitializedNewValue) {
+                    error("Unexpected POP2 instruction for ${value.newInsn.debugText}: ${insn.debugText}")
+                }
+            }
+        }
     }
 }

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

@@ -0,0 +1,152 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import helpers.*
+import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.*
+
+class Controller {
+    suspend fun suspendHere(): String = suspendCoroutineOrReturn { x ->
+        x.resume("K")
+        COROUTINE_SUSPENDED
+    }
+
+    suspend fun suspendWithArgument(v: String): String = suspendCoroutineOrReturn { x ->
+        x.resume(v)
+        COROUTINE_SUSPENDED
+    }
+
+    suspend fun suspendWithDouble(v: Double): Double = suspendCoroutineOrReturn { x ->
+        x.resume(v)
+        COROUTINE_SUSPENDED
+    }
+}
+
+fun builder(c: suspend Controller.() -> Unit) {
+    c.startCoroutine(Controller(), EmptyContinuation)
+}
+
+class A(val first: String, val second: String) {
+    override fun toString() = "$first$second"
+}
+
+class B(val first: String, val second: String, val third: String) {
+    override fun toString() = "$first$second$third"
+}
+
+class C(val a: Long, val b: Double, val c: Int, val d: String) {
+    override fun toString() = "$a#$b#$c#$d"
+}
+
+val condition = true
+
+fun box(): String {
+    var result = "OK"
+
+    builder {
+        for (count in 0..3) {
+            val local = A(if (count > 0) break else "O", suspendHere())
+
+            if (count > 0) {
+                result = "fail 1: count=$count"
+                return@builder
+            }
+
+            if (local.toString() != "OK") {
+                result = "fail 1: $local"
+                return@builder
+            }
+        }
+
+        for (count in 0..3) {
+            val local = B(if (count > 0) break else "#", suspendWithArgument("O"), suspendHere())
+
+            if (count > 0) {
+                result = "fail 2: count=$count"
+                return@builder
+            }
+
+            if (local.toString() != "#OK") {
+                result = "fail 2: $local"
+                return@builder
+            }
+        }
+
+        for (count in 0..3) {
+            val local = B(suspendWithArgument("#"), if (count > 0) break else "O", suspendHere())
+
+            if (count > 0) {
+                result = "fail 3: count=$count"
+                return@builder
+            }
+
+            if (local.toString() != "#OK") {
+                result = "fail 3: $local"
+                return@builder
+            }
+        }
+
+        for (count in 0..3) {
+            val local = B(
+                    "#",
+                    B("",
+                      if (count > 0) break else "O",
+                      suspendWithArgument("")
+                    ).toString(),
+                    suspendHere()
+            )
+
+            if (count > 0) {
+                result = "fail 4: count=$count"
+                return@builder
+            }
+
+            if (local.toString() != "#OK") {
+                result = "fail 4: $local"
+                return@builder
+            }
+        }
+
+        loop@for (count in 0..3) {
+            val local = B(
+                    if (!condition) "1" else suspendWithArgument("#"),
+                    when {
+                        count > 0 -> break@loop
+                        condition -> suspendWithArgument("O")
+                        else -> "2"
+                    },
+                    if (condition) suspendHere() else suspendWithArgument("3")
+            )
+
+            if (count > 0) {
+                result = "fail 5: count=$count"
+                return@builder
+            }
+
+            if (local.toString() != "#OK") {
+                result = "fail 5: $local"
+                return@builder
+            }
+        }
+
+        for (count in 0..3) {
+            val local = C(
+                    1234567890123L,
+                    suspendWithDouble(3.14),
+                    42,
+                    if (count > 0) break else suspendWithArgument("OK")
+            )
+
+            if (count > 0) {
+                result = "fail 6: count=$count"
+                return@builder
+            }
+
+            if (local.toString() != "1234567890123#3.14#42#OK") {
+                result = "fail 6: $local"
+                return@builder
+            }
+        }
+    }
+
+    return result
+}

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

@@ -5504,6 +5504,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("suspendInTheMiddleOfObjectConstructionWithJumpOut.kt")
+        public void testSuspendInTheMiddleOfObjectConstructionWithJumpOut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendInTheMiddleOfObjectConstructionWithJumpOut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("suspensionInsideSafeCall.kt")
         public void testSuspensionInsideSafeCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");

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

@@ -5504,6 +5504,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("suspendInTheMiddleOfObjectConstructionWithJumpOut.kt")
+        public void testSuspendInTheMiddleOfObjectConstructionWithJumpOut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendInTheMiddleOfObjectConstructionWithJumpOut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("suspensionInsideSafeCall.kt")
         public void testSuspensionInsideSafeCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");

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

@@ -5504,6 +5504,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("suspendInTheMiddleOfObjectConstructionWithJumpOut.kt")
+        public void testSuspendInTheMiddleOfObjectConstructionWithJumpOut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendInTheMiddleOfObjectConstructionWithJumpOut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("suspensionInsideSafeCall.kt")
         public void testSuspensionInsideSafeCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");

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

@@ -6152,6 +6152,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("suspendInTheMiddleOfObjectConstructionWithJumpOut.kt")
+        public void testSuspendInTheMiddleOfObjectConstructionWithJumpOut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendInTheMiddleOfObjectConstructionWithJumpOut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("suspensionInsideSafeCall.kt")
         public void testSuspensionInsideSafeCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");