FIR IDE: allow creating inspections by extended checkers

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/AbstractHLDiagnosticBasedInspection.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2021 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.api
+
+import org.jetbrains.kotlin.idea.fir.api.applicator.HLApplicatorInput
+import org.jetbrains.kotlin.idea.fir.api.applicator.HLApplicatorInputProvider
+import org.jetbrains.kotlin.idea.fir.api.applicator.inputProvider
+import org.jetbrains.kotlin.idea.frontend.api.components.KtDiagnosticCheckerFilter
+import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
+import org.jetbrains.kotlin.psi.KtElement
+import kotlin.reflect.KClass
+
+abstract class AbstractHLDiagnosticBasedInspection<PSI : KtElement, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>, INPUT : HLApplicatorInput>(
+    elementType: KClass<PSI>,
+    private val diagnosticType: KClass<DIAGNOSTIC>,
+) : AbstractHLInspection<PSI, INPUT>(elementType) {
+    abstract val inputByDiagnosticProvider: HLInputByDiagnosticProvider<PSI, DIAGNOSTIC, INPUT>
+
+    final override val inputProvider: HLApplicatorInputProvider<PSI, INPUT> = inputProvider { psi ->
+        val diagnostics = psi.getDiagnostics(KtDiagnosticCheckerFilter.ONLY_EXTENDED_CHECKERS)
+        val suitableDiagnostics = diagnostics.filterIsInstance(diagnosticType.java)
+        val diagnostic = suitableDiagnostics.firstOrNull() ?: return@inputProvider null
+        // TODO handle case with multiple diagnostics on single element
+        with(inputByDiagnosticProvider) { createInfo(diagnostic) }
+    }
+}
+
+

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/HLInputByDiagnosticProvider.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2021 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.api
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.fir.api.applicator.HLApplicatorInput
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
+
+sealed class HLInputByDiagnosticProvider<PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>, INPUT : HLApplicatorInput> {
+    abstract fun KtAnalysisSession.createInfo(diagnostic: DIAGNOSTIC): INPUT?
+}
+
+private class HLInputByDiagnosticProviderImpl<PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>, INPUT : HLApplicatorInput>(
+    private val createInfo: KtAnalysisSession.(DIAGNOSTIC) -> INPUT?
+) : HLInputByDiagnosticProvider<PSI, DIAGNOSTIC, INPUT>() {
+    override fun KtAnalysisSession.createInfo(diagnostic: DIAGNOSTIC): INPUT? =
+        createInfo.invoke(this, diagnostic)
+}
+
+fun <PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>, INPUT : HLApplicatorInput> inputByDiagnosticProvider(
+    createInfo: KtAnalysisSession.(DIAGNOSTIC) -> INPUT?
+): HLInputByDiagnosticProvider<PSI, DIAGNOSTIC, INPUT> =
+    HLInputByDiagnosticProviderImpl(createInfo)

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

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
 import org.jetbrains.kotlin.idea.frontend.api.diagnostics.getDefaultMessageWithFactoryName
 import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
 import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixService
