From 771823c9e76a521ff99645a3b4268e5078653ec2 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Mon, 21 Nov 2022 14:14:38 +0100
Subject: [PATCH] [AA] AbstractSymbolTest: improve implicit checks

^KT-54826
---
 .../test/cases/symbols/AbstractSymbolTest.kt  | 48 ++++++++++++-------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/symbols/AbstractSymbolTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/symbols/AbstractSymbolTest.kt
index 5acf3db23d8..2a617ff1e89 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/symbols/AbstractSymbolTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/symbols/AbstractSymbolTest.kt
@@ -85,20 +85,22 @@ abstract class AbstractSymbolTest : AbstractAnalysisApiSingleFileTest() {
             analyseForTest(ktFile) {
                 val (symbols, symbolForPrettyRendering) = collectSymbols(ktFile, testServices)
 
-                val pointerWithRenderedSymbol = symbols.flatMap { symbol ->
-                    listOf(
+                val pointerWithRenderedSymbol = symbols
+                    .asSequence()
+                    .flatMap { symbol ->
+                        sequenceOf(symbol to true) + symbol.withImplicitSymbols().map { implicitSymbol ->
+                            implicitSymbol to false
+                        }
+                    }
+                    .distinctBy { it.first }
+                    .map { (symbol, shouldBeRendered) ->
                         PointerWithRenderedSymbol(
-                            symbol.safePointer(),
-                            renderSymbolForComparison(symbol),
-                        )
-                    ) + symbol.withImplicitSymbols().map { implicitSymbol ->
-                        PointerWithRenderedSymbol(
-                            implicitSymbol.safePointer(),
-                            renderSymbolForComparison(implicitSymbol),
-                            shouldBeRendered = false,
+                            pointer = symbol.safePointer(),
+                            rendered = renderSymbolForComparison(symbol),
+                            shouldBeRendered = shouldBeRendered,
                         )
                     }
-                }
+                    .toList()
 
                 val pointerWithPrettyRenderedSymbol = symbolForPrettyRendering.map { symbol ->
                     PointerWithRenderedSymbol(
@@ -281,10 +283,22 @@ private data class PointerWrapper(
     val pointerWithoutPsiAnchor: KtSymbolPointer<*>?,
 )
 
-private fun KtSymbol?.withImplicitSymbols(): List<KtSymbol> = when (this) {
-    null -> emptyList()
-    is KtPropertySymbol -> listOf(this) + setter.withImplicitSymbols() + getter.withImplicitSymbols()
-    is KtPropertySetterSymbol -> listOf(this) + parameter
-    is KtValueParameterSymbol -> listOf(this) + generatedPrimaryConstructorProperty.withImplicitSymbols()
-    else -> listOf(this)
+private fun KtSymbol?.withImplicitSymbols(): Sequence<KtSymbol> {
+    val ktSymbol = this ?: return emptySequence()
+    return sequence {
+        yield(ktSymbol)
+
+        if (ktSymbol is KtPropertySymbol) {
+            yieldAll(ktSymbol.getter.withImplicitSymbols())
+            yieldAll(ktSymbol.setter.withImplicitSymbols())
+        }
+
+        if (ktSymbol is KtPropertySetterSymbol) {
+            yieldAll(ktSymbol.parameter.withImplicitSymbols())
+        }
+
+        if (ktSymbol is KtValueParameterSymbol) {
+            yieldAll(ktSymbol.generatedPrimaryConstructorProperty.withImplicitSymbols())
+        }
+    }
 }