[FIR] Some control flow graph building fixes

- fix edges for `continue` and labeled return
- fix deadness for init block and properties

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

@@ -349,7 +349,9 @@ class FirDataFlowAnalyzerImpl(transformer: FirBodyResolveTransformer) : FirDataF
     // ----------------------------------- While Loop -----------------------------------
 
     override fun enterWhileLoop(loop: FirLoop) {
-        graphBuilder.enterWhileLoop(loop).passFlow()
+        val (loopEnterNode, loopConditionEnterNode) = graphBuilder.enterWhileLoop(loop)
+        loopEnterNode.passFlow()
+        loopConditionEnterNode.passFlow()
     }
 
     override fun exitWhileLoopCondition(loop: FirLoop) {

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

@@ -70,8 +70,11 @@ class NodeStorage<T : FirElement, N : CFGNode<T>>(
     }
 }
 
-class SymbolBasedNodeStorage<T, N : CFGNode<T>> : Stack<N> where T : FirElement {
+class SymbolBasedNodeStorage<T, N : CFGNode<T>>(
-    private val stack: Stack<N> = stackOf()
+    pushCallback: ((N) -> Unit)? = null,
+    popCallback: ((N) -> Unit)? = null
+) : Stack<N> where T : FirElement {
+    private val stack: Stack<N> = StackImpl(pushCallback = pushCallback, popCallback = popCallback)
     private val map: MutableMap<FirBasedSymbol<*>, N> = mutableMapOf()
 
     override val size: Int get() = stack.size

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

@@ -25,10 +25,7 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
 
     private val exitNodes: Stack<CFGNode<*>> = stackOf()
 
-    private val functionExitNodes: NodeStorage<FirFunction<*>, FunctionExitNode> = NodeStorage(
+    private val functionExitNodes: SymbolBasedNodeStorage<FirFunction<*>, FunctionExitNode> = SymbolBasedNodeStorage()
-        pushCallback = { exitNodes.push(it) },
-        popCallback = { exitNodes.pop() }
-    )
 
     private val whenExitNodes: NodeStorage<FirWhenExpression, WhenExitNode> = NodeStorage()
 
@@ -88,6 +85,9 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
         }
 
         functionExitNodes.push(exitNode)
+        if (!isInplace) {
+            exitNodes.push(exitNode)
+        }
         levelCounter++
         return enterNode
     }
@@ -96,12 +96,16 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
         levelCounter--
         val exitNode = functionExitNodes.pop()
         val isInplace = function.isInplace()
+        if (!isInplace) {
+            exitNodes.pop()
+        }
         if (isInplace) {
             addNewSimpleNode(exitNode)
         } else {
             addEdge(lastNodes.pop(), exitNode)
             lexicalScopes.pop()
         }
