FIR IDE: move diagnostic highlighting to separate

pass

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt

@@ -52,6 +52,7 @@ abstract class AbstractDiagnosticCollector(
 
     @Suppress("LeakingThis")
     private var context = PersistentCheckerContext(this, returnTypeCalculator)
+    private var currentAction = DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
 
     fun initializeComponents(vararg components: AbstractDiagnosticCollectorComponent) {
         if (componentsInitialized) {
@@ -75,7 +76,6 @@ abstract class AbstractDiagnosticCollector(
 
         private fun visitJump(loopJump: FirLoopJump) {
             loopJump.runComponents()
-            loopJump.acceptChildren(this, null)
             loopJump.target.labeledElement.takeIf { it is FirErrorLoop }?.accept(this, null)
         }
 
@@ -157,16 +157,26 @@ abstract class AbstractDiagnosticCollector(
             }
         }
 
-        private fun visitWithDeclaration(declaration: FirDeclaration) {
-            declaration.runComponents()
-            withDeclaration(declaration) {
-                declaration.acceptChildren(this, null)
+        private inline fun visitWithDeclaration(
+            declaration: FirDeclaration,
+            block: () -> Unit = { declaration.acceptChildren(this, null) }
+        ) {
+            if (!currentAction.checkNested) return
+
+            val action = onDeclarationEnter(declaration)
+            if (action.checkCurrentDeclaration) {
+                declaration.runComponents()
             }
+            withDiagnosticsAction(action) {
+                withDeclaration(declaration) {
+                    block()
+                }
+            }
+            onDeclarationExit(declaration)
         }
 
         private fun visitWithDeclarationAndReceiver(declaration: FirDeclaration, labelName: Name?, receiverTypeRef: FirTypeRef?) {
-            declaration.runComponents()
-            withDeclaration(declaration) {
+            visitWithDeclaration(declaration) {
                 withLabelAndReceiverType(
                     labelName,
                     declaration,
@@ -178,6 +188,11 @@ abstract class AbstractDiagnosticCollector(
         }
     }
 
+    protected open fun onDeclarationEnter(declaration: FirDeclaration): DiagnosticCollectorDeclarationAction =
+        DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+
+    protected open fun onDeclarationExit(declaration: FirDeclaration) {}
+
     private inline fun <R> withDeclaration(declaration: FirDeclaration, block: () -> R): R {
         val existingContext = context
         context = context.addDeclaration(declaration)
@@ -208,6 +223,23 @@ abstract class AbstractDiagnosticCollector(
             context = existingContext
         }
     }
+
+    private inline fun <R> withDiagnosticsAction(action: DiagnosticCollectorDeclarationAction, block: () -> R): R {
+        val oldAction = currentAction
+        currentAction = action
+        return try {
+            block()
+        } finally {
+            currentAction = oldAction
+        }
+    }
+}
+
+enum class DiagnosticCollectorDeclarationAction(val checkCurrentDeclaration: Boolean, val checkNested: Boolean) {
+    CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED(checkCurrentDeclaration = true, checkNested = true),
+    CHECK_CURRENT_DECLARATION_AND_SKIP_NESTED(checkCurrentDeclaration = true, checkNested = false),
+    SKIP_CURRENT_DECLARATION_AND_CHECK_NESTED(checkCurrentDeclaration = false, checkNested = true),
+    SKIP(checkCurrentDeclaration = false, checkNested = false),
 }
 
 fun AbstractDiagnosticCollector.registerAllComponents() {

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/KotlinFirPsiChecker.kt

@@ -30,52 +30,9 @@ class KotlinFirPsiChecker : AbstractKotlinPsiChecker() {
             throw ProcessCanceledException()
         }
         analyze(element) {
-            this.highlightDiagnostics(element, holder)
-
             FirAfterResolveHighlightingVisitor
                 .createListOfVisitors(this, holder)
                 .forEach(element::accept)
         }
     }
-
-    private fun KtAnalysisSession.highlightDiagnostics(element: KtElement, holder: AnnotationHolder) {
-        val diagnostics = element.getDiagnostics()
-        if (diagnostics.isEmpty()) return
-
-        if (diagnostics.none(Diagnostic::isValid)) return
-
-        if (shouldHighlightErrors(element)) {
-            highlightDiagnostics(diagnostics, holder)
-        }
-    }
-
-    private fun highlightDiagnostics(diagnostics: Collection<Diagnostic>, holder: AnnotationHolder) {
-        diagnostics.groupBy { it.factory }.forEach { group -> registerDiagnosticAnnotations(group.value, holder) }
-    }
-
-    private fun registerDiagnosticAnnotations(diagnostics: List<Diagnostic>, holder: AnnotationHolder) {
-        assert(diagnostics.isNotEmpty())
-        val diagnostic = diagnostics.first()
-
-        val ranges = diagnostic.textRanges
-
-        ranges.forEach { range ->
-            diagnostics.forEach { diagnostic ->
-                Diagnostic2Annotation.createAnnotation(
-                    diagnostic,
-                    range,
-                    holder,
-                    nonDefaultMessage = null,
-                    textAttributes = null,
-                    highlightType = null,
-                    renderMessage = IdeErrorMessages::render
-                )
-            }
-        }
-
-    }
-
-    private fun shouldHighlightErrors(element: KtElement): Boolean {
-        return true // todo
-    }
 }

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/KotlinHighLevelDiagnosticHighlightingPass.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.highlighter
+
+import com.intellij.codeHighlighting.TextEditorHighlightingPass
+import com.intellij.codeHighlighting.TextEditorHighlightingPassFactory
+import com.intellij.codeHighlighting.TextEditorHighlightingPassFactoryRegistrar
+import com.intellij.codeHighlighting.TextEditorHighlightingPassRegistrar
+import com.intellij.codeInsight.daemon.impl.HighlightInfo
+import com.intellij.codeInsight.daemon.impl.HighlightInfoType
+import com.intellij.codeInsight.daemon.impl.UpdateHighlightersUtil
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.progress.ProgressIndicator
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiElementVisitor
+import com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.highlighter.Diagnostic2Annotation
+import org.jetbrains.kotlin.idea.highlighter.IdeErrorMessages
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtVisitorVoid
+
+class KotlinHighLevelDiagnosticHighlightingPass(
+    private val ktFile: KtFile,
+    document: Document,
+) : TextEditorHighlightingPass(ktFile.project, document) {
+    private val diagnosticInfos = mutableListOf<HighlightInfo>()
+
+    override fun doCollectInformation(progress: ProgressIndicator) = analyze(ktFile) {
+        ktFile.collectDiagnosticsForFile().forEach { diagnostic ->
+            if (!diagnostic.isValid) return@forEach
+            diagnostic.textRanges.forEach { range ->
+                HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR/*TODO*/)
+                    .descriptionAndTooltip(Diagnostic2Annotation.getMessage(diagnostic, IdeErrorMessages::render))
+                    .range(range)
+                    .create()
+                    ?.let(diagnosticInfos::add)
+            }
+        }
+    }
+
+
+    override fun doApplyInformationToEditor() {
+        UpdateHighlightersUtil.setHighlightersToEditor(
+            myProject, myDocument!!, /*startOffset=*/0, ktFile.textLength, diagnosticInfos, colorsScheme, id
+        )
+    }
+}
+
+internal class KotlinDiagnosticHighlightingPassFactory : TextEditorHighlightingPassFactory, TextEditorHighlightingPassFactoryRegistrar {
+    override fun registerHighlightingPassFactory(registrar: TextEditorHighlightingPassRegistrar, project: Project) {
+        registrar.registerTextEditorHighlightingPass(this, null, null, false, -1)
+    }
+
+    override fun createHighlightingPass(file: PsiFile, editor: Editor): TextEditorHighlightingPass? {
+        if (file !is KtFile) return null
+        if (file.isCompiled) return null
+        return KotlinHighLevelDiagnosticHighlightingPass(file, editor.document)
+    }
+}

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -60,6 +60,8 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
 
     fun KtElement.getDiagnostics(): Collection<Diagnostic> = diagnosticProvider.getDiagnosticsForElement(this)
 
+    fun KtFile.collectDiagnosticsForFile(): Collection<Diagnostic> = diagnosticProvider.collectDiagnosticsForFile(this)
+
     fun KtSymbolWithKind.getContainingSymbol(): KtSymbolWithKind? = containingDeclarationProvider.getContainingDeclaration(this)
 
     fun KtClassOrObjectSymbol.getMemberScope(): KtMemberScope = scopeProvider.getMemberScope(this)

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtDiagnosticProvider.kt

@@ -7,7 +7,9 @@ package org.jetbrains.kotlin.idea.frontend.api.components
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 
 abstract class KtDiagnosticProvider : KtAnalysisSessionComponent() {
     abstract fun getDiagnosticsForElement(element: KtElement): Collection<Diagnostic>
-}
+    abstract fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic>
+}

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

@@ -39,6 +39,8 @@ abstract class FirModuleResolveState {
 
     abstract fun getDiagnostics(element: KtElement): List<Diagnostic>
 
+    abstract fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic>
+
     // todo temporary, used only in completion
     abstract fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile)
 
@@ -80,6 +82,9 @@ internal class FirModuleResolveStateImpl(
     override fun getDiagnostics(element: KtElement): List<Diagnostic> =
         diagnosticsCollector.getDiagnosticsFor(element)
 
+    override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic> =
+        diagnosticsCollector.collectDiagnosticsForFile(ktFile)
+
     override fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
         error("Should be called only from FirModuleResolveStateForCompletion")
     }

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

@@ -64,4 +64,8 @@ internal class FirModuleResolveStateForCompletion(
     override fun getDiagnostics(element: KtElement): List<Diagnostic> {
         error("Diagnostics should not be retrieved in completion")
     }
+
+    override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic> {
+        error("Diagnostics should not be retrieved in completion")
+    }
 }

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

@@ -9,44 +9,39 @@ import com.intellij.openapi.diagnostic.Logger
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
-import org.jetbrains.kotlin.fir.analysis.collectors.registerAllComponents
+import org.jetbrains.kotlin.fir.analysis.collectors.components.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
-import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.psi.KtElement
 
-internal class FirIdeDiagnosticsCollector private constructor(
+internal abstract class AbstractFirIdeDiagnosticsCollector(
     session: FirSession,
 ) : AbstractDiagnosticCollector(
     session,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(session, ScopeSession())
 ) {
-    private val result = mutableMapOf<KtElement, MutableList<Diagnostic>>()
-
     init {
         registerAllComponents()
     }
 
+    protected abstract fun onDiagnostic(diagnostic: Diagnostic)
+
+
     private inner class Reporter : DiagnosticReporter() {
         override fun report(diagnostic: FirDiagnostic<*>?) {
-            checkCanceled()
-            try {
-                if (diagnostic !is FirPsiDiagnostic<*>) return
-                val psi = diagnostic.element.psi as? KtElement ?: return
-                result.addValueFor(psi, diagnostic.asPsiBasedDiagnostic())
-            } catch (e: Throwable) {
-                LOG.error(e)
-            }
+            if (diagnostic !is FirPsiDiagnostic<*>) return
+            if (diagnostic.element.psi !is KtElement) return
+            onDiagnostic(diagnostic.asPsiBasedDiagnostic())
         }
     }
 
     override var reporter: DiagnosticReporter = Reporter()
-        private set
 
     override fun initializeCollector() {
         reporter = Reporter()
@@ -58,16 +53,6 @@ internal class FirIdeDiagnosticsCollector private constructor(
     }
 
     companion object {
-        private val LOG = Logger.getInstance(FirIdeDiagnosticsCollector::class.java)
-
-        /**
-         * Collects diagnostics for given [firFile]
-         * Should be called under [firFile]-based lock
-         */
-        fun collect(firFile: FirFile): Map<KtElement, List<Diagnostic>> =
-            FirIdeDiagnosticsCollector(firFile.session).let { collector ->
-                collector.collectDiagnostics(firFile)
-                collector.result
-            }
+        private val LOG = Logger.getInstance(AbstractFirIdeDiagnosticsCollector::class.java)
     }
 }

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

@@ -6,48 +6,23 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import java.util.concurrent.ConcurrentHashMap
 
 internal class DiagnosticsCollector(
-    private val firFileBuilder: FirFileBuilder,
+    private val fileStructureCache: FileStructureCache,
     private val cache: ModuleFileCache,
 ) {
-    private val diagnosticsForFile = ConcurrentHashMap<KtFile, DiagnosticsForFile>()
+    fun getDiagnosticsFor(element: KtElement): List<Diagnostic> =
+        fileStructureCache
+            .getFileStructure(element.containingKtFile, cache)
+            .getStructureElementFor(element)
+            .diagnostics.diagnosticsFor(element)
 
-    fun getDiagnosticsFor(element: KtElement): List<Diagnostic> {
-        val ktFile = element.containingKtFile
-        val diagnostics = diagnosticsForFile.computeIfAbsent(ktFile) {
-            val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(
-                ktFile,
-                cache,
-                toPhase = FirResolvePhase.BODY_RESOLVE,
-                checkPCE = true
-            )
-            DiagnosticsForFile.collectDiagnosticsForFile(firFile)
-        }
-        return diagnostics.getDiagnosticsFor(element)
-    }
-}
-
-private class DiagnosticsForFile private constructor(private val diagnostics: Map<KtElement, List<Diagnostic>>) {
-    fun getDiagnosticsFor(element: KtElement): List<Diagnostic> = diagnostics[element].orEmpty()
-
-    companion object {
-        /**
-         * Collects diagnostics for given [firFile]
-         * Should be called under [firFile]-based lock
-         */
-        fun collectDiagnosticsForFile(firFile: FirFile): DiagnosticsForFile {
-            require(firFile.resolvePhase >= FirResolvePhase.BODY_RESOLVE) {
-                "To collect diagnostics at least FirResolvePhase.BODY_RESOLVE is needed, but file ${firFile.name} was resolved to ${firFile.resolvePhase}"
-            }
-            return DiagnosticsForFile(FirIdeDiagnosticsCollector.collect(firFile))
-        }
-    }
+    fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic> =
+        fileStructureCache
+            .getFileStructure(ktFile, cache)
+            .getAllDiagnosticsForFile()
 }

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
+
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirFile
+
+internal class FirIdeFileDiagnosticsCollector private constructor(
+    session: FirSession,
+) : AbstractFirIdeDiagnosticsCollector(
+    session,
+) {
+    private val result = mutableListOf<Diagnostic>()
+
+    override fun onDiagnostic(diagnostic: Diagnostic) {
+        result += diagnostic
+    }
+
+    companion object {
+        fun collectForFile(firFile: FirFile): List<Diagnostic> =
+            FirIdeFileDiagnosticsCollector(firFile.session).let { collector ->
+                collector.collectDiagnostics(firFile)
+                collector.result
+            }
+    }
+}

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
+
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureElementDiagnostics
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
+import org.jetbrains.kotlin.psi.KtAnnotated
+import org.jetbrains.kotlin.psi.KtElement
+
+internal class FirIdeStructureElementDiagnosticsCollector private constructor(
+    session: FirSession,
+    private val onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction,
+    private val onDeclarationExit: (FirDeclaration) -> Unit
+) : AbstractFirIdeDiagnosticsCollector(
+    session,
+) {
+    private val result = mutableMapOf<KtElement, MutableList<Diagnostic>>()
+
+    override fun onDiagnostic(diagnostic: Diagnostic) {
+        (diagnostic.psiElement as? KtElement)?.let { ktElement ->
+            result.addValueFor(ktElement, diagnostic)
+        }
+    }
+
+    override fun onDeclarationEnter(
+        declaration: FirDeclaration,
+    ): DiagnosticCollectorDeclarationAction =
+        onDeclarationEnter.invoke(declaration)
+
+    override fun onDeclarationExit(declaration: FirDeclaration) {
+        onDeclarationExit.invoke(declaration)
+    }
+
+
+    companion object {
+        fun collectForStructureElement(
+            firFile: FirFile,
+            onDeclarationExit: (FirDeclaration) -> Unit = {},
+            onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction,
+        ): FileStructureElementDiagnostics =
+            FirIdeStructureElementDiagnosticsCollector(firFile.session, onDeclarationEnter, onDeclarationExit).let { collector ->
+                collector.collectDiagnostics(firFile)
+                FileStructureElementDiagnostics(collector.result)
+            }
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
@@ -43,19 +44,28 @@ internal class WithInBlockModificationFileStructureElement(
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
-
     fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
 
     override val diagnostics: FileStructureElementDiagnostics by lazy {
-        val innerDeclarations = psi
-            .collectDescendantsOfTypeTo<KtDeclaration, HashSet<KtDeclaration>>(
-                hashSetOf(psi),
-                canGoInside = { true }
-            )
+        var inCurrentDeclaration = false
+
         FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(
-            innerDeclarations,
             firFile,
-            substitution = null
+            onDeclarationEnter = { firDeclaration ->
+                when {
+                    firDeclaration == firSymbol.fir -> {
+                        inCurrentDeclaration = true
+                        DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                    }
+                    inCurrentDeclaration -> DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                    else -> DiagnosticCollectorDeclarationAction.SKIP_CURRENT_DECLARATION_AND_CHECK_NESTED
+                }
+            },
+            onDeclarationExit = { declaration ->
+                if (declaration == firSymbol.fir) {
+                    inCurrentDeclaration = false
+                }
+            }
         )
     }
 
@@ -73,14 +83,30 @@ internal class NonLocalDeclarationFileStructureElement(
         FirElementsRecorder.recordElementsFrom(fir, recorder)
 
     override val diagnostics: FileStructureElementDiagnostics by lazy {
-        @Suppress("RemoveExplicitTypeArguments")
-        val innerDeclarations = psi
-            .collectDescendantsOfTypeTo<KtDeclaration, HashSet<KtDeclaration>>(
-                hashSetOf(psi),
-                canGoInside = { it.safeAs<KtNamedFunction>()?.hasExplicitTypeOrUnit != true }
-            )
-
-        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(innerDeclarations, firFile)
+        var inCurrentDeclaration = false
+        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(
+            firFile,
+            onDeclarationEnter = { firDeclaration ->
+                when {
+                    firDeclaration == fir -> {
+                        inCurrentDeclaration = true
+                        DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                    }
+                    (firDeclaration.psi as? KtNamedFunction)?.hasExplicitTypeOrUnit == true -> {
+                        DiagnosticCollectorDeclarationAction.SKIP
+                    }
+                    inCurrentDeclaration -> {
+                        DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                    }
+                    else -> DiagnosticCollectorDeclarationAction.SKIP_CURRENT_DECLARATION_AND_CHECK_NESTED
+                }
+            },
+            onDeclarationExit = { firDeclaration ->
+                if (firDeclaration == fir) {
+                    inCurrentDeclaration = false
+                }
+            },
+        )
     }
 
     companion object {
@@ -104,7 +130,10 @@ internal data class FileWithoutDeclarationsFileStructureElement(
         FirElementsRecorder.recordElementsFrom(firFile, recorder)
 
     override val diagnostics: FileStructureElementDiagnostics by lazy {
-        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(hashSetOf(psi), firFile)
+        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(firFile) { firDeclaration ->
+            if (firDeclaration is FirFile) DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_SKIP_NESTED
+            else DiagnosticCollectorDeclarationAction.SKIP
+        }
     }
 
     companion object {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirDiagnosticProvider.kt

@@ -6,20 +6,13 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
-import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartCast
-import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartcastKind
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.components.KtDiagnosticProvider
-import org.jetbrains.kotlin.idea.frontend.api.components.KtSmartCastProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
-import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFile
 
 internal class KtFirDiagnosticProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -28,4 +21,7 @@ internal class KtFirDiagnosticProvider(
     override fun getDiagnosticsForElement(element: KtElement): Collection<Diagnostic> = withValidityAssertion {
         LowLevelFirApiFacade.getDiagnosticsFor(element, firResolveState)
     }
+
+    override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic> =
+        firResolveState.collectDiagnosticsForFile(ktFile)
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/Diagnostic2Annotation.kt

@@ -33,6 +33,14 @@ object Diagnostic2Annotation {
         return message
     }
 
+    fun getMessage(diagnostic: Diagnostic, renderMessage: (Diagnostic) -> String): String {
+        val message = renderMessage(diagnostic)
+        if (ApplicationManager.getApplication().isInternal || ApplicationManager.getApplication().isUnitTestMode) {
+            return "[${diagnostic.factory.name}] $message"
+        }
+        return message
+    }
+
     fun getDefaultMessage(diagnostic: Diagnostic): String {
         val message = DefaultErrorMessages.render(diagnostic)
         if (ApplicationManager.getApplication().isInternal || ApplicationManager.getApplication().isUnitTestMode) {
@@ -41,6 +49,8 @@ object Diagnostic2Annotation {
         return message
     }
 
+
+
     fun createAnnotation(
         diagnostic: Diagnostic,
         range: TextRange,

idea/resources-fir/META-INF/plugin.xml

@@ -106,6 +106,8 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
                             implementationClass="org.jetbrains.kotlin.idea.completion.KotlinFirCompletionContributor"/>
     <codeInsight.typeInfo language="kotlin"
                           implementationClass="org.jetbrains.kotlin.idea.codeInsight.KotlinHighLevelExpressionTypeProvider"/>
+    <highlightingPassFactory implementation="org.jetbrains.kotlin.idea.fir.highlighter.KotlinDiagnosticHighlightingPassFactory"/>
+
 
   </extensions>