[FIR] Add more utils for traversing control flow graph

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgTraverser.kt

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.cfa
+
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
+
+enum class TraverseDirection {
+    Forward, Backward
+}
+
+fun <D> ControlFlowGraph.traverse(
+    direction: TraverseDirection,
+    visitor: ControlFlowGraphVisitor<*, D>,
+    data: D
+) {
+    for (node in getNodesInOrder(direction)) {
+        node.accept(visitor, data)
+        if (node is CFGNodeWithCfgOwner<*>) {
+            node.subGraph?.traverse(direction, visitor, data)
+        }
+    }
+}
+
+fun ControlFlowGraph.traverse(
+    direction: TraverseDirection,
+    visitor: ControlFlowGraphVisitorVoid
+) {
+    traverse(direction, visitor, null)
+}
+
+fun <I : ControlFlowInfo<I, K, V>, K : Any, V : Any> ControlFlowGraph.collectDataForNode(
+    direction: TraverseDirection,
+    initialInfo: I,
+    visitor: ControlFlowGraphVisitor<I, Collection<I>>
+): Map<CFGNode<*>, I> {
+    val nodeMap = LinkedHashMap<CFGNode<*>, I>()
+    val startNode = getEnterNode(direction)
+    nodeMap[startNode] = initialInfo
+
+    val changed = mutableMapOf<CFGNode<*>, Boolean>()
+    do {
+        collectDataForNodeInternal(direction, initialInfo, visitor, nodeMap, changed)
+    } while (changed.any { it.value })
+
+    return nodeMap
+}
+
+private fun <I : ControlFlowInfo<I, K, V>, K : Any, V : Any> ControlFlowGraph.collectDataForNodeInternal(
+    direction: TraverseDirection,
+    initialInfo: I,
+    visitor: ControlFlowGraphVisitor<I, Collection<I>>,
+    nodeMap: MutableMap<CFGNode<*>, I>,
+    changed: MutableMap<CFGNode<*>, Boolean>
+) {
+    val nodes = getNodesInOrder(direction)
+    for (node in nodes) {
+        if (!(node.isEnterNode(direction) && node.owner.owner == null)) {
+            val nextNodes = when (direction) {
+                TraverseDirection.Forward -> node.previousCfgNodes
+                TraverseDirection.Backward -> node.followingCfgNodes
+            }
+            val previousData = nextNodes.mapNotNull { nodeMap[it] }
+            val data = nodeMap[node]
+            val newData = node.accept(visitor, previousData)
+            val hasChanged = newData != data
+            changed[node] = hasChanged
+            if (hasChanged) {
+                nodeMap[node] = newData
+            }
+        }
+        if (node is CFGNodeWithCfgOwner<*>) {
+            node.subGraph?.collectDataForNodeInternal(direction, initialInfo, visitor, nodeMap, changed)
+        }
+    }
+}

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt

@@ -7,31 +7,34 @@ package org.jetbrains.kotlin.fir.analysis.cfa
 
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraphVisitor
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraphVisitorVoid
 
-enum class TraverseDirection {
-    Forward, Backward
+fun ControlFlowGraph.getEnterNode(direction: TraverseDirection): CFGNode<*> = when (direction) {
+    TraverseDirection.Forward -> enterNode
+    TraverseDirection.Backward -> exitNode
 }
 
-@OptIn(ExperimentalStdlibApi::class)
-fun <D> ControlFlowGraph.traverse(
-    direction: TraverseDirection,
-    visitor: ControlFlowGraphVisitor<*, D>,
-    data: D
-) {
-    val nodes = when (direction) {
-        TraverseDirection.Forward -> nodes
-        TraverseDirection.Backward -> nodes.reversed()
-    }
-    for (node in nodes) {
-        node.accept(visitor, data)
-    }
+fun ControlFlowGraph.getNodesInOrder(direction: TraverseDirection): List<CFGNode<*>> = when (direction) {
+    TraverseDirection.Forward -> nodes
+    TraverseDirection.Backward -> nodes.asReversed()
 }
 
-fun ControlFlowGraph.traverse(
-    direction: TraverseDirection,
-    visitor: ControlFlowGraphVisitorVoid
-) {
-    traverse(direction, visitor, null)
-}
+fun CFGNode<*>.isEnterNode(direction: TraverseDirection): Boolean = when (direction) {
+    TraverseDirection.Forward -> owner.enterNode == this
+    TraverseDirection.Backward -> owner.exitNode == this
+}
+
+val CFGNode<*>.previousCfgNodes: List<CFGNode<*>>
+    get() = previousNodes.filter {
+        val kind = incomingEdges.getValue(it)
+        if (this.isDead) {
+            kind.usedInCfa
+        } else {
+            kind.usedInCfa && !kind.isDead
+        }
+    }
+
+val CFGNode<*>.followingCfgNodes: List<CFGNode<*>>
+    get() = followingNodes.filter {
+        val kind = outgoingEdges.getValue(it)
+        kind.usedInCfa && !kind.isDead
+    }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/ControlFlowInfo.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.cfa
+
+import kotlinx.collections.immutable.PersistentMap
+
+abstract class ControlFlowInfo<S : ControlFlowInfo<S, K, V>, K : Any, V : Any> protected constructor(
+    protected val map: PersistentMap<K, V>,
+) : PersistentMap<K, V> by map {
+
+    protected abstract val constructor: (PersistentMap<K, V>) -> S
+
+    override fun equals(other: Any?): Boolean {
+        return map == (other as? ControlFlowInfo<*, *, *>)?.map
+    }
+
+    override fun hashCode(): Int {
+        return map.hashCode()
+    }
+
+    override fun toString(): String {
+        return map.toString()
+    }
+
+    override fun put(key: K, value: V): S {
+        return constructor(map.put(key, value))
+    }
+}