FIR IDE: report diagnostics for non-kt elements

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)
         }
     }

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(

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)
 }