[FIR] Add control flow graph to FirAnonymousInitializer

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

@@ -15,12 +15,61 @@ digraph initBlock_kt {
 
     subgraph cluster_1 {
         color=red
-        3 [label="Enter function <init>" style="filled" fillcolor=red];
+        3 [label="Enter init block" style="filled" fillcolor=red];
-        4 [label="Delegated constructor call: super<R|kotlin/Any|>()"];
+        subgraph cluster_2 {
-        5 [label="Exit function <init>" style="filled" fillcolor=red];
+            color=blue
+            4 [label="Enter block"];
+            5 [label="Const: Int(1)"];
+            6 [label="Variable declaration: lval x: R|kotlin/Int|"];
+            7 [label="Exit block"];
+        }
+        8 [label="Exit init block" style="filled" fillcolor=red];
     }
 
     3 -> {4};
     4 -> {5};
+    5 -> {6};
+    6 -> {7};
+    7 -> {8};
+
+    subgraph cluster_3 {
+        color=red
+        9 [label="Enter function <init>" style="filled" fillcolor=red];
+        10 [label="Delegated constructor call: super<R|kotlin/Any|>()"];
+        11 [label="Exit function <init>" style="filled" fillcolor=red];
+    }
+
+    9 -> {10};
+    10 -> {11};
+
+    subgraph cluster_4 {
+        color=red
+        12 [label="Enter init block" style="filled" fillcolor=red];
+        subgraph cluster_5 {
+            color=blue
+            13 [label="Enter block"];
+            14 [label="Const: Int(1)"];
+            15 [label="Variable declaration: lval x: R|kotlin/Int|"];
+            16 [label="Function call: R|java/lang/Exception.Exception|()"];
+            17 [label="Throw: throw R|java/lang/Exception.Exception|()"];
+            18 [label="Stub" style="filled" fillcolor=gray];
+            19 [label="Const: Int(2)" style="filled" fillcolor=gray];
+            20 [label="Variable declaration: lval y: R|kotlin/Int|" style="filled" fillcolor=gray];
+            21 [label="Exit block" style="filled" fillcolor=gray];
+        }
+        22 [label="Exit init block" style="filled" fillcolor=red];
+    }
+
+    12 -> {13};
+    13 -> {14};
+    14 -> {15};
+    15 -> {16};
+    16 -> {17};
+    17 -> {22};
+    17 -> {18} [style=dotted];
+    18 -> {19} [style=dotted];
+    19 -> {20} [style=dotted];
+    20 -> {21} [style=dotted];
+    21 -> {22} [style=dotted];
 
 }

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

@@ -29,4 +29,52 @@ digraph initBlockAndInPlaceLambda_kt {
     4 -> {5};
     5 -> {6};
 
+    subgraph cluster_3 {
+        color=red
+        7 [label="Enter init block" style="filled" fillcolor=red];
+        subgraph cluster_4 {
+            color=blue
+            8 [label="Enter block"];
+            9 [label="Access variable R|<local>/a|"];
+            10 [label="Access variable R|/A.b|"];
+            11 [label="Enter safe call"];
+            12 [label="Postponed enter to lambda"];
+            subgraph cluster_5 {
+                color=blue
+                13 [label="Enter function anonymousFunction"];
+                14 [label="Access variable R|<local>/a|"];
+                15 [label="Access variable R|<local>/it|"];
+                16 [label="Function call: R|/C.C|(...)"];
+                17 [label="Exit function anonymousFunction"];
+            }
+            18 [label="Call arguments union" style="filled" fillcolor=yellow];
+            19 [label="Postponed exit from lambda"];
+            20 [label="Function call: R|<local>/a|.R|/A.b|?.R|kotlin/let|<R|B|, R|C|>(...)"];
+            21 [label="Exit safe call"];
+            22 [label="Variable declaration: lval c: R|C?|"];
+            23 [label="Exit block"];
+        }
+        24 [label="Exit init block" style="filled" fillcolor=red];
+    }
+
+    7 -> {8};
+    8 -> {9};
+    9 -> {10};
+    10 -> {11 21};
+    11 -> {12};
+    12 -> {13};
+    12 -> {19} [color=red];
+    13 -> {14};
+    14 -> {15};
+    15 -> {16};
+    16 -> {17};
+    17 -> {19} [color=green];
+    17 -> {18} [color=red];
+    18 -> {20} [color=red];
+    19 -> {20} [color=green];
+    20 -> {21};
+    21 -> {22};
+    22 -> {23};
+    23 -> {24};
+
 }

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

