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 ControlFlowGraph.traverse( + direction: TraverseDirection, + visitor: ControlFlowGraphVisitor<*, D>, + data: D +) { + val visitedNodes = mutableSetOf>() + // used to prevent infinite cycle + val delayedNodes = mutableSetOf>() + val stack = ArrayDeque>() + 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