From acbb67f851c4bb664bbe4f76d69ea8ac9506c285 Mon Sep 17 00:00:00 2001
From: Oleg Ivanov <oleg.ivanov@jetbrains.com>
Date: Fri, 14 Aug 2020 17:26:39 +0300
Subject: [PATCH] [FIR] Fix collectDataForNode backward traverse

---
 .../kotlin/fir/analysis/cfa/CfgTraverser.kt        |  5 ++++-
 .../jetbrains/kotlin/fir/analysis/cfa/CfgUtils.kt  | 14 +++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

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
index 416ed8b20be..a9c97806dc2 100644
--- 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
@@ -55,6 +55,9 @@ private fun <I : ControlFlowInfo<I, K, V>, K : Any, V : Any> ControlFlowGraph.co
 ) {
     val nodes = getNodesInOrder(direction)
     for (node in nodes) {
+        if (direction == TraverseDirection.Backward && node is CFGNodeWithCfgOwner<*>) {
+            node.subGraphs.forEach { it.collectDataForNodeInternal(direction, initialInfo, visitor, nodeMap, changed) }
+        }
         if (!(node.isEnterNode(direction) && node.owner.owner == null)) {
             val previousNodes = when (direction) {
                 TraverseDirection.Forward -> node.previousCfgNodes
@@ -69,7 +72,7 @@ private fun <I : ControlFlowInfo<I, K, V>, K : Any, V : Any> ControlFlowGraph.co
                 nodeMap[node] = newData
             }
         }
-        if (node is CFGNodeWithCfgOwner<*>) {
+        if (direction == TraverseDirection.Forward && node is CFGNodeWithCfgOwner<*>) {
             node.subGraphs.forEach { it.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 86bc6cf65cf..1c5cfb3f0a4 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
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.analysis.cfa
 
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNodeWithCfgOwner
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
 
 fun ControlFlowGraph.getEnterNode(direction: TraverseDirection): CFGNode<*> = when (direction) {
@@ -34,7 +35,14 @@ val CFGNode<*>.previousCfgNodes: List<CFGNode<*>>
     }
 
 val CFGNode<*>.followingCfgNodes: List<CFGNode<*>>
-    get() = followingNodes.filter {
-        val kind = outgoingEdges.getValue(it)
-        kind.usedInCfa && !kind.isDead
+    get() {
+        val nodes = mutableListOf<CFGNode<*>>()
+
+        followingNodes.filterTo(nodes) {
+            val kind = outgoingEdges.getValue(it)
+            kind.usedInCfa && !kind.isDead
+        }
+        (this as? CFGNodeWithCfgOwner<*>)?.subGraphs?.mapTo(nodes) { it.enterNode }
+
+        return nodes
     }
\ No newline at end of file