From c9bc5884dd60058ae5ac4d11d1bcf4d67abbc576 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Tue, 16 Jun 2020 22:27:40 +0300
Subject: [PATCH] [FIR] Add more utils for traversing control flow graph

---
 .../kotlin/fir/analysis/cfa/CfgTraverser.kt   | 78 +++++++++++++++++++
 .../kotlin/fir/analysis/cfa/CfgUtils.kt       | 49 ++++++------
 .../fir/analysis/cfa/ControlFlowInfo.kt       | 31 ++++++++
 3 files changed, 135 insertions(+), 23 deletions(-)
 create mode 100644 compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgTraverser.kt
 create mode 100644 compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/ControlFlowInfo.kt

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgTraverser.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgTraverser.kt
new file mode 100644
index 00000000000..81d4d05a09e
--- /dev/null
+++ b/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)
+        }
+    }
+}
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt
index 022db68b395..86bc6cf65cf 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt
+++ b/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)
-}
\ No newline at end of file
+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
+    }
\ No newline at end of file
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/ControlFlowInfo.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/ControlFlowInfo.kt
new file mode 100644
index 00000000000..f045b5bc0d1
--- /dev/null
+++ b/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))
+    }
+}
\ No newline at end of file