FIR DFA: make continue jump to condition entry, not loop entry

It's also not a backwards jump in do-while, unless it's in the loop's
condition, which is a stupid "feature" IMO. As you can probably tell
from the comments added in this commit.

compiler/fir/analysis-tests/testData/resolve/cfg/jumps.dot

@@ -354,7 +354,7 @@ digraph jumps_kt {
     113 -> {114};
     114 -> {115};
     115 -> {116} [style=dotted];
-    115 -> {102} [color=green style=dashed];
+    115 -> {103} [color=green style=dashed];
     116 -> {117} [style=dotted];
     117 -> {118} [style=dotted];
     118 -> {119} [style=dotted];

compiler/fir/analysis-tests/testData/resolve/cfg/tryCatch.dot

@@ -324,7 +324,7 @@ digraph tryCatch_kt {
     96 -> {97};
     97 -> {98};
     98 -> {99} [style=dotted];
-    98 -> {46} [color=green style=dashed];
+    98 -> {47} [color=green style=dashed];
     99 -> {100} [style=dotted];
     100 -> {101} [style=dotted];
     101 -> {102} [style=dotted];

compiler/fir/analysis-tests/testData/resolve/smartcasts/boundSmartcasts/boundSmartcastsInBranches.dot

@@ -136,7 +136,7 @@ digraph boundSmartcastsInBranches_kt {
     36 -> {37};
     37 -> {38};
     38 -> {39} [style=dotted];
-    38 -> {17} [color=green style=dashed];
+    38 -> {18} [color=green style=dashed];
     39 -> {40} [style=dotted];
     40 -> {41} [style=dotted];
     41 -> {42} [style=dotted];

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -85,7 +85,7 @@ class ControlFlowGraphBuilder {
     private val exitsFromPostponedAnonymousFunctions: MutableMap<FirFunctionSymbol<*>, PostponedLambdaExitNode> = mutableMapOf()
     private val parentGraphForAnonymousFunctions: MutableMap<FirFunctionSymbol<*>, ControlFlowGraph> = mutableMapOf()
 
-    private val loopEnterNodes: NodeStorage<FirElement, CFGNode<FirElement>> = NodeStorage()
+    private val loopConditionEnterNodes: NodeStorage<FirElement, LoopConditionEnterNode> = NodeStorage()
     private val loopExitNodes: NodeStorage<FirLoop, LoopExitNode> = NodeStorage()
 
     private val exitsFromCompletedPostponedAnonymousFunctions: MutableList<MutableList<CFGNode<*>>> = mutableListOf()
@@ -670,7 +670,7 @@ class ControlFlowGraphBuilder {
         val node = createJumpNode(jump)
         val nextNode = when (jump) {
             is FirReturnExpression -> exitTargetsForReturn[jump.target.labeledElement.symbol]
-            is FirContinueExpression -> loopEnterNodes[jump.target.labeledElement]
+            is FirContinueExpression -> loopConditionEnterNodes[jump.target.labeledElement.condition]
             is FirBreakExpression -> loopExitNodes[jump.target.labeledElement]
             else -> throw IllegalArgumentException("Unknown jump type: ${jump.render()}")
         }
@@ -681,12 +681,17 @@ class ControlFlowGraphBuilder {
             else -> NormalPath
         }
 
+        // while (...) continue // <- jump back
+        // do continue while (...) // <- jump forward (block exit node not created yet)
+        // do ... while (continue) // <- jump back (to itself), believe it or not
+        val isBack = nextNode is LoopConditionEnterNode &&
+                (nextNode.loop !is FirDoWhileLoop || nextNode.previousNodes.any { it is LoopBlockExitNode })
         addNodeWithJump(
             node,
             nextNode,
-            isBack = jump is FirContinueExpression,
+            isBack = isBack,
             trackJump = jump is FirReturnExpression,
-            label = NormalPath,
+            label = if (isBack) LoopBackPath else NormalPath,
             labelForFinallyBLock = labelForFinallyBLock
         )
         return node
@@ -758,14 +763,12 @@ class ControlFlowGraphBuilder {
     fun enterWhileLoop(loop: FirLoop): Pair<LoopEnterNode, LoopConditionEnterNode> {
         val loopEnterNode = createLoopEnterNode(loop).also {
             addNewSimpleNode(it)
-            loopEnterNodes.push(it)
         }
         loopExitNodes.push(createLoopExitNode(loop))
         levelCounter++
         val conditionEnterNode = createLoopConditionEnterNode(loop.condition, loop).also {
             addNewSimpleNode(it)
-            // put conditional node twice so we can refer it after exit from loop block
+            loopConditionEnterNodes.push(it)
-            lastNodes.push(it)
         }
         levelCounter++
         return loopEnterNode to conditionEnterNode
@@ -784,15 +787,11 @@ class ControlFlowGraphBuilder {
     }
 
     fun exitWhileLoop(loop: FirLoop): Pair<LoopBlockExitNode, LoopExitNode> {
-        loopEnterNodes.pop()
         levelCounter--
         val loopBlockExitNode = createLoopBlockExitNode(loop)
         popAndAddEdge(loopBlockExitNode)
-        if (lastNodes.isNotEmpty) {
+        val conditionEnterNode = loopConditionEnterNodes.pop()
-            val conditionEnterNode = lastNodes.pop()
+        addBackEdge(loopBlockExitNode, conditionEnterNode, label = LoopBackPath)
-            require(conditionEnterNode is LoopConditionEnterNode) { loop.render() }
-            addBackEdge(loopBlockExitNode, conditionEnterNode, label = LoopBackPath)
-        }
         val loopExitNode = loopExitNodes.pop()
         loopExitNode.updateDeadStatus()
         lastNodes.push(loopExitNode)
@@ -811,7 +810,7 @@ class ControlFlowGraphBuilder {
         addNewSimpleNode(blockEnterNode)
         // put block enter node twice so we can refer it after exit from loop condition
         lastNodes.push(blockEnterNode)
-        loopEnterNodes.push(blockEnterNode)
+        loopConditionEnterNodes.push(createLoopConditionEnterNode(loop.condition, loop))
         levelCounter++
         return loopEnterNode to blockEnterNode
     }
@@ -819,13 +818,15 @@ class ControlFlowGraphBuilder {
     fun enterDoWhileLoopCondition(loop: FirLoop): Pair<LoopBlockExitNode, LoopConditionEnterNode> {
         levelCounter--
         val blockExitNode = createLoopBlockExitNode(loop).also { addNewSimpleNode(it) }
-        val conditionEnterNode = createLoopConditionEnterNode(loop.condition, loop).also { addNewSimpleNode(it) }
+        // This may sound shocking, but `do...while` conditions can `continue` to themselves,
+        // so we can't pop the node off the stack here.
+        val conditionEnterNode = loopConditionEnterNodes.top().also { addNewSimpleNode(it) }
         levelCounter++
         return blockExitNode to conditionEnterNode
     }
 
     fun exitDoWhileLoop(loop: FirLoop): Pair<LoopConditionExitNode, LoopExitNode> {
-        loopEnterNodes.pop()
+        loopConditionEnterNodes.pop()
         levelCounter--
         val conditionExitNode = createLoopConditionExitNode(loop.condition)
         val conditionBooleanValue = conditionExitNode.booleanConstValue
@@ -1556,12 +1557,7 @@ class ControlFlowGraphBuilder {
         popAndAddEdge(node, preferredKind)
         if (targetNode != null) {
             if (isBack) {
-                if (targetNode is LoopEnterNode) {
+                addBackEdge(node, targetNode, label = label)
-                    // `continue` to the loop header
-                    addBackEdge(node, targetNode, label = LoopBackPath)
-                } else {
-                    addBackEdge(node, targetNode, label = label)
-                }
             } else {
                 // go through all final nodes between node and target
                 val finallyNodes = finallyBefore(targetNode)

compiler/testData/diagnostics/tests/smartCasts/loops/doWhileEarlyContinue.fir.kt

@@ -1,13 +0,0 @@
-fun x(): Boolean { return true }
-
-fun y(): Boolean { return false }
-
-public fun foo(p: String?): Int {
-    do {
-        if (y()) continue
-        // We do not always reach this statement
-        p!!.length
-    } while (!x())
-    // Here we have do while loop but p is still nullable due to continue before
-    return p.length
-}

compiler/testData/diagnostics/tests/smartCasts/loops/doWhileEarlyContinue.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun x(): Boolean { return true }
 
 fun y(): Boolean { return false }

compiler/testData/diagnostics/tests/smartCasts/loops/nestedDoWhileWithLongContinue.fir.kt

@@ -5,10 +5,10 @@ public fun foo(p: String?, r: String?): Int {
         do {
             p!!.length
             if (!x()) continue@outer
-        } while (r == null)  
+        } while (r == null)
     } while (!x())
     // Auto cast NOT possible due to long continue
-    r.length
+    r<!UNSAFE_CALL!>.<!>length
     // Auto cast possible
     return p.length
 }

compiler/testData/diagnostics/tests/smartCasts/loops/nestedDoWhileWithLongContinue.kt

@@ -5,7 +5,7 @@ public fun foo(p: String?, r: String?): Int {
         do {
             p!!.length
             if (!x()) continue@outer
-        } while (r == null)  
+        } while (r == null)
     } while (!x())
     // Auto cast NOT possible due to long continue
     r<!UNSAFE_CALL!>.<!>length

compiler/tests-spec/testData/diagnostics/notLinked/dfa/pos/35.fir.kt

@@ -70,7 +70,7 @@ fun case_6(x: Any?) {
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any? & kotlin.Any")!>x<!>.equals(10)
     } while (false)
 
-    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any? & kotlin.Any")!>x<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any?")!>x<!>
 }
 
 // TESTCASE NUMBER: 7