From f5ec37db95ca56ede1b8c34907a085f7192a60a5 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Fri, 29 Jan 2021 19:50:40 +0100
Subject: [PATCH] FIR IDE: report diagnostics for non-kt elements

---
 .../api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt  | 1 -
 .../FirIdeStructureElementDiagnosticsCollector.kt          | 7 +++----
 .../low/level/api/file/structure/FileStructureElement.kt   | 5 +++--
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt
index 0bb4e9d8f8a..e8015c33fe8 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt
@@ -42,7 +42,6 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
     private inner class Reporter : DiagnosticReporter() {
         override fun report(diagnostic: FirDiagnostic<*>?, context: CheckerContext) {
             if (diagnostic !is FirPsiDiagnostic<*>) return
-            if (diagnostic.element.psi !is KtElement) return
             onDiagnostic(diagnostic)
         }
     }
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeStructureElementDiagnosticsCollector.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeStructureElementDiagnosticsCollector.kt
index cac6b26af64..4c476d7e156 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeStructureElementDiagnosticsCollector.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeStructureElementDiagnosticsCollector.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
@@ -24,12 +25,10 @@ internal class FirIdeStructureElementDiagnosticsCollector private constructor(
 ) : AbstractFirIdeDiagnosticsCollector(
     session,
 ) {
-    private val result = mutableMapOf<KtElement, MutableList<FirPsiDiagnostic<*>>>()
+    private val result = mutableMapOf<PsiElement, MutableList<FirPsiDiagnostic<*>>>()
 
     override fun onDiagnostic(diagnostic: FirPsiDiagnostic<*>) {
-        (diagnostic.psiElement as? KtElement)?.let { ktElement ->
-            result.addValueFor(ktElement, diagnostic)
-        }
+        result.addValueFor(diagnostic.psiElement, diagnostic)
     }
 
     override fun onDeclarationEnter(
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt
index c1c28f4a054..6d4722b34b4 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
@@ -25,9 +26,9 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.ktDeclaration
 import org.jetbrains.kotlin.psi.*
 
 internal class FileStructureElementDiagnostics(
-    private val map: Map<KtElement, List<FirPsiDiagnostic<*>>>
+    private val map: Map<PsiElement, List<FirPsiDiagnostic<*>>>
 ) {
-    fun diagnosticsFor(element: KtElement): List<FirPsiDiagnostic<*>> = map[element] ?: emptyList()
+    fun diagnosticsFor(element: PsiElement): List<FirPsiDiagnostic<*>> = map[element] ?: emptyList()
 
     inline fun forEach(action: (List<FirPsiDiagnostic<*>>) -> Unit) = map.values.forEach(action)
 }