[JVM] Do not extend local ranges across control-flow merges.

The coroutine method transformer extends the range of locals
across code where the local is not live when it is safe to do
so. However, it only bailed out for one case of control-flow
merging, namely backwards branches for loops. That is not
sufficient as there can be control flow merges without loops
where the local is only defined on one control-flow path.

This change generalizes the bailout to any control-flow merge.

^KT-49834 Fixed

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

@@ -1251,45 +1251,24 @@ private fun updateLvtAccordingToLiveness(method: MethodNode, isForNamedFunction:
                 // No variable in LVT -> do not add one
                 val lvtRecord = oldLvt.findRecord(insnIndex, variableIndex) ?: continue
                 if (lvtRecord.name == CONTINUATION_VARIABLE_NAME || lvtRecord.name == SUSPEND_CALL_RESULT_NAME) continue
-                // End the local when it is no longer live. Since it is not live, we will not spill and unspill it across
-                // suspension points. It is tempting to keep it alive until the next suspension point to leave it visible in
-                // the debugger for as long as possible. However, in the case of loops, the resumption after suspension can
-                // have a backwards edge targeting instruction between the point of death and the next suspension point.
-                //
-                // For example, code such as the following:
-                //
-                //    listOf<String>.forEach {
-                //       yield(it)
-                //    }
-                //
-                // Generates code of this form with a back edge after resumption that will lead to invalid locals tables
-                // if the local range is extended to the next suspension point.
-                //
-                //        iterator = iterable.iterator()
-                //    L1: (iterable dies here)
-                //        load iterator.next if there
-                //        yield suspension point
-                //
-                //    L2: (resumption point)
-                //        restore live variables (not including iterable)
-                //        goto L1 (iterator not restored here, so we cannot not have iterator live at L1)
+                // End the local when it is no longer live and then attempt to extend its range when safe.
                 val endLabel = nextLabel(insn.next)?.let { min(lvtRecord.end, it) } ?: lvtRecord.end
                 // startLabel can be null in case of parameters
                 @Suppress("NAME_SHADOWING") val startLabel = startLabel ?: lvtRecord.start
 
                 // Attempt to extend existing local variable node corresponding to the record in
-                // the original local variable table, if there is no back-edge
+                // the original local variable table if there are no control-flow merges.
                 val latest = oldLvtNodeToLatestNewLvtNode[lvtRecord]
-                // if we can extend the previous range to where the local variable dies, we do not need a
+                // If we can extend the previous range to where the local variable dies, we do not need a
                 // new entry, we know we cannot extend it to the lvt.endOffset, if we could we would have
                 // done so when we added it below.
-                val extended = latest?.extendRecordIfPossible(method, suspensionPoints, lvtRecord.end) ?: false
+                val extended = latest?.extendRecordIfPossible(method, suspensionPoints, lvtRecord.end, liveness) ?: false
                 if (!extended) {
                     val new = LocalVariableNode(lvtRecord.name, lvtRecord.desc, lvtRecord.signature, startLabel, endLabel, lvtRecord.index)
                     oldLvtNodeToLatestNewLvtNode[lvtRecord] = new
                     method.localVariables.add(new)
-                    // see if we can extend it all the way to the old end
-                    new.extendRecordIfPossible(method, suspensionPoints, lvtRecord.end)
+                    // See if we can extend it all the way to the old end.
+                    new.extendRecordIfPossible(method, suspensionPoints, lvtRecord.end, liveness)
                 }
             }
         }
@@ -1313,26 +1292,71 @@ private fun updateLvtAccordingToLiveness(method: MethodNode, isForNamedFunction:
     }
 }
 
