[FIR] Add different states for CFG and some assertions for graph modification

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

@@ -3,6 +3,8 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
+@file:Suppress("Reformat")
+
 package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 
 import org.jetbrains.kotlin.fir.FirElement

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

@@ -3,8 +3,6 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-@file:Suppress("Reformat")
-
 package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
@@ -15,20 +13,29 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
     val nodes: List<CFGNode<*>> get() = _nodes
 
     internal fun addNode(node: CFGNode<*>) {
+        assertState(State.Building)
         _nodes += node
     }
 
     lateinit var enterNode: CFGNode<*>
         internal set
+
     lateinit var exitNode: CFGNode<*>
         internal set
-
     var owner: ControlFlowGraph? = null
         private set
 
+    var state: State = State.Building
+        private set
+
     private val _subGraphs: MutableList<ControlFlowGraph> = mutableListOf()
     val subGraphs: List<ControlFlowGraph> get() = _subGraphs
 
+    internal fun complete() {
+        assertState(State.Building)
+        state = State.Completed
+    }
+
     internal fun addSubGraph(graph: ControlFlowGraph) {
         assert(graph.owner == null) {
             "SubGraph already has owner"
@@ -46,6 +53,17 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
         CFGNode.removeAllOutgoingEdges(graph.exitNode)
     }
 
+    private fun assertState(state: State) {
+        assert(this.state == state) {
+            "This action can not be performed at $this state"
+        }
+    }
+
+    enum class State {
+        Building,
+        Completed;
+    }
+
     enum class Kind {
         Function, ClassInitializer, TopLevel
     }

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

@@ -173,6 +173,7 @@ class ControlFlowGraphBuilder {
         exitNode.updateDeadStatus()
         val graph = graphs.pop().also { graph ->
             exitsFromCompletedPostponedAnonymousFunctions.removeAll { it.owner == graph }
+            graph.complete()
         }
         val previousGraph = parentGraphForAnonymousFunctions.remove(function.symbol) ?: graphs.top()
         if (previousGraph.kind == ControlFlowGraph.Kind.Function) {
@@ -252,7 +253,7 @@ class ControlFlowGraphBuilder {
         }
         val exitNode = createClassExitNode(klass)
         addEdge(node, exitNode, preferredKind = EdgeKind.Cfg)
-        return graphs.pop()
+        return graphs.pop().also { it.complete() }
     }
 
     fun prepareForLocalClassMembers(members: Collection<FirSymbolOwner<*>>) {
@@ -361,7 +362,7 @@ class ControlFlowGraphBuilder {
         levelCounter--
         exitNodes.pop()
         lexicalScopes.pop()
-        return topLevelVariableExitNode to graphs.pop()
+        return topLevelVariableExitNode to graphs.pop().also { it.complete() }
     }
 
     // ----------------------------------- Delegate -----------------------------------
@@ -854,7 +855,7 @@ class ControlFlowGraphBuilder {
         }
         lexicalScopes.pop()
         exitNodes.pop()
-        val graph = graphs.pop()
+        val graph = graphs.pop().also { it.complete() }
         return exitNode to graph
     }