From c4b8e9406192f3aa1a7054a2d5bb7f97aaeba12a Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 19 Mar 2020 16:04:26 +0300
Subject: [PATCH] [FIR] Add utils for traversing control flow graph

---
 .../kotlin/fir/analysis/cfa/CfgUtils.kt       | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)
 create mode 100644 compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt

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
new file mode 100644
index 00000000000..14712032e65
--- /dev/null
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt
@@ -0,0 +1,63 @@
+/*
+ * 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.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
+}
+
+@OptIn(ExperimentalStdlibApi::class)
+fun <D> ControlFlowGraph.traverse(
+    direction: TraverseDirection,
+    visitor: ControlFlowGraphVisitor<*, D>,
+    data: D
+) {
+    val visitedNodes = mutableSetOf<CFGNode<*>>()
+    // used to prevent infinite cycle
+    val delayedNodes = mutableSetOf<CFGNode<*>>()
+    val stack = ArrayDeque<CFGNode<*>>()
+    val initialNode = when (direction) {
+        TraverseDirection.Forward -> enterNode
+        TraverseDirection.Backward -> exitNode
+    }
+    stack.addFirst(initialNode)
+    while (stack.isNotEmpty()) {
+        val node = stack.removeFirst()
+        val previousNodes = when (direction) {
+            TraverseDirection.Forward -> node.previousNodes
+            TraverseDirection.Backward -> node.followingNodes
+        }
+        if (!previousNodes.all { it in visitedNodes }) {
+            if (!delayedNodes.add(node)) {
+                throw IllegalArgumentException("Infinite loop")
+            }
+            stack.addLast(node)
+        }
+
+        node.accept(visitor, data)
+
+        val followingNodes = when (direction) {
+            TraverseDirection.Forward -> node.followingNodes
+            TraverseDirection.Backward -> node.previousNodes
+        }
+
+        followingNodes.forEach {
+            stack.addFirst(it)
+        }
+    }
+}
+
+fun ControlFlowGraph.traverse(
+    direction: TraverseDirection,
+    visitor: ControlFlowGraphVisitorVoid
+) {
+    traverse(direction, visitor, null)
+}
\ No newline at end of file