@@ -103,121 +103,167 @@ digraph propertiesAndInitBlocks_kt {
 
     subgraph cluster_8 {
         color=red
-        32 [label="Enter function getter" style="filled" fillcolor=red];
+        32 [label="Enter init block" style="filled" fillcolor=red];
-        33 [label="Exit function getter" style="filled" fillcolor=red];
+        subgraph cluster_9 {
+            color=blue
+            33 [label="Enter block"];
+            34 [label="Function call: R|java/lang/Exception.Exception|()"];
+            35 [label="Throw: throw R|java/lang/Exception.Exception|()"];
+            36 [label="Stub" style="filled" fillcolor=gray];
+            37 [label="Const: Int(1)" style="filled" fillcolor=gray];
+            38 [label="Exit block" style="filled" fillcolor=gray];
+        }
+        39 [label="Exit init block" style="filled" fillcolor=red];
     }
 
     32 -> {33};
-
+    33 -> {34};
-    subgraph cluster_9 {
-        color=red
-        34 [label="Enter function <init>" style="filled" fillcolor=red];
-        35 [label="Delegated constructor call: super<R|kotlin/Any|>()"];
-        36 [label="Exit function <init>" style="filled" fillcolor=red];
-    }
-
     34 -> {35};
-    35 -> {36};
+    35 -> {39};
+    35 -> {36} [style=dotted];
+    36 -> {37} [style=dotted];
+    37 -> {38} [style=dotted];
+    38 -> {39} [style=dotted];
 
     subgraph cluster_10 {
         color=red
-        37 [label="Enter property" style="filled" fillcolor=red];
+        40 [label="Enter function getter" style="filled" fillcolor=red];
-        38 [label="Postponed enter to lambda"];
+        41 [label="Exit function getter" style="filled" fillcolor=red];
-        subgraph cluster_11 {
-            color=blue
-            39 [label="Enter function anonymousFunction"];
-            40 [label="Function call: R|java/lang/Exception.Exception|()"];
-            41 [label="Throw: throw R|java/lang/Exception.Exception|()"];
-            42 [label="Stub" style="filled" fillcolor=gray];
-            43 [label="Exit function anonymousFunction"];
-        }
-        44 [label="Postponed exit from lambda"];
-        45 [label="Function call: R|/run|(...)"];
-        46 [label="Exit property" style="filled" fillcolor=red];
     }
 
-    37 -> {38};
-    38 -> {39};
-    38 -> {44} [color=red];
-    39 -> {43 40};
     40 -> {41};
-    41 -> {46};
+
-    41 -> {42} [style=dotted];
+    subgraph cluster_11 {
-    42 -> {43} [style=dotted];
+        color=red
-    43 -> {39};
+        42 [label="Enter function <init>" style="filled" fillcolor=red];
-    43 -> {44} [color=green];
+        43 [label="Delegated constructor call: super<R|kotlin/Any|>()"];
-    44 -> {45};
+        44 [label="Exit function <init>" style="filled" fillcolor=red];
-    45 -> {46};
+    }
+
+    42 -> {43};
+    43 -> {44};
 
     subgraph cluster_12 {
         color=red
-        47 [label="Enter function getter" style="filled" fillcolor=red];
+        45 [label="Enter init block" style="filled" fillcolor=red];
-        48 [label="Exit function getter" style="filled" fillcolor=red];
+        subgraph cluster_13 {
-    }
-
-    47 -> {48};
-
-    subgraph cluster_13 {
-        color=red
-        49 [label="Enter property" style="filled" fillcolor=red];
-        subgraph cluster_14 {
             color=blue
-            50 [label="Try expression enter"];
+            46 [label="Enter block"];
-            subgraph cluster_15 {
+            47 [label="Function call: R|java/lang/Exception.Exception|()"];
-                color=blue
+            48 [label="Throw: throw R|java/lang/Exception.Exception|()"];
-                51 [label="Try main block enter"];
+            49 [label="Stub" style="filled" fillcolor=gray];
-                subgraph cluster_16 {
+            50 [label="Exit block" style="filled" fillcolor=gray];
-                    color=blue
-                    52 [label="Enter block"];
-                    53 [label="Const: Int(1)"];
-                    54 [label="Exit block"];
-                }
-                55 [label="Try main block exit"];
-            }
-            subgraph cluster_17 {
-                color=blue
-                56 [label="Enter finally"];
-                subgraph cluster_18 {
-                    color=blue
-                    57 [label="Enter block"];
-                    58 [label="Const: Int(0)"];
-                    59 [label="Exit block"];
-                }
-                60 [label="Exit finally"];
-            }
-            subgraph cluster_19 {
-                color=blue
-                61 [label="Catch enter"];
-                subgraph cluster_20 {
-                    color=blue
-                    62 [label="Enter block"];
-                    63 [label="Const: Int(2)"];
-                    64 [label="Exit block"];
-                }
-                65 [label="Catch exit"];
-            }
-            66 [label="Try expression exit"];
         }
-        67 [label="Exit property" style="filled" fillcolor=red];
+        51 [label="Exit init block" style="filled" fillcolor=red];
+    }
+
+    45 -> {46};
+    46 -> {47};
+    47 -> {48};
+    48 -> {51};
+    48 -> {49} [style=dotted];
+    49 -> {50} [style=dotted];
+    50 -> {51} [style=dotted];
+
+    subgraph cluster_14 {
+        color=red
+        52 [label="Enter property" style="filled" fillcolor=red];
+        53 [label="Postponed enter to lambda"];
+        subgraph cluster_15 {
+            color=blue
+            54 [label="Enter function anonymousFunction"];
+            55 [label="Function call: R|java/lang/Exception.Exception|()"];
+            56 [label="Throw: throw R|java/lang/Exception.Exception|()"];
+            57 [label="Stub" style="filled" fillcolor=gray];
+            58 [label="Exit function anonymousFunction"];
+        }
+        59 [label="Postponed exit from lambda"];
+        60 [label="Function call: R|/run|(...)"];
+        61 [label="Exit property" style="filled" fillcolor=red];
     }
 
-    49 -> {50};
-    50 -> {51};
-    51 -> {67 61 56 52};
     52 -> {53};
     53 -> {54};
-    54 -> {55};
+    53 -> {59} [color=red];
-    55 -> {66};
+    54 -> {58 55};
-    56 -> {57};
+    55 -> {56};
-    57 -> {58};
+    56 -> {61};
-    58 -> {59};
+    56 -> {57} [style=dotted];
+    57 -> {58} [style=dotted];
+    58 -> {54};
+    58 -> {59} [color=green];
     59 -> {60};
-    60 -> {66};
+    60 -> {61};
-    61 -> {67 62};
+
+    subgraph cluster_16 {
+        color=red
+        62 [label="Enter function getter" style="filled" fillcolor=red];
+        63 [label="Exit function getter" style="filled" fillcolor=red];
+    }
+
     62 -> {63};
-    63 -> {64};
+
+    subgraph cluster_17 {
+        color=red
+        64 [label="Enter property" style="filled" fillcolor=red];
+        subgraph cluster_18 {
+            color=blue
+            65 [label="Try expression enter"];
+            subgraph cluster_19 {
+                color=blue
+                66 [label="Try main block enter"];
+                subgraph cluster_20 {
+                    color=blue
+                    67 [label="Enter block"];
+                    68 [label="Const: Int(1)"];
+                    69 [label="Exit block"];
+                }
+                70 [label="Try main block exit"];
+            }
+            subgraph cluster_21 {
+                color=blue
+                71 [label="Enter finally"];
+                subgraph cluster_22 {
+                    color=blue
+                    72 [label="Enter block"];
+                    73 [label="Const: Int(0)"];
+                    74 [label="Exit block"];
+                }
+                75 [label="Exit finally"];
+            }
+            subgraph cluster_23 {
+                color=blue
+                76 [label="Catch enter"];
+                subgraph cluster_24 {
+                    color=blue
+                    77 [label="Enter block"];
+                    78 [label="Const: Int(2)"];
+                    79 [label="Exit block"];
+                }
+                80 [label="Catch exit"];
+            }
+            81 [label="Try expression exit"];
+        }
+        82 [label="Exit property" style="filled" fillcolor=red];
+    }
+
     64 -> {65};
     65 -> {66};
-    66 -> {67};
+    66 -> {82 76 71 67};
+    67 -> {68};
+    68 -> {69};
+    69 -> {70};
+    70 -> {81};
+    71 -> {72};
+    72 -> {73};
+    73 -> {74};
+    74 -> {75};
+    75 -> {81};
+    76 -> {82 77};
+    77 -> {78};
+    78 -> {79};
+    79 -> {80};
+    80 -> {81};
+    81 -> {82};
 
 }

compiler/fir/analysis-tests/testData/resolve/smartcasts/smartCastInInit.dot

@@ -54,4 +54,27 @@ digraph smartCastInInit_kt {
 
     13 -> {14};
 
+    subgraph cluster_5 {
+        color=red
+        15 [label="Enter init block" style="filled" fillcolor=red];
+        subgraph cluster_6 {
+            color=blue
+            16 [label="Enter block"];
+            17 [label="Function call: R|/s|()"];
+            18 [label="Assignmenet: R|/Main.x|"];
+            19 [label="Access variable R|/Main.x|"];
+            20 [label="Function call: this@R|/Main|.R|/Main.x|.R|/S.foo|()"];
+            21 [label="Exit block"];
+        }
+        22 [label="Exit init block" style="filled" fillcolor=red];
+    }
+
+    15 -> {16};
+    16 -> {17};
+    17 -> {18};
+    18 -> {19};
+    19 -> {20};
+    20 -> {21};
+    21 -> {22};
+
 }

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

@@ -924,8 +924,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    fun exitInitBlock(initBlock: FirAnonymousInitializer) {
+    fun exitInitBlock(initBlock: FirAnonymousInitializer): ControlFlowGraph {
-        graphBuilder.exitInitBlock(initBlock).mergeIncomingFlow()
+        val (node, controlFlowGraph) = graphBuilder.exitInitBlock(initBlock)
+        node.mergeIncomingFlow()
+        return controlFlowGraph
     }
 
     // ------------------------------------------------------ Utils ------------------------------------------------------

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

@@ -772,15 +772,16 @@ class ControlFlowGraphBuilder {
         return enterNode to lastNode
     }
 
-    fun exitInitBlock(initBlock: FirAnonymousInitializer): InitBlockExitNode {
+    fun exitInitBlock(initBlock: FirAnonymousInitializer): Pair<InitBlockExitNode, ControlFlowGraph> {
         levelCounter--
-        return initBlockExitNodes.pop().also {
+        val exitNode = initBlockExitNodes.pop().also {
             addNewSimpleNode(it)
             it.updateDeadStatus()
-            lexicalScopes.pop()
-            exitNodes.pop()
-            graphs.pop()
         }
+        lexicalScopes.pop()
+        exitNodes.pop()
+        val graph = graphs.pop()
+        return exitNode to graph
     }
 
     // ----------------------------------- Safe calls -----------------------------------

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

@@ -15,11 +15,15 @@ fun ControlFlowGraphBuilder.createLoopExitNode(fir: FirLoop): LoopExitNode = Loo
 
 fun ControlFlowGraphBuilder.createLoopEnterNode(fir: FirLoop): LoopEnterNode = LoopEnterNode(graph, fir, levelCounter, createId())
 
-fun ControlFlowGraphBuilder.createInitBlockExitNode(fir: FirAnonymousInitializer): InitBlockExitNode =
-    InitBlockExitNode(graph, fir, levelCounter, createId())
-
 fun ControlFlowGraphBuilder.createInitBlockEnterNode(fir: FirAnonymousInitializer): InitBlockEnterNode =
-    InitBlockEnterNode(graph, fir, levelCounter, createId())
+    InitBlockEnterNode(graph, fir, levelCounter, createId()).also {
+        graph.enterNode = it
+    }
+
+fun ControlFlowGraphBuilder.createInitBlockExitNode(fir: FirAnonymousInitializer): InitBlockExitNode =
+    InitBlockExitNode(graph, fir, levelCounter, createId()).also {
+        graph.exitNode = it
+    }
 
 fun ControlFlowGraphBuilder.createTypeOperatorCallNode(fir: FirTypeOperatorCall): TypeOperatorCallNode =
     TypeOperatorCallNode(graph, fir, levelCounter, createId())

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -442,9 +442,9 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             dataFlowAnalyzer.enterInitBlock(anonymousInitializer)
             addLocalScope(primaryConstructorParametersScope)
             addLocalScope(FirLocalScope())
-            transformDeclarationContent(anonymousInitializer, ResolutionMode.ContextIndependent).also {
+            val result = transformDeclarationContent(anonymousInitializer, ResolutionMode.ContextIndependent).single as FirAnonymousInitializer
-                dataFlowAnalyzer.exitInitBlock(it.single as FirAnonymousInitializer)
+            val graph = dataFlowAnalyzer.exitInitBlock(result)
-            }
+            result.transformControlFlowGraphReference(ControlFlowGraphReferenceTransformer, graph).compose()
         }
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirAnonymousInitializer.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -21,6 +22,9 @@ abstract class FirAnonymousInitializer : FirPureAbstractElement(), FirDeclaratio
     abstract override val session: FirSession
     abstract override val resolvePhase: FirResolvePhase
     abstract val body: FirBlock?
+    abstract val controlFlowGraphReference: FirControlFlowGraphReference
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitAnonymousInitializer(this, data)
+
+    abstract fun <D> transformControlFlowGraphReference(transformer: FirTransformer<D>, data: D): FirAnonymousInitializer
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirAnonymousInitializerBuilder.kt

@@ -13,6 +13,8 @@ import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.impl.FirAnonymousInitializerImpl
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
+import org.jetbrains.kotlin.fir.references.impl.FirEmptyControlFlowGraphReference
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirAnonymousInitializerImpl.kt

@@ -10,6 +10,8 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
+import org.jetbrains.kotlin.fir.references.impl.FirEmptyControlFlowGraphReference
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -23,12 +25,21 @@ internal class FirAnonymousInitializerImpl(
     override var resolvePhase: FirResolvePhase,
     override var body: FirBlock?,
 ) : FirAnonymousInitializer() {
+    override var controlFlowGraphReference: FirControlFlowGraphReference = FirEmptyControlFlowGraphReference
+
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         body?.accept(visitor, data)
+        controlFlowGraphReference.accept(visitor, data)
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirAnonymousInitializerImpl {
         body = body?.transformSingle(transformer, data)
+        transformControlFlowGraphReference(transformer, data)
+        return this
+    }
+
+    override fun <D> transformControlFlowGraphReference(transformer: FirTransformer<D>, data: D): FirAnonymousInitializerImpl {
+        controlFlowGraphReference = controlFlowGraphReference.transformSingle(transformer, data)
         return this
     }
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -367,6 +367,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
 
         anonymousInitializer.configure {
             +body(nullable = true)
+            +controlFlowGraphReferenceField
         }
 
         file.configure {