-/* We cannot extend a record if there is STORE instruction or a back-edge.
- * STORE instructions can signify a unspilling operation, in which case, the variable will become visible before it unspilled,
- * back-edges occur in loops.
+/* We cannot extend a record if there is STORE instruction or a control-flow merge.
+ *
+ * STORE instructions can signify a unspilling operation, in which case, the variable will become visible before it unspilled.
+ *
+ * If there is a control-flow merge point in a range where a variable is dead, it might not have been restored on one of the paths
+ * and therefore it is not safe to extend the record across the control flow merge point.
+ *
+ * For example, code such as the following:
+ *
+ *    listOf<String>.forEach {
+ *       yield(it)
+ *    }
+ *
+ * Generates code of this form with a back edge after resumption that will lead to invalid locals tables
+ * if the local range is extended to the next suspension point. L1 is a merge point and therefore, we do
+ * not extend.
+ *
+ *        iterator = iterable.iterator()
+ *    L1: (iterable dies here)
+ *        load iterator.next if there
+ *        yield suspension point
+ *
+ *    L2: (resumption point)
+ *        restore live variables (not including iterable)
+ *        goto L1 (iterator not restored here, so we cannot not have iterator live at L1)
+ *
+ * Code such as:
+ *
+ *    val value = getValue()
+ *    return if (value == null) {
+ *        computeValueAsync()  // suspension point
+ *    } else {
+ *        value
+ *    } + "K"
+ *
+ * Generates code of this form, where it is not safe to extend the `value` local variable across the control-flow
+ * merge because it is dead and will not have been restored after the suspend point in one of the branches.
+ *
+ *      value = getValue()
+ *      if (value != null) goto L2
+ *  L1: (value dead here)
+ *      temp = computeValueAsync() // suspension point and resumption point, value NOT restored as it is dead
+ *      load temp
+ *      goto L3
+ *  L2: (value alive here)
+ *      load value
+ *  L3: (merge point, cannot extend `value` local across as it is not defined on one of the paths)
+ *      load "K"
+ *      add strings
+ *      return
  *
  * @return true if the range has been extended
  */
 private fun LocalVariableNode.extendRecordIfPossible(
     method: MethodNode,
     suspensionPoints: List<LabelNode>,
-    endLabel: LabelNode
+    endLabel: LabelNode,
+    liveness: List<VariableLivenessFrame>
 ): Boolean {
     val nextSuspensionPointLabel = suspensionPoints.find { it in InsnSequence(end, endLabel) } ?: endLabel
 
     var current: AbstractInsnNode? = end
+    var index = method.instructions.indexOf(current)
     while (current != null && current != nextSuspensionPointLabel) {
-        if (current is JumpInsnNode) {
-            if (method.instructions.indexOf(current.label) < method.instructions.indexOf(current)) {
-                return false
-            }
-        }
+        if (liveness[index].isControlFlowMerge()) return false
         // TODO: HACK
         // TODO: Find correct label, which is OK to be used as end label.
         if (current.opcode == Opcodes.ARETURN && nextSuspensionPointLabel != endLabel) return false
@@ -1340,6 +1364,7 @@ private fun LocalVariableNode.extendRecordIfPossible(
             return false
         }
         current = current.next
+        ++index
     }
     end = nextSuspensionPointLabel
     return true

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/backwardAnalysis.kt

@@ -21,6 +21,7 @@ import org.jetbrains.org.objectweb.asm.tree.MethodNode
 
 interface VarFrame<F : VarFrame<F>> {
     fun mergeFrom(other: F)
+    fun markControlFlowMerge()
     override fun equals(other: Any?): Boolean
     override fun hashCode(): Int
 }
@@ -48,6 +49,7 @@ fun <F : VarFrame<F>> analyze(node: MethodNode, interpreter: BackwardAnalysisInt
             for (successorIndex in graph.getSuccessorsIndices(insn)) {
                 newFrame.mergeFrom(frames[successorIndex])
             }
+            if (graph.getPredecessorsIndices(insn).size > 1) newFrame.markControlFlowMerge()
 
             interpreter.def(newFrame, insn)
             interpreter.use(newFrame, insn)

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/variableLiveness.kt

@@ -14,11 +14,16 @@ import java.util.*
 
 class VariableLivenessFrame(val maxLocals: Int) : VarFrame<VariableLivenessFrame> {
     private val bitSet = BitSet(maxLocals)
+    private var controlFlowMerge = false
 
     override fun mergeFrom(other: VariableLivenessFrame) {
         bitSet.or(other.bitSet)
     }
 
+    override fun markControlFlowMerge() {
+        controlFlowMerge = true
+    }
+
     fun markAlive(varIndex: Int) {
         bitSet.set(varIndex, true)
     }
@@ -29,14 +34,17 @@ class VariableLivenessFrame(val maxLocals: Int) : VarFrame<VariableLivenessFrame
 
     fun isAlive(varIndex: Int): Boolean = bitSet.get(varIndex)
 
+    fun isControlFlowMerge(): Boolean = controlFlowMerge
+
     override fun equals(other: Any?): Boolean {
         if (other !is VariableLivenessFrame) return false
-        return bitSet == other.bitSet
+        return bitSet == other.bitSet && controlFlowMerge == other.controlFlowMerge
     }
 
-    override fun hashCode() = bitSet.hashCode()
+    override fun hashCode() = bitSet.hashCode() * 31 + controlFlowMerge.hashCode()
 
-    override fun toString(): String = (0 until maxLocals).map { if (bitSet[it]) '@' else '_' }.joinToString(separator = "")
+    override fun toString(): String =
+        (if (controlFlowMerge) "*" else " ") + (0 until maxLocals).map { if (bitSet[it]) '@' else '_' }.joinToString(separator = "")
 }
 
 fun analyzeLiveness(method: MethodNode): List<VariableLivenessFrame> =
@@ -59,4 +67,4 @@ private fun useVar(frame: VariableLivenessFrame, insn: AbstractInsnNode) {
     } else if (insn is IincInsnNode) {
         frame.markAlive(insn.`var`)
     }
-}
+}

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -12623,6 +12623,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @Test
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @Test
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {

compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt

@@ -0,0 +1,18 @@
+// WITH_STDLIB
+
+fun getValue() : String? = null
+suspend fun computeValue() = "O"
+
+suspend fun repro() : String {
+    val value = getValue()
+    return if (value == null) {
+        computeValue()
+    } else {
+        value
+    } + "K"
+}
+
+// This test is checking that the local variable table for `repro` is valid.
+// This is checked because the D8 dexer is run on the produced code and
+// we fail the tests on warnings because of invalid locals.
+fun box() = "OK"

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -12545,6 +12545,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @Test
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @Test
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -12623,6 +12623,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @Test
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @Test
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {

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

@@ -10121,6 +10121,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/lvtWithInlineOnly.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java

@@ -9403,6 +9403,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @Test
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @Test
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java

@@ -9445,6 +9445,12 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @Test
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @Test
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -8243,6 +8243,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
             }
 
+            @TestMetadata("kt49834.kt")
+            public void testKt49834() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+            }
+
             @TestMetadata("lvtWithInlineOnly.kt")
             public void testLvtWithInlineOnly() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/varSpilling/lvtWithInlineOnly.kt");

native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeExtBlackBoxTestGenerated.java

@@ -12759,6 +12759,12 @@ public class NativeExtBlackBoxTestGenerated extends AbstractNativeBlackBoxTest {
                     runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt38925.kt");
                 }
 
+                @Test
+                @TestMetadata("kt49834.kt")
+                public void testKt49834() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/varSpilling/kt49834.kt");
+                }
+
                 @Test
                 @TestMetadata("lvtWithInlineOnly.kt")
                 public void testLvtWithInlineOnly() throws Exception {