From 238488d4fb571f1b1c6292edbf46a83f4528ee61 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 3 Apr 2020 15:19:08 +0300
Subject: [PATCH] [FIR] Fix collection subgraphs to drop from delegate calls

---
 .../resolve/dfa/cfg/ControlFlowGraphBuilder.kt   | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
index 3862f7eecb6..e510e0d8050 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSymbolOwner
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.dfa.*
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
@@ -927,15 +928,24 @@ class ControlFlowGraphBuilder {
 
     fun dropSubgraphFromCall(call: FirFunctionCall) {
         val graphs = mutableListOf<ControlFlowGraph>()
+
         call.acceptChildren(object : FirDefaultVisitorVoid() {
             override fun visitElement(element: FirElement) {
                 element.acceptChildren(this)
             }
 
             override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-                anonymousFunction.controlFlowGraphReference.controlFlowGraph?.let {
-                    graphs += it
-                }
+                anonymousFunction.controlFlowGraphReference.accept(this)
+            }
+
+            override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) {
+                anonymousObject.controlFlowGraphReference.accept(this)
+            }
+
+            override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference) {
+                val graph = controlFlowGraphReference.controlFlowGraph ?: return
+                if (graph.owner == null) return
+                graphs += graph
             }
         }, null)