From ad9e6517c7fe0c58c8f401b2732079596621a708 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 26 Mar 2020 11:35:28 +0300
Subject: [PATCH] [FIR] Add cfg nodes for classes

---
 .../fir/resolve/dfa/cfg/ControlFlowGraph.kt   | 20 ++++++++++++++++++-
 .../dfa/cfg/ControlFlowGraphNodeBuilder.kt    | 13 ++++++++++++
 .../dfa/cfg/ControlFlowGraphVisitor.kt        | 14 ++++++++++++-
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt
index 48feb694946..cd2f53aa9f9 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt
@@ -133,7 +133,25 @@ class UnionFunctionCallArgumentsNode(owner: ControlFlowGraph, override val fir:
     }
 }
 
-// ----------------------------------- Anonymous object -----------------------------------
+// ----------------------------------- Classes -----------------------------------
+
+class ClassEnterNode(owner: ControlFlowGraph, override val fir: FirClass<*>, level: Int, id: Int) : CFGNode<FirClass<*>>(owner, level, id) {
+    override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
+        return visitor.visitClassEnterNode(this, data)
+    }
+}
+
+class ClassExitNode(owner: ControlFlowGraph, override val fir: FirClass<*>, level: Int, id: Int) : CFGNode<FirClass<*>>(owner, level, id) {
+    override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
+        return visitor.visitClassExitNode(this, data)
+    }
+}
+
+class LocalClassExitNode(owner: ControlFlowGraph, override val fir: FirRegularClass, level: Int, id: Int) : CFGNode<FirRegularClass>(owner, level, id) {
+    override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
+        return visitor.visitLocalClassExitNode(this, data)
+    }
+}
 
 class AnonymousObjectExitNode(owner: ControlFlowGraph, override val fir: FirAnonymousObject, level: Int, id: Int) : CFGNode<FirAnonymousObject>(owner, level, id) {
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphNodeBuilder.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphNodeBuilder.kt
index a5dd1c71037..d95fb8168fd 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphNodeBuilder.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphNodeBuilder.kt
@@ -201,3 +201,16 @@ fun ControlFlowGraphBuilder.createAnonymousObjectExitNode(fir: FirAnonymousObjec
 
 fun ControlFlowGraphBuilder.createUnionFunctionCallArgumentsNode(fir: FirElement): UnionFunctionCallArgumentsNode =
     UnionFunctionCallArgumentsNode(graph, fir, levelCounter, createId())
+
+fun ControlFlowGraphBuilder.createClassEnterNode(fir: FirClass<*>): ClassEnterNode =
+    ClassEnterNode(graph, fir, levelCounter, createId()).also {
+        graph.enterNode = it
+    }
+
+fun ControlFlowGraphBuilder.createClassExitNode(fir: FirClass<*>): ClassExitNode =
+    ClassExitNode(graph, fir, levelCounter, createId()).also {
+        graph.exitNode = it
+    }
+
+fun ControlFlowGraphBuilder.createLocalClassExitNode(fir: FirRegularClass): LocalClassExitNode =
+    LocalClassExitNode(graph, fir, levelCounter, createId())
\ No newline at end of file
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphVisitor.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphVisitor.kt
index d228c81418d..8085ad3a827 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphVisitor.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphVisitor.kt
@@ -32,12 +32,24 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
         return visitNode(node, data)
     }
 
-    // ----------------------------------- Anonymous object -----------------------------------
+    // ----------------------------------- Classes -----------------------------------
 
     open fun visitAnonymousObjectExitNode(node: AnonymousObjectExitNode, data: D): R {
         return visitNode(node, data)
     }
 
+    open fun visitClassEnterNode(node: ClassEnterNode, data: D): R {
+        return visitNode(node, data)
+    }
+
+    open fun visitClassExitNode(node: ClassExitNode, data: D): R {
+        return visitNode(node, data)
+    }
+
+    open fun visitLocalClassExitNode(node: LocalClassExitNode, data: D): R {
+        return visitNode(node, data)
+    }
+
     // ----------------------------------- Property -----------------------------------
 
     open fun visitPropertyInitializerEnterNode(node: PropertyInitializerEnterNode, data: D): R {