+import org.jetbrains.kotlin.idea.frontend.api.components.KtDiagnosticCheckerFilter
 import org.jetbrains.kotlin.psi.KtFile
 
 class KotlinHighLevelDiagnosticHighlightingPass(
@@ -43,7 +44,7 @@ class KotlinHighLevelDiagnosticHighlightingPass(
 
     override fun doCollectInformation(progress: ProgressIndicator) {
         analyze(ktFile) {
-            ktFile.collectDiagnosticsForFile().forEach { diagnostic ->
+            ktFile.collectDiagnosticsForFile(KtDiagnosticCheckerFilter.ONLY_COMMON_CHECKERS).forEach { diagnostic ->
                 addDiagnostic(diagnostic)
             }
         }

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import kotlin.reflect.KClass
 
 /**
  * The entry point into all frontend-related work. Has the following contracts:
@@ -91,9 +92,11 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
 
     fun KtClassOrObjectSymbol.buildSelfClassType(): KtType = typeProvider.buildSelfClassType(this)
 
-    fun KtElement.getDiagnostics(): Collection<KtDiagnostic> = diagnosticProvider.getDiagnosticsForElement(this)
+    fun KtElement.getDiagnostics(filter: KtDiagnosticCheckerFilter): Collection<KtDiagnostic> =
+        diagnosticProvider.getDiagnosticsForElement(this, filter)
 
-    fun KtFile.collectDiagnosticsForFile(): Collection<KtDiagnosticWithPsi<*>> = diagnosticProvider.collectDiagnosticsForFile(this)
+    fun KtFile.collectDiagnosticsForFile(filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>> =
+        diagnosticProvider.collectDiagnosticsForFile(this, filter)
 
     fun KtSymbolWithKind.getContainingSymbol(): KtSymbolWithKind? = containingDeclarationProvider.getContainingDeclaration(this)
 

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

@@ -10,6 +10,12 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
 abstract class KtDiagnosticProvider : KtAnalysisSessionComponent() {
-    abstract fun getDiagnosticsForElement(element: KtElement): Collection<KtDiagnosticWithPsi<*>>
+    abstract fun getDiagnosticsForElement(element: KtElement, filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>>
-    abstract fun collectDiagnosticsForFile(ktFile: KtFile): Collection<KtDiagnosticWithPsi<*>>
+    abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>>
+}
+
+enum class KtDiagnosticCheckerFilter {
+    ONLY_COMMON_CHECKERS,
+    ONLY_EXTENDED_CHECKERS,
+    EXTENDED_AND_COMMON_CHECKERS,
 }

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

@@ -17,11 +17,13 @@ import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
+import org.jetbrains.kotlin.idea.frontend.api.components.KtDiagnosticCheckerFilter
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -84,11 +86,11 @@ internal class FirModuleResolveStateForCompletion(
     }
 
 
-    override fun getDiagnostics(element: KtElement): List<FirPsiDiagnostic<*>> {
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>> {
         error("Diagnostics should not be retrieved in completion")
     }
 
-    override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<FirPsiDiagnostic<*>> {
+    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> {
         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/FirModuleResolveStateImpl.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
@@ -69,11 +70,11 @@ internal class FirModuleResolveStateImpl(
     override fun getFirFile(ktFile: KtFile): FirFile =
         firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, lazyBodiesMode = false)
 
-    override fun getDiagnostics(element: KtElement): List<FirPsiDiagnostic<*>> =
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>> =
-        diagnosticsCollector.getDiagnosticsFor(element)
+        diagnosticsCollector.getDiagnosticsFor(element, filter)
 
-    override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<FirPsiDiagnostic<*>> =
+    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
-        diagnosticsCollector.collectDiagnosticsForFile(ktFile)
+        diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
 
     override fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile? {
         val cache = sessionProvider.getModuleCache(ktFile.getModuleInfo() as ModuleSourceInfo)
@@ -119,7 +120,7 @@ internal class FirModuleResolveStateImpl(
                 container,
                 cache,
                 FirResolvePhase.BODY_RESOLVE,
-                checkPCE = false /*TODO*/,
+                checkPCE = false, /*TODO*/
                 towerDataContextCollector = collector,
             )
         }

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

@@ -0,0 +1,12 @@
+/*
+ * Copyright 2010-2021 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.api
+
+enum class DiagnosticCheckerFilter(val runCommonCheckers: Boolean, val runExtendedCheckers: Boolean) {
+    ONLY_COMMON_CHECKERS(runCommonCheckers = true, runExtendedCheckers = false),
+    ONLY_EXTENDED_CHECKERS(runCommonCheckers = false, runExtendedCheckers = true),
+    EXTENDED_AND_COMMON_CHECKERS(runCommonCheckers = true, runExtendedCheckers = true),
+}

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

@@ -37,9 +37,9 @@ abstract class FirModuleResolveState {
 
     internal abstract fun isFirFileBuilt(ktFile: KtFile): Boolean
 
-    internal abstract fun getDiagnostics(element: KtElement): List<FirPsiDiagnostic<*>>
+    internal abstract fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>>
 
-    internal abstract fun collectDiagnosticsForFile(ktFile: KtFile): Collection<FirPsiDiagnostic<*>>
+    internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>>
 
     @TestOnly
     internal abstract fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile?

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
-import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
@@ -17,7 +16,6 @@ import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ktDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
@@ -76,13 +74,13 @@ inline fun <reified F : FirDeclaration, R> KtDeclaration.withFirDeclarationOfTyp
 }
 
 /**
-* Creates [FirDeclaration] by [KtLambdaExpression] and executes an [action] on it
+ * Creates [FirDeclaration] by [KtLambdaExpression] and executes an [action] on it
-*
+ *
-* If resulted [FirDeclaration] is not [F] throws [InvalidFirElementTypeException]
+ * If resulted [FirDeclaration] is not [F] throws [InvalidFirElementTypeException]
-*
+ *
-* [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
+ * [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
-* Otherwise, some threading problems may arise,
+ * Otherwise, some threading problems may arise,
-*/
+ */
 @OptIn(InternalForInline::class)
 inline fun <reified F : FirDeclaration, R> KtLambdaExpression.withFirDeclarationOfType(
     resolveState: FirModuleResolveState,
@@ -121,14 +119,17 @@ fun <D : FirDeclaration, R> D.withFirDeclaration(
 /**
  * Returns a list of Diagnostics compiler finds for given [KtElement]
  */
-fun KtElement.getDiagnostics(resolveState: FirModuleResolveState): Collection<FirPsiDiagnostic<*>> =
+fun KtElement.getDiagnostics(resolveState: FirModuleResolveState, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
-    resolveState.getDiagnostics(this)
+    resolveState.getDiagnostics(this, filter)
 
 /**
  * Returns a list of Diagnostics compiler finds for given [KtFile]
  */
-fun KtFile.collectDiagnosticsForFile(resolveState: FirModuleResolveState): Collection<FirPsiDiagnostic<*>> =
+fun KtFile.collectDiagnosticsForFile(
-    resolveState.collectDiagnosticsForFile(this)
+    resolveState: FirModuleResolveState,
+    filter: DiagnosticCheckerFilter
+): Collection<FirPsiDiagnostic<*>> =
+    resolveState.collectDiagnosticsForFile(this, filter)
 
 /**
  * Resolves a given [FirDeclaration] to [phase] and returns resolved declaration

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

@@ -5,24 +5,33 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
-import com.intellij.openapi.diagnostic.Logger
-import org.jetbrains.kotlin.diagnostics.Diagnostic
 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.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.registerAllComponents
+import org.jetbrains.kotlin.fir.analysis.collectors.components.ControlFlowAnalysisDiagnosticComponent
+import org.jetbrains.kotlin.fir.analysis.collectors.components.DeclarationCheckersDiagnosticComponent
+import org.jetbrains.kotlin.fir.analysis.collectors.components.ErrorNodeDiagnosticCollectorComponent
+import org.jetbrains.kotlin.fir.analysis.collectors.components.ExpressionCheckersDiagnosticComponent
 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.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
-import org.jetbrains.kotlin.psi.KtElement
 
 internal abstract class AbstractFirIdeDiagnosticsCollector(
     session: FirSession,
+    useExtendedCheckers: Boolean,
 ) : AbstractDiagnosticCollector(
     session,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
@@ -33,7 +42,16 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
     )
 ) {
     init {
-        registerAllComponents()
+        val declarationCheckers = CheckersFactory.createDeclarationCheckers(useExtendedCheckers)
+        val expressionCheckers = CheckersFactory.createExpressionCheckers(useExtendedCheckers)
+
+        @Suppress("LeakingThis")
+        initializeComponents(
+            DeclarationCheckersDiagnosticComponent(this, declarationCheckers),
+            ExpressionCheckersDiagnosticComponent(this, expressionCheckers),
+            ErrorNodeDiagnosticCollectorComponent(this),
+            ControlFlowAnalysisDiagnosticComponent(this, declarationCheckers),
+        )
     }
 
     protected abstract fun onDiagnostic(diagnostic: FirPsiDiagnostic<*>)
@@ -60,8 +78,23 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
         // Not necessary in IDE
         return emptyList()
     }
-
+}
-    companion object {
+
-        private val LOG = Logger.getInstance(AbstractFirIdeDiagnosticsCollector::class.java)
+
-    }
+private object CheckersFactory {
+    private val extendedDeclarationCheckers = createDeclarationCheckers(ExtendedDeclarationCheckers)
+    private val commonDeclarationCheckers = createDeclarationCheckers(CommonDeclarationCheckers)
+
+    fun createDeclarationCheckers(useExtendedCheckers: Boolean): DeclarationCheckers =
+        if (useExtendedCheckers) extendedDeclarationCheckers else commonDeclarationCheckers
+
+    fun createExpressionCheckers(useExtendedCheckers: Boolean): ExpressionCheckers =
+        if (useExtendedCheckers) ExtendedExpressionCheckers else CommonExpressionCheckers
+
+    // TODO hack to have all checkers present in DeclarationCheckers.memberDeclarationCheckers and similar
+    // If use ExtendedDeclarationCheckers directly when DeclarationCheckers.memberDeclarationCheckers will not contain basicDeclarationCheckers
+    @OptIn(SessionConfiguration::class)
+    private fun createDeclarationCheckers(declarationCheckers: DeclarationCheckers): DeclarationCheckers =
+        CheckersComponent().apply { register(declarationCheckers) }.declarationCheckers
+
 }

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

@@ -5,8 +5,8 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
-import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 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
@@ -16,14 +16,14 @@ internal class DiagnosticsCollector(
     private val fileStructureCache: FileStructureCache,
     private val cache: ModuleFileCache,
 ) {
-    fun getDiagnosticsFor(element: KtElement): List<FirPsiDiagnostic<*>> =
+    fun getDiagnosticsFor(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>> =
         fileStructureCache
             .getFileStructure(element.containingKtFile, cache)
             .getStructureElementFor(element)
-            .diagnostics.diagnosticsFor(element)
+            .diagnostics.diagnosticsFor(filter, element)
 
-    fun collectDiagnosticsForFile(ktFile: KtFile): Collection<FirPsiDiagnostic<*>> =
+    fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
         fileStructureCache
             .getFileStructure(ktFile, cache)
-            .getAllDiagnosticsForFile()
+            .getAllDiagnosticsForFile(filter)
 }

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

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2010-2021 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 com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+
+internal class FileStructureElementDiagnosticList(
+    private val map: Map<PsiElement, List<FirPsiDiagnostic<*>>>
+) {
+    fun diagnosticsFor(element: PsiElement): List<FirPsiDiagnostic<*>> = map[element] ?: emptyList()
+
+    inline fun forEach(action: (List<FirPsiDiagnostic<*>>) -> Unit) = map.values.forEach(action)
+}

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

@@ -0,0 +1,12 @@
+/*
+ * Copyright 2010-2021 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.fir.declarations.FirFile
+
+internal abstract class FileStructureElementDiagnosticRetriever {
+    abstract fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList
+}

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2010-2021 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 com.intellij.psi.PsiElement
+import com.intellij.util.SmartList
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
+
+internal class FileStructureElementDiagnostics(
+    private val firFile: FirFile,
+    private val retriever: FileStructureElementDiagnosticRetriever
+) {
+    private val diagnosticByCommonCheckers: FileStructureElementDiagnosticList by lazy {
+        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.USUAL_COLLECTOR)
+    }
+
+    private val diagnosticByExtendedCheckers: FileStructureElementDiagnosticList by lazy {
+        retriever.retrieve(firFile, FileStructureElementDiagnosticsCollector.EXTENDED_COLLECTOR)
+    }
+
+    fun diagnosticsFor(filter: DiagnosticCheckerFilter, element: PsiElement): List<FirPsiDiagnostic<*>> =
+        SmartList<FirPsiDiagnostic<*>>().apply {
+            if (filter.runCommonCheckers) {
+                addAll(diagnosticByCommonCheckers.diagnosticsFor(element))
+            }
+            if (filter.runExtendedCheckers) {
+                addAll(diagnosticByExtendedCheckers.diagnosticsFor(element))
+            }
+        }
+
+
+    inline fun forEach(filter: DiagnosticCheckerFilter, action: (List<FirPsiDiagnostic<*>>) -> Unit) {
+        if (filter.runCommonCheckers) {
+            diagnosticByCommonCheckers.forEach(action)
+        }
+        if (filter.runExtendedCheckers) {
+            diagnosticByExtendedCheckers.forEach(action)
+        }
+    }
+}

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

@@ -0,0 +1,61 @@
+/*
+ * 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 com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
+
+internal class FileStructureElementDiagnosticsCollector private constructor(private val useExtendedCheckers: Boolean) {
+    companion object {
+        val USUAL_COLLECTOR = FileStructureElementDiagnosticsCollector(useExtendedCheckers = false)
+        val EXTENDED_COLLECTOR = FileStructureElementDiagnosticsCollector(useExtendedCheckers = true)
+    }
+
+    fun collectForStructureElement(
+        firFile: FirFile,
+        onDeclarationExit: (FirDeclaration) -> Unit = {},
+        onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction,
+    ): FileStructureElementDiagnosticList =
+        FirIdeStructureElementDiagnosticsCollector(
+            firFile.session,
+            useExtendedCheckers,
+            onDeclarationEnter,
+            onDeclarationExit
+        ).let { collector ->
+            collector.collectDiagnostics(firFile)
+            FileStructureElementDiagnosticList(collector.result)
+        }
+
+    private class FirIdeStructureElementDiagnosticsCollector(
+        session: FirSession,
+        useExtendedCheckers: Boolean,
+        private val onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction,
+        private val onDeclarationExit: (FirDeclaration) -> Unit
+    ) : AbstractFirIdeDiagnosticsCollector(
+        session,
+        useExtendedCheckers,
+    ) {
+        val result = mutableMapOf<PsiElement, MutableList<FirPsiDiagnostic<*>>>()
+
+        override fun onDiagnostic(diagnostic: FirPsiDiagnostic<*>) {
+            result.addValueFor(diagnostic.psiElement, diagnostic)
+        }
+
+        override fun onDeclarationEnter(
+            declaration: FirDeclaration,
+        ): DiagnosticCollectorDeclarationAction =
+            onDeclarationEnter.invoke(declaration)
+
+        override fun onDeclarationExit(declaration: FirDeclaration) {
+            onDeclarationExit.invoke(declaration)
+        }
+    }
+}

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

@@ -5,15 +5,16 @@
 
 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.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
 
 internal class FirIdeFileDiagnosticsCollector private constructor(
     session: FirSession,
+    useExtendedCheckers: Boolean,
 ) : AbstractFirIdeDiagnosticsCollector(
     session,
+    useExtendedCheckers,
 ) {
     private val result = mutableListOf<FirPsiDiagnostic<*>>()
 
@@ -22,8 +23,8 @@ internal class FirIdeFileDiagnosticsCollector private constructor(
     }
 
     companion object {
-        fun collectForFile(firFile: FirFile): List<FirPsiDiagnostic<*>> =
+        fun collectForFile(firFile: FirFile, useExtendedCheckers: Boolean): List<FirPsiDiagnostic<*>> =
-            FirIdeFileDiagnosticsCollector(firFile.session).let { collector ->
+            FirIdeFileDiagnosticsCollector(firFile.session, useExtendedCheckers).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

@@ -1,78 +0,0 @@
-/*
- * 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 com.intellij.psi.PsiElement
-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.analysis.diagnostics.FirPsiDiagnostic
-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<PsiElement, MutableList<FirPsiDiagnostic<*>>>()
-
-    override fun onDiagnostic(diagnostic: FirPsiDiagnostic<*>) {
-        result.addValueFor(diagnostic.psiElement, 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)
-            }
-
-        fun collectForSingleDeclaration(firFile: FirFile, declaration: FirDeclaration): FileStructureElementDiagnostics {
-            var inCurrentDeclaration = false
-
-            return collectForStructureElement(
-                firFile,
-                onDeclarationEnter = { firDeclaration ->
-                    when {
-                        firDeclaration == declaration -> {
-                            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 = { firDeclaration ->
-                    if (declaration == firDeclaration) {
-                        inCurrentDeclaration = false
-                    }
-                }
-            )
-        }
-    }
-}

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -56,10 +57,21 @@ internal class FileStructure(
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    fun getAllDiagnosticsForFile(): Collection<FirPsiDiagnostic<*>> {
+    fun getAllDiagnosticsForFile(diagnosticCheckerFilter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> {
         val structureElements = getAllStructureElements()
         return buildSet {
-            structureElements.forEach { it.diagnostics.forEach { diagnostics -> addAll(diagnostics) } }
+            collectDiagnosticsFromStructureElements(structureElements, diagnosticCheckerFilter)
+        }
+    }
+
+    private fun MutableSet<FirPsiDiagnostic<*>>.collectDiagnosticsFromStructureElements(
+        structureElements: Collection<FileStructureElement>,
+        diagnosticCheckerFilter: DiagnosticCheckerFilter
+    ) {
+        structureElements.forEach { structureElement ->
+            structureElement.diagnostics.forEach(diagnosticCheckerFilter) { diagnostics ->
+                addAll(diagnostics)
+            }
         }
     }
 

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

@@ -5,18 +5,17 @@
 
 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
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
-import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FirIdeStructureElementDiagnosticsCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnosticList
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnosticRetriever
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnostics
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
@@ -25,22 +24,13 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.isGeneratedDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.ktDeclaration
 import org.jetbrains.kotlin.psi.*
 
-internal class FileStructureElementDiagnostics(
+internal sealed class FileStructureElement(val firFile: FirFile) {
-    private val map: Map<PsiElement, List<FirPsiDiagnostic<*>>>
-) {
-    fun diagnosticsFor(element: PsiElement): List<FirPsiDiagnostic<*>> = map[element] ?: emptyList()
-
-    inline fun forEach(action: (List<FirPsiDiagnostic<*>>) -> Unit) = map.values.forEach(action)
-}
-
-internal sealed class FileStructureElement {
-    abstract val firFile: FirFile
     abstract val psi: KtAnnotated
     abstract val mappings: Map<KtElement, FirElement>
     abstract val diagnostics: FileStructureElementDiagnostics
 }
 
-internal sealed class ReanalyzableStructureElement<KT : KtDeclaration> : FileStructureElement() {
+internal sealed class ReanalyzableStructureElement<KT : KtDeclaration>(firFile: FirFile) : FileStructureElement(firFile) {
     abstract override val psi: KtDeclaration
     abstract val firSymbol: AbstractFirBasedSymbol<*>
     abstract val timestamp: Long
@@ -58,8 +48,31 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration> : FileStr
 
     fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
 
-    override val diagnostics: FileStructureElementDiagnostics by lazy {
+    override val diagnostics = FileStructureElementDiagnostics(firFile, FileStructureElementSingleDeclarationDiagnosticRetriever())
-        FirIdeStructureElementDiagnosticsCollector.collectForSingleDeclaration(firFile, firSymbol.fir as FirDeclaration)
+
+    inner class FileStructureElementSingleDeclarationDiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
+        override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
+            var inCurrentDeclaration = false
+            val declaration = firSymbol.fir as FirDeclaration
+            return collector.collectForStructureElement(
+                firFile,
+                onDeclarationEnter = { firDeclaration ->
+                    when {
+                        firDeclaration == declaration -> {
+                            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 = { firDeclaration ->
+                    if (declaration == firDeclaration) {
+                        inCurrentDeclaration = false
+                    }
+                }
+            )
+        }
     }
 
     companion object {
@@ -68,11 +81,11 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration> : FileStr
 }
 
 internal class ReanalyzableFunctionStructureElement(
-    override val firFile: FirFile,
+    firFile: FirFile,
     override val psi: KtNamedFunction,
     override val firSymbol: FirFunctionSymbol<*>,
     override val timestamp: Long
-) : ReanalyzableStructureElement<KtNamedFunction>() {
+) : ReanalyzableStructureElement<KtNamedFunction>(firFile) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
@@ -106,11 +119,11 @@ internal class ReanalyzableFunctionStructureElement(
 }
 
 internal class ReanalyzablePropertyStructureElement(
-    override val firFile: FirFile,
+    firFile: FirFile,
     override val psi: KtProperty,
     override val firSymbol: FirPropertySymbol,
     override val timestamp: Long
-) : ReanalyzableStructureElement<KtProperty>() {
+) : ReanalyzableStructureElement<KtProperty>(firFile) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
@@ -144,42 +157,47 @@ internal class ReanalyzablePropertyStructureElement(
 }
 
 internal class NonReanalyzableDeclarationStructureElement(
-    override val firFile: FirFile,
+    firFile: FirFile,
-    fir: FirDeclaration,
+    private val fir: FirDeclaration,
     override val psi: KtDeclaration,
-) : FileStructureElement() {
+) : FileStructureElement(firFile) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(fir, recorder)
 
-    override val diagnostics: FileStructureElementDiagnostics by lazy {
+    override val diagnostics = FileStructureElementDiagnostics(firFile, DiagnosticRetriever())
-        var inCurrentDeclaration = false
+
-        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(
+    private inner class DiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
-            firFile,
+        override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
-            onDeclarationEnter = { firDeclaration ->
+            var inCurrentDeclaration = false
-                when {
+            return collector.collectForStructureElement(
-                    firDeclaration.isGeneratedDeclaration -> DiagnosticCollectorDeclarationAction.SKIP
+                firFile,
-                    firDeclaration is FirFile -> DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                onDeclarationEnter = { firDeclaration ->
-                    firDeclaration == fir -> {
+                    when {
-                        inCurrentDeclaration = true
+                        firDeclaration.isGeneratedDeclaration -> DiagnosticCollectorDeclarationAction.SKIP
-                        DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                        firDeclaration is FirFile -> DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                        firDeclaration == fir -> {
+                            inCurrentDeclaration = true
+                            DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                        }
+                        FileElementFactory.isReanalyzableContainer(firDeclaration.ktDeclaration) -> {
+                            DiagnosticCollectorDeclarationAction.SKIP
+                        }
+                        inCurrentDeclaration -> {
+                            DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                        }
+                        else -> DiagnosticCollectorDeclarationAction.SKIP_CURRENT_DECLARATION_AND_CHECK_NESTED
                     }
-                    FileElementFactory.isReanalyzableContainer(firDeclaration.ktDeclaration) -> {
+                },
-                        DiagnosticCollectorDeclarationAction.SKIP
+                onDeclarationExit = { firDeclaration ->
+                    if (firDeclaration == fir) {
+                        inCurrentDeclaration = false
                     }
-                    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 {
         private val recorder = object : FirElementsRecorder() {
             override fun visitProperty(property: FirProperty, data: MutableMap<KtElement, FirElement>) {
@@ -200,17 +218,21 @@ internal class NonReanalyzableDeclarationStructureElement(
 }
 
 
-internal data class RootStructureElement(
+internal class RootStructureElement(
-    override val firFile: FirFile,
+    firFile: FirFile,
     override val psi: KtFile,
-) : FileStructureElement() {
+) : FileStructureElement(firFile) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firFile, recorder)
 
-    override val diagnostics: FileStructureElementDiagnostics by lazy {
+    override val diagnostics = FileStructureElementDiagnostics(firFile, DiagnosticRetriever)
-        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(firFile) { firDeclaration ->
+
-            if (firDeclaration is FirFile) DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_SKIP_NESTED
+    private object DiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
-            else DiagnosticCollectorDeclarationAction.SKIP
+        override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
+            return collector.collectForStructureElement(firFile) { firDeclaration ->
+                if (firDeclaration is FirFile) DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_SKIP_NESTED
+                else DiagnosticCollectorDeclarationAction.SKIP
+            }
         }
     }
 

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

@@ -0,0 +1,7 @@
+/*
+ * Copyright 2010-2021 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.file.structure
+

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

@@ -0,0 +1,7 @@
+/*
+ * Copyright 2010-2021 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.file.structure
+

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

@@ -0,0 +1,7 @@
+/*
+ * Copyright 2010-2021 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.file.structure
+

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmTypeMapper
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
 import org.jetbrains.kotlin.fir.checkers.registerCommonCheckers
+import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.dependenciesWithoutSelf
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
@@ -133,7 +134,7 @@ internal object FirIdeSessionFactory {
             registerJavaSpecificResolveComponents()
             FirSessionFactory.FirSessionConfigurator(this).apply {
                 if (isRootModule) {
-                    registerCommonCheckers()
+                    registerExtendedCommonCheckers()
                 }
             }.configure()
         }