From dbb0496d239824ef71995691f2cc6a589fed20fa Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Mon, 25 Oct 2021 11:48:01 -0700
Subject: [PATCH] FIR: add bfs helper and use it

It's also used in Kotlin quickfixes in intellij repo
---
 .../api/fir/components/KtFirTypeProvider.kt   | 10 ++------
 .../jetbrains/kotlin/util/containerUtils.kt   | 23 +++++++++++++++++++
 2 files changed, 25 insertions(+), 8 deletions(-)
 create mode 100644 core/compiler.common/src/org/jetbrains/kotlin/util/containerUtils.kt

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeProvider.kt
index 2f8adc682d6..77710d82a08 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeProvider.kt
@@ -46,6 +46,7 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtTypeReference
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.model.CaptureStatus
+import org.jetbrains.kotlin.util.bfs
 
 internal class KtFirTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -176,14 +177,7 @@ internal class KtFirTypeProvider(
 
     override fun getAllSuperTypes(type: KtType, shouldApproximate: Boolean): List<KtType> {
         require(type is KtFirType)
-        val queue = type.coneType.getDirectSuperTypes(shouldApproximate).toCollection(ArrayDeque())
-        return buildSet {
-            while (queue.isNotEmpty()) {
-                val current = queue.removeFirst()
-                if (!add(current)) continue
-                queue.addAll(current.getDirectSuperTypes(shouldApproximate))
-            }
-        }.map { it.asKtType() }
+        return listOf(type.coneType).bfs { it.getDirectSuperTypes(shouldApproximate).iterator() }.mapTo(mutableListOf()) { it.asKtType() }
     }
 }
 
diff --git a/core/compiler.common/src/org/jetbrains/kotlin/util/containerUtils.kt b/core/compiler.common/src/org/jetbrains/kotlin/util/containerUtils.kt
new file mode 100644
index 00000000000..fc48124ed12
--- /dev/null
+++ b/core/compiler.common/src/org/jetbrains/kotlin/util/containerUtils.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2021 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.util
+
+/**
+ * Perform BFS on the given collection with neighbors created by the given function.
+ */
+fun <T> Collection<T>.bfs(getNeighbors: (T) -> Iterator<T>): Sequence<T> {
+    val queue = ArrayDeque(this)
+    val visited = mutableSetOf<T>()
+    return sequence {
+        while (queue.isNotEmpty()) {
+            val current = queue.removeFirst()
+            if (current in visited) continue
+            visited.add(current)
+            yield(current)
+            getNeighbors(current).forEach(queue::add)
+        }
+    }
+}