From e78e8f8c79472e84a672c9765089e1bb6d9dd9ea Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Fri, 9 Apr 2021 11:40:00 +0200
Subject: [PATCH] FIR IDE: fix behaviour of DiagnosticTraversalCounterTest

---
 .../AbstractDiagnosticTraversalCounterTest.kt   | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/AbstractDiagnosticTraversalCounterTest.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/AbstractDiagnosticTraversalCounterTest.kt
index 6588352b2be..382f525c936 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/AbstractDiagnosticTraversalCounterTest.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/AbstractDiagnosticTraversalCounterTest.kt
@@ -7,7 +7,9 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostic
 
 import com.intellij.openapi.util.io.FileUtil
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirRealSourceElementKind
 import org.jetbrains.kotlin.fir.SessionConfiguration
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
@@ -15,6 +17,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDiagnosticsForFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.createResolveStateForNoCaching
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.BeforeElementDiagnosticCollectionHandler
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.renderWithClassName
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.psi.KtFile
@@ -72,7 +75,17 @@ abstract class AbstractDiagnosticTraversalCounterTest : KotlinLightCodeInsightFi
         handler: BeforeElementTestDiagnosticCollectionHandler
     ): List<Pair<FirElement, Int>> {
         val errorElements = mutableListOf<Pair<FirElement, Int>>()
-        val nonDuplicatingElements = findNonDuplicatingFirElements(firFile)
+        val nonDuplicatingElements = findNonDuplicatingFirElements(firFile).filter { element ->
+            when {
+                element is FirTypeRef && element.source?.kind != FirRealSourceElementKind -> {
+                    // AbstractDiagnosticCollectorVisitor do not visit such elements
+                    false
+                }
+                element.source?.kind == FirRealSourceElementKind -> true
+                SingleNonLocalDeclarationDiagnosticRetriever.shouldDiagnosticsAlwaysBeCheckedOn(element) -> true
+                else -> false
+            }
+        }
         firFile.accept(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {
                 if (element !in nonDuplicatingElements) return
@@ -86,6 +99,7 @@ abstract class AbstractDiagnosticTraversalCounterTest : KotlinLightCodeInsightFi
         return errorElements
     }
 
+
     private fun findNonDuplicatingFirElements(
         firFile: FirElement,
     ): Set<FirElement> {
@@ -93,6 +107,7 @@ abstract class AbstractDiagnosticTraversalCounterTest : KotlinLightCodeInsightFi
         firFile.accept(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {
                 elementUsageCount.compute(element) { _, count -> (count ?: 0) + 1 }
+                element.acceptChildren(this)
             }
         })
         return elementUsageCount.filterValues { it == 1 }.keys