+        exitNode.markAsDeadIfNecessary()
         val graph = if (!isInplace) {
             graphs.pop()
         } else {
@@ -136,8 +140,10 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
     }
 
     fun exitProperty(property: FirProperty): Pair<PropertyExitNode, ControlFlowGraph> {
-        val topLevelVariableExitNode = topLevelVariableExitNodes.pop()
+        val topLevelVariableExitNode = topLevelVariableExitNodes.pop().also {
-        addNewSimpleNode(topLevelVariableExitNode)
+            addNewSimpleNode(it)
+            it.markAsDeadIfNecessary()
+        }
         levelCounter--
         lexicalScopes.pop()
         return topLevelVariableExitNode to graphs.pop()
@@ -158,7 +164,7 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
     fun exitJump(jump: FirJump<*>): JumpNode {
         val node = createJumpNode(jump)
         val nextNode = when (jump) {
-            is FirReturnExpression -> functionExitNodes[jump.target.labeledElement]
+            is FirReturnExpression -> functionExitNodes[jump.target.labeledElement.symbol]
             is FirContinueExpression -> loopEnterNodes[jump.target.labeledElement]
             is FirBreakExpression -> loopExitNodes[jump.target.labeledElement]
             else -> throw IllegalArgumentException("Unknown jump type: ${jump.render()}")
@@ -213,17 +219,20 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
 
     // ----------------------------------- While Loop -----------------------------------
 
-    fun enterWhileLoop(loop: FirLoop): LoopConditionEnterNode {
+    fun enterWhileLoop(loop: FirLoop): Pair<LoopEnterNode, LoopConditionEnterNode> {
-        addNewSimpleNode(createLoopEnterNode(loop))
+        val loopEnterNode = createLoopEnterNode(loop).also {
+            addNewSimpleNode(it)
+            loopEnterNodes.push(it)
+        }
         loopExitNodes.push(createLoopExitNode(loop))
         levelCounter++
-        val node = createLoopConditionEnterNode(loop.condition)
+        val conditionEnterNode = createLoopConditionEnterNode(loop.condition).also {
+            addNewSimpleNode(it)
+            // put conditional node twice so we can refer it after exit from loop block
+            lastNodes.push(it)
+        }
         levelCounter++
-        addNewSimpleNode(node)
+        return loopEnterNode to conditionEnterNode
-        // put conditional node twice so we can refer it after exit from loop block
-        lastNodes.push(node)
-        loopEnterNodes.push(node)
-        return node
     }
 
     fun exitWhileLoopCondition(loop: FirLoop): LoopConditionExitNode {
@@ -509,6 +518,7 @@ class ControlFlowGraphBuilder : ControlFlowGraphNodeBuilder() {
         levelCounter--
         return initBlockExitNodes.pop().also {
             addNewSimpleNode(it)
+            it.markAsDeadIfNecessary()
             lexicalScopes.pop()
         }
     }

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

@@ -197,4 +197,106 @@ subgraph cluster_test_4 {
   88 -> {89};
 }
 
+subgraph cluster_test_5 {
+  90 [shape=box label="Enter function test_5" style="filled"];
+  91 [shape=box label="Enter block"];
+  92 [shape=box label="Enter while loop"];
+  93 [shape=box label="Enter loop condition"];
+  94 [shape=box label="Access variable R|<local>/b|"];
+  95 [shape=box label="Exit loop condition"];
+  96 [shape=box label="Enter loop block"];
+  97 [shape=box label="Enter block"];
+  98 [shape=box label="Enter when"];
+  99 [shape=box label="Enter when branch condition "];
+  100 [shape=box label="Access variable R|<local>/b|"];
+  101 [shape=box label="Exit when branch condition"];
+  102 [shape=box label="Enter block"];
+  103 [shape=box label="Jump: continue@@@[R|<local>/b|] "];
+  104 [shape=box label="Stub[DEAD]"];
+  105 [shape=box label="Exit block[DEAD]"];
+  106 [shape=box label="Exit when branch result[DEAD]"];
+  107 [shape=box label="Enter when branch condition else"];
+  108 [shape=box label="Exit when branch condition"];
+  109 [shape=box label="Enter block"];
+  110 [shape=box label="Exit block"];
+  111 [shape=box label="Exit when branch result"];
+  112 [shape=box label="Exit when"];
+  113 [shape=box label="Exit block"];
+  114 [shape=box label="Exit loop block"];
+  115 [shape=box label="Exit whileloop"];
+  116 [shape=box label="Exit block"];
+  117 [shape=box label="Exit function test_5" style="filled"];
+
+  90 -> {91};
+  91 -> {92};
+  92 -> {93};
+  93 -> {94};
+  94 -> {95};
+  95 -> {115 96};
+  96 -> {97};
+  97 -> {98};
+  98 -> {99};
+  99 -> {100};
+  100 -> {101};
+  101 -> {102 107};
+  102 -> {103};
+  103 -> {92};
+  103 -> {104} [style=dotted];
+  104 -> {105} [style=dotted];
+  105 -> {106} [style=dotted];
+  106 -> {112} [style=dotted];
+  107 -> {108};
+  108 -> {109};
+  109 -> {110};
+  110 -> {111};
+  111 -> {112};
+  112 -> {113};
+  113 -> {114};
+  114 -> {93};
+  115 -> {116};
+  116 -> {117};
+}
+
+subgraph cluster_run {
+  118 [shape=box label="Enter function run" style="filled"];
+  119 [shape=box label="Enter block"];
+  120 [shape=box label="Function call: R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
+  121 [shape=box label="Exit block"];
+  122 [shape=box label="Exit function run" style="filled"];
+
+  118 -> {119};
+  119 -> {120};
+  120 -> {121};
+  121 -> {122};
+}
+
+subgraph cluster_test_6 {
+  123 [shape=box label="Enter function test_6" style="filled"];
+  124 [shape=box label="Enter block"];
+  125 [shape=box label="Enter function anonymousFunction"];
+  126 [shape=box label="Enter block"];
+  127 [shape=box label="Jump: ^@run Unit"];
+  128 [shape=box label="Stub[DEAD]"];
+  129 [shape=box label="Exit block[DEAD]"];
+  130 [shape=box label="Exit function anonymousFunction"];
+  131 [shape=box label="Function call: R|/run|(<L> = run@fun <anonymous>(): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    ^@run Unit
+}
+)"];
+  132 [shape=box label="Exit block"];
+  133 [shape=box label="Exit function test_6" style="filled"];
+
+  123 -> {124};
+  124 -> {125};
+  125 -> {126};
+  126 -> {127};
+  127 -> {130};
+  127 -> {128} [style=dotted];
+  128 -> {129} [style=dotted];
+  129 -> {130} [style=dotted];
+  130 -> {131};
+  131 -> {132};
+  132 -> {133};
+}
+
 }

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

@@ -30,4 +30,22 @@ fun test_4(x: Int?) {
         break
     } while (true)
     x.inc()
+}
+
+fun test_5(b: Boolean) {
+    while (b) {
+        if (b) {
+            continue
+        }
+    }
+}
+
+fun run(block: () -> Unit) {
+    block()
+}
+
+fun test_6() {
+    run {
+        return@run
+    }
 }

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

@@ -39,3 +39,25 @@ FILE: jumps.kt
         while(Boolean(true))
         R|<local>/x|.R|kotlin/Int.inc|()
     }
+    public final fun test_5(b: R|kotlin/Boolean|): R|kotlin/Unit| {
+        while(R|<local>/b|) {
+            when () {
+                R|<local>/b| ->  {
+                    continue@@@[R|<local>/b|]
+                }
+                else ->  {
+                }
+            }
+
+        }
+
+    }
+    public final fun run(block: R|kotlin/Function0<kotlin/Unit>|): R|kotlin/Unit| {
+        R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun test_6(): R|kotlin/Unit| {
+        R|/run|(<L> = run@fun <anonymous>(): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+            ^@run Unit
+        }
+        )
+    }

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

@@ -30,12 +30,13 @@ subgraph cluster__getter_ {
   11 [shape=box label="Jump: ^ Int(1)"];
   12 [shape=box label="Stub[DEAD]"];
   13 [shape=box label="Exit block[DEAD]"];
-  14 [shape=box label="Exit function getter[DEAD]" style="filled"];
+  14 [shape=box label="Exit function getter" style="filled"];
 
   8 -> {9};
   9 -> {10};
   10 -> {11};
-  11 -> {14 12} [style=dotted];
+  11 -> {14};
+  11 -> {12} [style=dotted];
   12 -> {13} [style=dotted];
   13 -> {14} [style=dotted];
 }
@@ -75,7 +76,7 @@ subgraph cluster_foo {
   31 [shape=box label="Throw: throw <Ambiguity: Exception, [java/lang/Exception.Exception, java/lang/Exception.Exception]>#()"];
   32 [shape=box label="Stub[DEAD]"];
   33 [shape=box label="Exit block[DEAD]"];
-  34 [shape=box label="Exit function foo[DEAD]" style="filled"];
+  34 [shape=box label="Exit function foo" style="filled"];
 
   24 -> {25};
   25 -> {26};
@@ -84,7 +85,8 @@ subgraph cluster_foo {
   28 -> {29};
   29 -> {30};
   30 -> {31};
-  31 -> {34 32} [style=dotted];
+  31 -> {34};
+  31 -> {32} [style=dotted];
   32 -> {33} [style=dotted];
   33 -> {34} [style=dotted];
 }
@@ -103,7 +105,7 @@ subgraph cluster__getter_ {
   40 [shape=box label="Throw: throw <Ambiguity: Exception, [java/lang/Exception.Exception, java/lang/Exception.Exception]>#()"];
   41 [shape=box label="Stub[DEAD]"];
   42 [shape=box label="Exit block[DEAD]"];
-  43 [shape=box label="Exit init block[DEAD]"];
+  43 [shape=box label="Exit init block"];
   44 [shape=box label="Enter function getter" style="filled"];
   45 [shape=box label="Enter block"];
   46 [shape=box label="Exit block"];
@@ -112,7 +114,8 @@ subgraph cluster__getter_ {
   37 -> {38};
   38 -> {39};
   39 -> {40};
-  40 -> {43 41} [style=dotted];
+  40 -> {43};
+  40 -> {41} [style=dotted];
   41 -> {42} [style=dotted];
   42 -> {43} [style=dotted];
   44 -> {45};
@@ -135,7 +138,7 @@ subgraph cluster_val_x3 {
   54 [shape=box label="Stub[DEAD]"];
   55 [shape=box label="Const: Int(1)[DEAD]"];
   56 [shape=box label="Exit block[DEAD]"];
-  57 [shape=box label="Exit init block[DEAD]"];
+  57 [shape=box label="Exit init block"];
   58 [shape=box label="Enter property" style="filled"];
   59 [shape=box label="Enter function anonymousFunction"];
   60 [shape=box label="Enter block"];
@@ -165,12 +168,13 @@ subgraph cluster_val_x3 {
     throw <Ambiguity: Exception, [java/lang/Exception.Exception, java/lang/Exception.Exception]>#()
 }
 )[DEAD]"];
-  67 [shape=box label="Exit property[DEAD]" style="filled"];
+  67 [shape=box label="Exit property" style="filled"];
 
   50 -> {51};
   51 -> {52};
   52 -> {53};
-  53 -> {57 54} [style=dotted];
+  53 -> {57};
+  53 -> {54} [style=dotted];
   54 -> {55} [style=dotted];
   55 -> {56} [style=dotted];
   56 -> {57} [style=dotted];
@@ -178,7 +182,8 @@ subgraph cluster_val_x3 {
   59 -> {60};
   60 -> {61};
   61 -> {62};
-  62 -> {65 63} [style=dotted];
+  62 -> {67};
+  62 -> {63} [style=dotted];
   63 -> {64} [style=dotted];
   64 -> {65} [style=dotted];
   65 -> {66} [style=dotted];

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

@@ -218,6 +218,7 @@ subgraph cluster_test_3 {
   92 -> {99} [style=dotted];
   93 -> {109 94};
   94 -> {95};
+  95 -> {43};
   95 -> {96} [style=dotted];
   96 -> {97} [style=dotted];
   97 -> {98} [style=dotted];

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

@@ -13,7 +13,7 @@ FILE: elvis.kt
                     ^test_1 Unit
                 }
                 else ->  {
-                    R|<local>/<elvis>|!
+                    R|<local>/<elvis>|
                 }
             }
              ->  {

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

@@ -328,7 +328,7 @@ subgraph cluster_test_6 {
   146 [shape=box label="Access variable R|/Q.data|"];
   147 [shape=box label="Access variable R|/MyData.s|"];
   148 [shape=box label="Function call: R|<local>/q|?.R|/Q.data|?.R|/MyData.s|?.R|kotlin/Int.inc|()"];
-  149 [shape=box label="Variable declaration: lval <elvis>: R|kotlin/Int|"];
+  149 [shape=box label="Variable declaration: lval <elvis>: R|kotlin/Int?|"];
   150 [shape=box label="Enter when branch condition "];
   151 [shape=box label="Const: Null(null)"];
   152 [shape=box label="Operator =="];

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

@@ -49,7 +49,7 @@ FILE: nullability.kt
                 ^test_3 Unit
             }
             else ->  {
-                R|<local>/<elvis>|!
+                R|<local>/<elvis>|
             }
         }
 
@@ -79,12 +79,12 @@ FILE: nullability.kt
 
     }
     public final fun test_6(q: R|Q?|): R|kotlin/Unit| {
-        when (lval <elvis>: R|kotlin/Int| = R|<local>/q|?.R|/Q.data|?.R|/MyData.s|?.R|kotlin/Int.inc|()) {
+        when (lval <elvis>: R|kotlin/Int?| = R|<local>/q|?.R|/Q.data|?.R|/MyData.s|?.R|kotlin/Int.inc|()) {
             ==($subj$, Null(null)) ->  {
                 ^test_6 Unit
             }
             else ->  {
-                R|<local>/<elvis>|!
+                R|<local>/<elvis>|
             }
         }