FIR CFA: compute subgraph relationships automatically

No more `addSubGraph`. Also no more `owner` in graphs.

^KT-40526 Obsolete
^KT-40582 Obsolete

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

@@ -142,21 +142,12 @@ interface UnionNodeMarker
 // ----------------------------------- EnterNode for declaration with CFG -----------------------------------
 
 sealed class CFGNodeWithSubgraphs<out E : FirElement>(owner: ControlFlowGraph, level: Int, id: Int) : CFGNode<E>(owner, level, id) {
-    private val _subGraphs = mutableListOf<ControlFlowGraph>()
-
-    fun addSubGraph(graph: ControlFlowGraph){
-        _subGraphs += graph
-    }
-
-    open val subGraphs: List<ControlFlowGraph>
-        get() = _subGraphs
+    abstract val subGraphs: List<ControlFlowGraph>
 }
 
 sealed class CFGNodeWithCfgOwner<out E : FirControlFlowGraphOwner>(owner: ControlFlowGraph, level: Int, id: Int) : CFGNodeWithSubgraphs<E>(owner, level, id) {
-    override val subGraphs: List<ControlFlowGraph> by lazy {
-        fir.controlFlowGraphReference?.controlFlowGraph?.run(::addSubGraph)
-        super.subGraphs
-    }
+    final override val subGraphs: List<ControlFlowGraph>
+        get() = listOfNotNull(fir.controlFlowGraphReference?.controlFlowGraph)
 }
 
 // ----------------------------------- Named function -----------------------------------
@@ -222,6 +213,10 @@ class ExitValueParameterNode(owner: ControlFlowGraph, override val fir: FirValue
 
 class SplitPostponedLambdasNode(owner: ControlFlowGraph, override val fir: FirStatement, val lambdas: List<FirAnonymousFunction>, level: Int, id: Int)
     : CFGNodeWithSubgraphs<FirStatement>(owner, level, id) {
+
+    override val subGraphs: List<ControlFlowGraph>
+        get() = lambdas.mapNotNull { it.controlFlowGraphReference?.controlFlowGraph }
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitSplitPostponedLambdasNode(this, data)
     }
@@ -240,6 +235,9 @@ class MergePostponedLambdaExitsNode(owner: ControlFlowGraph, override val fir: F
 }
 
 class AnonymousFunctionExpressionNode(owner: ControlFlowGraph, override val fir: FirAnonymousFunctionExpression, level: Int, id: Int) : CFGNodeWithSubgraphs<FirAnonymousFunctionExpression>(owner, level, id) {
+    override val subGraphs: List<ControlFlowGraph>
+        get() = listOfNotNull(fir.anonymousFunction.controlFlowGraphReference?.controlFlowGraph)
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitAnonymousFunctionExpressionNode(this, data)
     }
@@ -266,6 +264,8 @@ class ClassExitNode(owner: ControlFlowGraph, override val fir: FirClass, level:
         owner.exitNode = this
     }
 
+    lateinit override var subGraphs: List<ControlFlowGraph>
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitClassExitNode(this, data)
     }

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

@@ -26,14 +26,14 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
     @set:CfgInternals
     lateinit var exitNode: CFGNode<*>
 
-    var owner: ControlFlowGraph? = null
-        private set
+    val isSubGraph: Boolean
+        get() = enterNode.previousNodes.isNotEmpty()
 
     var state: State = State.Building
         private set
 
-    private val _subGraphs: MutableList<ControlFlowGraph> = mutableListOf()
-    val subGraphs: List<ControlFlowGraph> get() = _subGraphs
+    val subGraphs: List<ControlFlowGraph>
+        get() = _nodes.flatMap { (it as? CFGNodeWithSubgraphs<*>)?.subGraphs ?: emptyList() }
 
     @CfgInternals
     fun complete() {
@@ -50,15 +50,6 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
         _nodes.addAll(sortedNodes)
     }
 
-    @CfgInternals
-    fun addSubGraph(graph: ControlFlowGraph) {
-        assert(graph.owner == null) {
-            "SubGraph already has owner"
-        }
-        graph.owner = this
-        _subGraphs += graph
-    }
-
     private fun assertState(state: State) {
         assert(this.state == state) {
             "This action can not be performed at $this state"
@@ -171,34 +162,16 @@ private fun ControlFlowGraph.orderNodes(): LinkedHashSet<CFGNode<*>> {
             continue
         }
         if (!visitedNodes.add(node)) continue
-        val followingNodes = node.followingNodes
 
-        for (followingNode in followingNodes) {
+        // NOTE: this intentionally does not walk through subgraphs. If the only path from A to B
+        //  is through a subgraph, add a dead edge between them to enforce ordering.
+        for (followingNode in node.followingNodes) {
             if (followingNode.owner == this) {
                 if (followingNode !in visitedNodes) {
                     stack.addFirst(followingNode)
                 }
-            } else {
-                walkThrowSubGraphs(followingNode.owner, visitedNodes, stack)
             }
         }
     }
     return visitedNodes
 }
-
-private fun ControlFlowGraph.walkThrowSubGraphs(
-    otherGraph: ControlFlowGraph,
-    visitedNodes: Set<CFGNode<*>>,
-    stack: ArrayDeque<CFGNode<*>>
-) {
-    if (otherGraph.owner != this) return
-    for (otherNode in otherGraph.exitNode.followingNodes) {
-        if (otherNode.owner == this) {
-            if (otherNode !in visitedNodes) {
-                stack.addFirst(otherNode)
-            }
-        } else if (otherNode.owner != otherGraph) {
-            walkThrowSubGraphs(otherNode.owner, visitedNodes, stack)
-        }
-    }
-}