Test: FIR CFA: index nodes in rendering order

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

@@ -32,10 +32,6 @@ class FirControlFlowGraphRenderVisitor(
 
     private var nodeCounter = 0
     private var clusterCounter = 0
-    private val indices = mutableMapOf<CFGNode<*>, Int>()
-
-    private val topLevelGraphs = mutableSetOf<ControlFlowGraph>()
-    private val allGraphs = mutableSetOf<ControlFlowGraph>()
 
     override fun visitFile(file: FirFile) {
         var name = file.name.replace(".", "_")
@@ -50,28 +46,14 @@ class FirControlFlowGraphRenderVisitor(
             .println("edge [penwidth=2]")
             .println()
         visitElement(file)
-
-        for (topLevelGraph in topLevelGraphs) {
-            printer.renderNodes(topLevelGraph)
-            printer.renderEdges(topLevelGraph)
-            printer.println()
-        }
-
         printer
             .popIndent()
             .println("}")
     }
 
-    private fun ControlFlowGraph.collectNodes() {
-        for (node in nodes) {
-            indices[node] = nodeCounter++
-        }
-    }
-
-    private fun Printer.renderNodes(graph: ControlFlowGraph) {
+    private fun Printer.renderNodes(nodes: Map<CFGNode<*>, Int>) {
         var color = RED
-        val sortedNodes = graph.sortedNodes()
-        for (node in sortedNodes) {
+        for ((node, index) in nodes) {
             if (node is EnterNodeMarker) {
                 enterCluster(color)
                 color = BLUE
@@ -94,7 +76,7 @@ class FirControlFlowGraphRenderVisitor(
                 attributes += "style=\"filled\""
                 attributes += "fillcolor=$it"
             }
-            println(indices.getValue(node), attributes.joinToString(separator = " ", prefix = " [", postfix = "];"))
+            println(index, attributes.joinToString(separator = " ", prefix = " [", postfix = "];"))
             if (node is ExitNodeMarker) {
                 exitCluster()
             }
@@ -116,26 +98,23 @@ class FirControlFlowGraphRenderVisitor(
             label.label?.let { "label=\"$it\"" }
         ).ifEmpty { null }?.joinToString(prefix = "[", separator = " ", postfix = "]")
 
-    private fun Printer.renderEdges(graph: ControlFlowGraph) {
-        for (node in graph.nodes) {
+    private fun Printer.renderEdges(nodes: Map<CFGNode<*>, Int>) {
+        for ((node, index) in nodes) {
             for ((style, group) in node.followingNodes.groupBy { node.edgeTo(it).style }.entries.sortedBy { it.key }) {
-                val mappedGroup = group.map { indices.getValue(it) }.sorted()
-                print(indices.getValue(node), EDGE, mappedGroup.joinToString(prefix = "{", postfix = "}", separator = " "))
+                val mappedGroup = group.map { nodes.getValue(it) }.sorted()
+                print(index, EDGE, mappedGroup.joinToString(prefix = "{", postfix = "}", separator = " "))
                 style?.let { printWithNoIndent(" $it") }
                 printlnWithNoIndent(";")
             }
 
             if (node is CFGNodeWithSubgraphs<*>) {
-                val subNodes = node.subGraphs.mapNotNull { indices[it.enterNode] }.sorted()
+                val subNodes = node.subGraphs.mapNotNull { nodes[it.enterNode] }.sorted()
                 if (subNodes.isNotEmpty()) {
-                    print(indices.getValue(node), EDGE, subNodes.joinToString(prefix = "{", postfix = "}", separator = " "))
+                    print(index, EDGE, subNodes.joinToString(prefix = "{", postfix = "}", separator = " "))
                     printlnWithNoIndent(" [style=dashed];")
                 }
             }
         }
-        for (subGraph in graph.subGraphs) {
-            renderEdges(subGraph)
-        }
     }
 
     override fun visitElement(element: FirElement) {
@@ -144,19 +123,15 @@ class FirControlFlowGraphRenderVisitor(
 
     override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference) {
         val controlFlowGraph = (controlFlowGraphReference as? FirControlFlowGraphReferenceImpl)?.controlFlowGraph ?: return
-        initializeNodes(controlFlowGraph)
-    }
+        if (controlFlowGraph.isSubGraph) return
 
-    private fun initializeNodes(graph: ControlFlowGraph) {
-        if (graph in allGraphs) return
-        graph.collectNodes()
-        if (!graph.isSubGraph) {
-            topLevelGraphs += graph
-        }
-        allGraphs += graph
-        for (subGraph in graph.subGraphs) {
-            initializeNodes(subGraph)
-        }
+        // TODO: nodes are already in a topological order, but grouping nodes into clusters requires something more.
+        //  But what exactly? And is there a way to do `renderNodes` differently so that any topological order is ok?
+        val nodes = DFS.topologicalOrder(listOf(controlFlowGraph.enterNode)) { it.followingNodes }
+            .associateWithTo(linkedMapOf()) { nodeCounter++ }
+        printer.renderNodes(nodes)
+        printer.renderEdges(nodes)
+        printer.println()
     }
 
     private fun Printer.enterCluster(color: String) {
@@ -170,29 +145,3 @@ class FirControlFlowGraphRenderVisitor(
         println("}")
     }
 }
-
-private fun ControlFlowGraph.sortedNodes(): List<CFGNode<*>> {
-    val nodesToSort = nodes.filterTo(mutableListOf()) { it != enterNode }
-    val graphs = mutableSetOf(this)
-    forEachSubGraph {
-        nodesToSort += it.nodes
-        graphs += it
-    }
-
-    val topologicalOrder = DFS.topologicalOrder(nodesToSort) {
-        val result = if (it !is WhenBranchConditionExitNode || it.followingNodes.size < 2) {
-            it.followingNodes
-        } else {
-            it.followingNodes.sortedBy { node -> if (node is BlockEnterNode) 1 else 0 }
-        }.filter { node -> node.owner in graphs }
-        result
-    }
-    return listOf(enterNode) + topologicalOrder
-}
-
-private fun ControlFlowGraph.forEachSubGraph(block: (ControlFlowGraph) -> Unit) {
-    for (subGraph in subGraphs) {
-        block(subGraph)
-        subGraph.forEachSubGraph(block)
-    }
-}