[FIR] Fix passing flow throw loops

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -367,20 +367,23 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     fun exitWhileLoopCondition(loop: FirLoop) {
-        graphBuilder.exitWhileLoopCondition(loop).mergeIncomingFlow()
+        val (loopConditionExitNode, loopBlockEnterNode) = graphBuilder.exitWhileLoopCondition(loop)
+        loopConditionExitNode.mergeIncomingFlow()
+        loopBlockEnterNode.mergeIncomingFlow()
     }
 
     fun exitWhileLoop(loop: FirLoop) {
-        graphBuilder.exitWhileLoop(loop).also { (blockExitNode, exitNode) ->
+        val (blockExitNode, exitNode) = graphBuilder.exitWhileLoop(loop)
-            blockExitNode.mergeIncomingFlow()
+        blockExitNode.mergeIncomingFlow()
-            exitNode.mergeIncomingFlow()
+        exitNode.mergeIncomingFlow()
-        }
     }
 
     // ----------------------------------- Do while Loop -----------------------------------
 
     fun enterDoWhileLoop(loop: FirLoop) {
-        graphBuilder.enterDoWhileLoop(loop).mergeIncomingFlow()
+        val (loopEnterNode, loopBlockEnterNode) = graphBuilder.enterDoWhileLoop(loop)
+        loopEnterNode.mergeIncomingFlow()
+        loopBlockEnterNode.mergeIncomingFlow()
     }
 
     fun enterDoWhileLoopCondition(loop: FirLoop) {
@@ -390,7 +393,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     fun exitDoWhileLoop(loop: FirLoop) {
-        graphBuilder.exitDoWhileLoop(loop).mergeIncomingFlow()
+        val (loopConditionExitNode, loopExitNode) = graphBuilder.exitDoWhileLoop(loop)
+        loopConditionExitNode.mergeIncomingFlow()
+        loopExitNode.mergeIncomingFlow()
     }
 
     // ----------------------------------- Try-catch-finally -----------------------------------

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

@@ -238,7 +238,7 @@ class ControlFlowGraphBuilder {
         return loopEnterNode to conditionEnterNode
     }
 
-    fun exitWhileLoopCondition(loop: FirLoop): LoopConditionExitNode {
+    fun exitWhileLoopCondition(loop: FirLoop): Pair<LoopConditionExitNode, LoopBlockEnterNode> {
         levelCounter--
         val conditionExitNode = createLoopConditionExitNode(loop.condition)
         addNewSimpleNode(conditionExitNode)
@@ -247,7 +247,7 @@ class ControlFlowGraphBuilder {
         val loopBlockEnterNode = createLoopBlockEnterNode(loop)
         addNewSimpleNode(loopBlockEnterNode, conditionConstBooleanValue == false)
         levelCounter++
-        return conditionExitNode
+        return conditionExitNode to loopBlockEnterNode
     }
 
     fun exitWhileLoop(loop: FirLoop): Pair<LoopBlockExitNode, LoopExitNode> {
@@ -269,8 +269,9 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Do while Loop -----------------------------------
 
-    fun enterDoWhileLoop(loop: FirLoop): LoopBlockEnterNode {
+    fun enterDoWhileLoop(loop: FirLoop): Pair<LoopEnterNode, LoopBlockEnterNode> {
-        addNewSimpleNode(createLoopEnterNode(loop))
+        val loopEnterNode = createLoopEnterNode(loop)
+        addNewSimpleNode(loopEnterNode)
         loopExitNodes.push(createLoopExitNode(loop))
         levelCounter++
         val blockEnterNode = createLoopBlockEnterNode(loop)
@@ -279,7 +280,7 @@ class ControlFlowGraphBuilder {
         lastNodes.push(blockEnterNode)
         loopEnterNodes.push(blockEnterNode)
         levelCounter++
-        return blockEnterNode
+        return loopEnterNode to blockEnterNode
     }
 
     fun enterDoWhileLoopCondition(loop: FirLoop): Pair<LoopBlockExitNode, LoopConditionEnterNode> {
@@ -290,7 +291,7 @@ class ControlFlowGraphBuilder {
         return blockExitNode to conditionEnterNode
     }
 
-    fun exitDoWhileLoop(loop: FirLoop): LoopExitNode {
+    fun exitDoWhileLoop(loop: FirLoop): Pair<LoopConditionExitNode, LoopExitNode> {
         loopEnterNodes.pop()
         levelCounter--
         val conditionExitNode = createLoopConditionExitNode(loop.condition)
@@ -304,7 +305,7 @@ class ControlFlowGraphBuilder {
         loopExit.markAsDeadIfNecessary()
         lastNodes.push(loopExit)
         levelCounter--
-        return loopExit
+        return conditionExitNode to loopExit
     }
 
     // ----------------------------------- Boolean operators -----------------------------------

compiler/fir/resolve/testData/resolve/cfg/complex.dot

@@ -285,4 +285,128 @@ digraph complex_kt {
   88 -> {89} [style=dotted];
   89 -> {90} [style=dotted];
 
+  subgraph cluster_25 {
+    color=red
+    91 [label="Enter function firstIsInstanceOrNull" style="filled" fillcolor=red];
+    subgraph cluster_26 {
+      color=blue
+      92 [label="Enter block"];
+      93 [label="Access variable this@R|/firstIsInstanceOrNull|"];
+      94 [label="Variable declaration: lval <range>: R|kotlin/sequences/Sequence<*>|"];
+      95 [label="Access variable R|<local>/<range>|"];
+      96 [label="Function call: R|<local>/<range>|.R|kotlin/sequences/Sequence.iterator|()"];
+      97 [label="Variable declaration: lval <iterator>: R|kotlin/collections/Iterator<T>|"];
+      subgraph cluster_27 {
+        color=blue
+        98 [label="Enter while loop"];
+        subgraph cluster_28 {
+          color=blue
+          99 [label="Enter loop condition"];
+          100 [label="Access variable R|<local>/<iterator>|"];
+          101 [label="Function call: R|<local>/<iterator>|.R|kotlin/collections/Iterator.hasNext|()"];
+          102 [label="Exit loop condition"];
+        }
+        subgraph cluster_29 {
+          color=blue
+          103 [label="Enter loop block"];
+          subgraph cluster_30 {
+            color=blue
+            104 [label="Enter block"];
+            105 [label="Access variable R|<local>/<iterator>|"];
+            106 [label="Function call: R|<local>/<iterator>|.R|FakeOverride<kotlin/collections/Iterator.next: R|T|>|()"];
+            107 [label="Variable declaration: lval element: R|T|"];
+            subgraph cluster_31 {
+              color=blue
+              108 [label="Enter when"];
+              subgraph cluster_32 {
+                color=blue
+                109 [label="Enter when branch condition "];
+                110 [label="Access variable R|<local>/element|"];
+                111 [label="Type operator: element is T"];
+                112 [label="Exit when branch condition"];
+              }
+              subgraph cluster_33 {
+                color=blue
+                113 [label="Enter when branch condition else"];
+                114 [label="Exit when branch condition"];
+              }
+              115 [label="Enter when branch result"];
+              subgraph cluster_34 {
+                color=blue
+                116 [label="Enter block"];
+                117 [label="Exit block"];
+              }
+              118 [label="Exit when branch result"];
+              119 [label="Enter when branch result"];
+              subgraph cluster_35 {
+                color=blue
+                120 [label="Enter block"];
+                121 [label="Access variable R|<local>/element|"];
+                122 [label="Jump: ^firstIsInstanceOrNull R|<local>/element|"];
+                123 [label="Stub" style="filled" fillcolor=gray];
+                124 [label="Exit block" style="filled" fillcolor=gray];
+              }
+              125 [label="Exit when branch result" style="filled" fillcolor=gray];
+              126 [label="Exit when"];
+            }
+            127 [label="Exit block"];
+          }
+          128 [label="Exit loop block"];
+        }
+        129 [label="Exit whileloop"];
+      }
+      130 [label="Const: Null(null)"];
+      131 [label="Jump: ^firstIsInstanceOrNull Null(null)"];
+      132 [label="Stub" style="filled" fillcolor=gray];
+      133 [label="Exit block" style="filled" fillcolor=gray];
+    }
+    134 [label="Exit function firstIsInstanceOrNull" style="filled" fillcolor=red];
+  }
+
+  91 -> {92};
+  92 -> {93};
+  93 -> {94};
+  94 -> {95};
+  95 -> {96};
+  96 -> {97};
+  97 -> {98};
+  98 -> {99};
+  99 -> {100};
+  100 -> {101};
+  101 -> {102};
+  102 -> {129 103};
+  103 -> {104};
+  104 -> {105};
+  105 -> {106};
+  106 -> {107};
+  107 -> {108};
+  108 -> {109};
+  109 -> {110};
+  110 -> {111};
+  111 -> {112};
+  112 -> {119 113};
+  113 -> {114};
+  114 -> {115};
+  115 -> {116};
+  116 -> {117};
+  117 -> {118};
+  118 -> {126};
+  119 -> {120};
+  120 -> {121};
+  121 -> {122};
+  122 -> {134};
+  122 -> {123} [style=dotted];
+  123 -> {124} [style=dotted];
+  124 -> {125} [style=dotted];
+  125 -> {126} [style=dotted];
+  126 -> {127};
+  127 -> {128};
+  128 -> {99};
+  129 -> {130};
+  130 -> {131};
+  131 -> {134};
+  131 -> {132} [style=dotted];
+  132 -> {133} [style=dotted];
+  133 -> {134} [style=dotted];
+
 }

compiler/fir/resolve/testData/resolve/cfg/complex.kt

@@ -27,3 +27,8 @@ internal fun AutoCloseable?.closeFinally(cause: Throwable?) = when {
             cause.addSuppressed(closeException)
         }
 }
+
+inline fun <reified T : Any> Sequence<*>.firstIsInstanceOrNull(): T? {
+    for (element in this) if (element is T) return element
+    return null
+}

compiler/fir/resolve/testData/resolve/cfg/complex.txt

@@ -38,3 +38,20 @@ FILE: complex.kt
         }
 
     }
+    public final inline fun <reified T : R|kotlin/Any|> R|kotlin/sequences/Sequence<*>|.firstIsInstanceOrNull(): R|T?| {
+        lval <range>: R|kotlin/sequences/Sequence<*>| = this@R|/firstIsInstanceOrNull|
+        lval <iterator>: R|kotlin/collections/Iterator<T>| = R|<local>/<range>|.R|kotlin/sequences/Sequence.iterator|()
+        while(R|<local>/<iterator>|.R|kotlin/collections/Iterator.hasNext|()) {
+            lval element: R|T| = R|<local>/<iterator>|.R|FakeOverride<kotlin/collections/Iterator.next: R|T|>|()
+            when () {
+                (R|<local>/element| is R|T|) ->  {
+                    ^firstIsInstanceOrNull R|<local>/element|
+                }
+                else ->  {
+                }
+            }
+
+        }
+
+        ^firstIsInstanceOrNull Null(null)
+    }

compiler/fir/resolve/testData/resolve/smartcasts/endlessLoops.dot

@@ -733,7 +733,7 @@ digraph endlessLoops_kt {
         244 [label="Exit do-whileloop" style="filled" fillcolor=gray];
       }
       245 [label="Access variable R|<local>/x|" style="filled" fillcolor=gray];
-      246 [label="Function call: R|<local>/x|.<Unresolved name: foo>#()" style="filled" fillcolor=gray];
+      246 [label="Function call: R|<local>/x|.R|/A.foo|()" style="filled" fillcolor=gray];
       247 [label="Exit block" style="filled" fillcolor=gray];
     }
     248 [label="Exit function test_7" style="filled" fillcolor=red style="filled" fillcolor=gray];

compiler/fir/resolve/testData/resolve/smartcasts/endlessLoops.kt

@@ -70,5 +70,5 @@ fun test_7(x: Any) {
     do {
         x as A
     } while (true)
-    x.foo() // No smartcast
+    x.foo()
 }

compiler/fir/resolve/testData/resolve/smartcasts/endlessLoops.txt

@@ -107,5 +107,5 @@ FILE: endlessLoops.kt
             (R|<local>/x| as R|A|)
         }
         while(Boolean(true))
-        R|<local>/x|.<Unresolved name: foo>#()
+        R|<local>/x|.R|/A.foo|()
     }