FIR IDE: support reporting of FirLightSourceElement based diagnostics

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt

@@ -200,20 +200,20 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
 
     override val treeStructure: FlyweightCapableTreeStructure<LighterASTNode> by lazy { WrappedTreeStructure(psi.containingFile) }
 
-    private class WrappedTreeStructure(file: PsiFile) : FlyweightCapableTreeStructure<LighterASTNode> {
+    internal class WrappedTreeStructure(file: PsiFile) : FlyweightCapableTreeStructure<LighterASTNode> {
         private val lighterAST = TreeBackedLighterAST(file.node)
 
-        private fun LighterASTNode.unwrap() = lighterAST.unwrap(this)
+        fun unwrap(node: LighterASTNode) = lighterAST.unwrap(node)
 
-        override fun toString(node: LighterASTNode): CharSequence = node.unwrap().text
+        override fun toString(node: LighterASTNode): CharSequence = unwrap(node).text
 
         override fun getRoot(): LighterASTNode = lighterAST.root
 
         override fun getParent(node: LighterASTNode): LighterASTNode? =
-            node.unwrap().psi.parent?.node?.let { TreeBackedLighterAST.wrap(it) }
+            unwrap(node).psi.parent?.node?.let { TreeBackedLighterAST.wrap(it) }
 
         override fun getChildren(node: LighterASTNode, nodesRef: Ref<Array<LighterASTNode>>): Int {
-            val psi = node.unwrap().psi
+            val psi = unwrap(node).psi
             val children = mutableListOf<PsiElement>()
             var child = psi.firstChild
             while (child != null) {
@@ -232,7 +232,7 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
         }
 
         override fun getStartOffset(node: LighterASTNode): Int {
-            return getStartOffset(node.unwrap().psi)
+            return getStartOffset(unwrap(node).psi)
         }
 
         private fun getStartOffset(element: PsiElement): Int {
@@ -249,7 +249,7 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
         }
 
         override fun getEndOffset(node: LighterASTNode): Int {
-            return getEndOffset(node.unwrap().psi)
+            return getEndOffset(unwrap(node).psi)
         }
 
         private fun getEndOffset(element: PsiElement): Int {
@@ -296,6 +296,19 @@ class FirLightSourceElement(
 ) : FirSourceElement() {
     override val elementType: IElementType
         get() = lighterASTNode.tokenType
+
+    /**
+     * We can create a [FirLightSourceElement] from a [FirPsiSourceElement] by using [FirPsiSourceElement.lighterASTNode];
+     * [unwrapToFirPsiSourceElement] allows to get original [FirPsiSourceElement] in such case.
+     *
+     * If it is `pure` [FirLightSourceElement], i.e, compiler created it in light tree mode, then return [unwrapToFirPsiSourceElement] `null`.
+     * Otherwise, return some not-null result.
+     */
+    fun unwrapToFirPsiSourceElement(): FirPsiSourceElement<*>? {
+        if (treeStructure !is FirPsiSourceElement.WrappedTreeStructure) return null
+        val node = treeStructure.unwrap(lighterASTNode)
+        return node.psi?.toFirPsiSourceElement(kind)
+    }
 }
 
 val FirSourceElement?.psi: PsiElement? get() = (this as? FirPsiSourceElement<*>)?.psi

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt

@@ -5,29 +5,18 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
-import org.jetbrains.kotlin.fir.FirElement
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirPsiSourceElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.CheckersComponent
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
-import org.jetbrains.kotlin.fir.analysis.collectors.CheckerRunningDiagnosticCollectorVisitor
 import org.jetbrains.kotlin.fir.analysis.collectors.components.*
-import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.fir.analysis.diagnostics.*
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
+import org.jetbrains.kotlin.fir.checkers.*
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
-import org.jetbrains.kotlin.fir.checkers.CommonDeclarationCheckers
-import org.jetbrains.kotlin.fir.checkers.CommonExpressionCheckers
-import org.jetbrains.kotlin.fir.checkers.ExtendedDeclarationCheckers
-import org.jetbrains.kotlin.fir.checkers.ExtendedExpressionCheckers
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 
 internal abstract class AbstractFirIdeDiagnosticsCollector(
     session: FirSession,
@@ -53,9 +42,16 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
 
     private inner class Reporter : DiagnosticReporter() {
         override fun report(diagnostic: FirDiagnostic<*>?, context: CheckerContext) {
-            if (diagnostic !is FirPsiDiagnostic<*>) return
+            if (diagnostic == null) return
             if (context.isDiagnosticSuppressed(diagnostic)) return
-            onDiagnostic(diagnostic)
+
+            val psiDiagnostic = when (diagnostic) {
+                is FirPsiDiagnostic<*> -> diagnostic
+                is FirLightDiagnostic -> diagnostic.toPsiDiagnostic()
+                else -> error("Unknown diagnostic type ${diagnostic::class.simpleName}")
+            }
+
+            onDiagnostic(psiDiagnostic)
         }
     }
 
@@ -72,6 +68,40 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
     }
 }
 
+private fun FirLightDiagnostic.toPsiDiagnostic(): FirPsiDiagnostic<*> {
+    val psiSourceElement = element.unwrapToFirPsiSourceElement()
+        ?: error("Diagnostic should be created from PSI in IDE")
+    @Suppress("UNCHECKED_CAST")
+    return when (this) {
+        is FirLightSimpleDiagnostic -> FirPsiSimpleDiagnostic(
+            psiSourceElement,
+            severity,
+            factory as FirDiagnosticFactory0<PsiElement>
+        )
+
+        is FirLightDiagnosticWithParameters1<*> -> FirPsiDiagnosticWithParameters1(
+            psiSourceElement,
+            a,
+            severity,
+            factory as FirDiagnosticFactory1<PsiElement, Any>
+        )
+
+        is FirLightDiagnosticWithParameters2<*, *> -> FirPsiDiagnosticWithParameters2(
+            psiSourceElement,
+            a, b,
+            severity,
+            factory as FirDiagnosticFactory2<PsiElement, Any, Any>
+        )
+
+        is FirLightDiagnosticWithParameters3<*, *, *> -> FirPsiDiagnosticWithParameters3(
+            psiSourceElement,
+            a, b, c,
+            severity,
+            factory as FirDiagnosticFactory3<PsiElement, Any, Any, Any>
+        )
+        else -> error("Unknown diagnostic type ${this::class.simpleName}")
+    }
+}
 
 private object CheckersFactory {
     private val extendedDeclarationCheckers = createDeclarationCheckers(ExtendedDeclarationCheckers)