From 914acd841f9bac38ab200a1ecd31dc64f47b97dc Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Wed, 15 Feb 2023 18:50:04 +0200
Subject: [PATCH] [FIR] Add special kind for graph of init blocks and local
 functions

---
 .../kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt    | 5 +++--
 .../jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt | 2 ++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
index dabd55efa6e..4eb7a93b9bb 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
@@ -155,7 +155,8 @@ class ControlFlowGraphBuilder {
         val localFunctionNode = runIf(function is FirSimpleFunction && function.isLocal && bodyBuildingMode) {
             createLocalFunctionDeclarationNode(function).also { addNewSimpleNode(it) }
         }
-        val enterNode = enterGraph(function, name, ControlFlowGraph.Kind.Function) {
+        val kind = if (localFunctionNode != null) ControlFlowGraph.Kind.LocalFunction else ControlFlowGraph.Kind.Function
+        val enterNode = enterGraph(function, name, kind) {
             createFunctionEnterNode(it) to createFunctionExitNode(it).also { exit -> exitTargetsForReturn[it.symbol] = exit }
         }
         if (localFunctionNode != null) {
@@ -1151,7 +1152,7 @@ class ControlFlowGraphBuilder {
 
     fun enterInitBlock(initBlock: FirAnonymousInitializer): InitBlockEnterNode {
         // TODO: questionable moment that we should pass data flow from init to init
-        return enterGraph(initBlock, "init block", ControlFlowGraph.Kind.Function) {
+        return enterGraph(initBlock, "init block", ControlFlowGraph.Kind.ClassInitializer) {
             createInitBlockEnterNode(it) to createInitBlockExitNode(it)
         }.also { addEdgeIfLocalClassMember(it) }
     }
diff --git a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt
index f1a14ca88e3..e996455f937 100644
--- a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt
+++ b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt
@@ -34,9 +34,11 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
     enum class Kind {
         Class,
         Function,
+        LocalFunction,
         AnonymousFunction,
         AnonymousFunctionCalledInPlace,
         PropertyInitializer,
+        ClassInitializer,
         FieldInitializer,
         FakeCall,
         DefaultArgument,