FIR IDE: separate diagnostics collection and checker context collection

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

@@ -20,12 +20,12 @@ abstract class AbstractDiagnosticCollector(
     override val session: FirSession,
     override val scopeSession: ScopeSession = ScopeSession(),
 ) : SessionHolder {
-    fun collectDiagnostics(firFile: FirFile): List<FirDiagnostic<*>> {
+    fun collectDiagnostics(firDeclaration: FirDeclaration): List<FirDiagnostic<*>> {
         if (!componentsInitialized) {
             throw IllegalStateException("Components are not initialized")
         }
         initializeCollector()
-        firFile.accept(visitor, null)
+        firDeclaration.accept(visitor, null)
         return getCollectedDiagnostics()
     }
 

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

@@ -7,24 +7,236 @@ package org.jetbrains.kotlin.fir.analysis.collectors
 
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
+import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.collectImplicitReceivers
+import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.name.Name
 
 abstract class AbstractDiagnosticCollectorVisitor(
-    @set:PrivateForInline var context: PersistentCheckerContext
+    @set:PrivateForInline var context: PersistentCheckerContext,
+    protected val components: List<AbstractDiagnosticCollectorComponent>,
 ) : FirDefaultVisitor<Unit, Nothing?>() {
+    protected var currentAction = DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
+
+    protected open fun beforeRunningAllComponentsOnElement(element: FirElement) {}
+    protected open fun beforeRunningSingleComponentOnElement(element: FirElement) {}
+
+    protected abstract fun goToNestedDeclarations(element: FirElement)
+    protected abstract fun runComponents(element: FirElement)
+
+    override fun visitElement(element: FirElement, data: Nothing?) {
+        if (element is FirAnnotationContainer) {
+            visitAnnotationContainer(element, data)
+            return
+        }
+        runComponents(element)
+        goToNestedDeclarations(element)
+    }
+
+    override fun visitAnnotationContainer(annotationContainer: FirAnnotationContainer, data: Nothing?) {
+        withSuppressedDiagnostics(annotationContainer) {
+            runComponents(annotationContainer)
+            goToNestedDeclarations(annotationContainer)
+        }
+    }
+
+    private fun visitJump(loopJump: FirLoopJump) {
+        withSuppressedDiagnostics(loopJump) {
+            runComponents(loopJump)
+            loopJump.target.labeledElement.takeIf { it is FirErrorLoop }?.accept(this, null)
+        }
+    }
+
+    override fun visitBreakExpression(breakExpression: FirBreakExpression, data: Nothing?) {
+        visitJump(breakExpression)
+    }
+
+    override fun visitContinueExpression(continueExpression: FirContinueExpression, data: Nothing?) {
+        visitJump(continueExpression)
+    }
+
+    private fun visitClassAndChildren(klass: FirClass<*>, type: ConeKotlinType) {
+        val typeRef = buildResolvedTypeRef {
+            this.type = type
+        }
+        visitWithDeclarationAndReceiver(klass, (klass as? FirRegularClass)?.name, typeRef)
+    }
+
+    override fun visitRegularClass(regularClass: FirRegularClass, data: Nothing?) {
+        withSuppressedDiagnostics(regularClass) {
+            visitClassAndChildren(regularClass, regularClass.defaultType())
+        }
+    }
+
+    override fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: Nothing?) {
+        withSuppressedDiagnostics(anonymousObject) {
+            visitClassAndChildren(anonymousObject, anonymousObject.defaultType())
+        }
+    }
+
+    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?) {
+        withSuppressedDiagnostics(simpleFunction) {
+            visitWithDeclarationAndReceiver(simpleFunction, simpleFunction.name, simpleFunction.receiverTypeRef)
+        }
+    }
+
+    override fun visitConstructor(constructor: FirConstructor, data: Nothing?) {
+        withSuppressedDiagnostics(constructor) {
+            visitWithDeclaration(constructor)
+        }
+    }
+
+    override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Nothing?) {
+        withSuppressedDiagnostics(anonymousFunction) {
+            val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }
+            visitWithDeclarationAndReceiver(
+                anonymousFunction,
+                labelName,
+                anonymousFunction.receiverTypeRef
+            )
+        }
+    }
+
+    override fun visitProperty(property: FirProperty, data: Nothing?) {
+        withSuppressedDiagnostics(property) {
+            visitWithDeclaration(property)
+        }
+    }
+
+    override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: Nothing?) {
+        if (propertyAccessor !is FirDefaultPropertyAccessor) {
+            val property = context.containingDeclarations.last() as FirProperty
+            withSuppressedDiagnostics(propertyAccessor) {
+                visitWithDeclarationAndReceiver(propertyAccessor, property.name, property.receiverTypeRef)
+            }
+        }
+    }
+
+    override fun visitValueParameter(valueParameter: FirValueParameter, data: Nothing?) {
+        withSuppressedDiagnostics(valueParameter) {
+            visitWithDeclaration(valueParameter)
+        }
+    }
+
+    override fun visitEnumEntry(enumEntry: FirEnumEntry, data: Nothing?) {
+        withSuppressedDiagnostics(enumEntry) {
+            visitWithDeclaration(enumEntry)
+        }
+    }
+
+    override fun visitFile(file: FirFile, data: Nothing?) {
+        withSuppressedDiagnostics(file) {
+            visitWithDeclaration(file)
+        }
+    }
+
+    override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer, data: Nothing?) {
+        visitWithDeclaration(anonymousInitializer)
+    }
+
+    override fun visitBlock(block: FirBlock, data: Nothing?) {
+        withSuppressedDiagnostics(block) {
+            visitExpression(block, data)
+        }
+    }
+
+    override fun visitTypeRef(typeRef: FirTypeRef, data: Nothing?) {
+        if (typeRef.source != null && typeRef.source?.kind !is FirFakeSourceElementKind) {
+            withSuppressedDiagnostics(typeRef) {
+                runComponents(typeRef)
+                goToNestedDeclarations(typeRef)
+            }
+        }
+    }
+
+    override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: Nothing?) {
+        super.visitResolvedTypeRef(resolvedTypeRef, data)
+        resolvedTypeRef.delegatedTypeRef?.accept(this, data)
+    }
+
+    override fun visitFunctionCall(functionCall: FirFunctionCall, data: Nothing?) {
+        visitWithQualifiedAccess(functionCall)
+    }
+
+    override fun visitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression, data: Nothing?) {
+        visitWithQualifiedAccess(qualifiedAccessExpression)
+    }
+
+    override fun visitGetClassCall(getClassCall: FirGetClassCall, data: Nothing?) {
+        visitWithGetClassCall(getClassCall)
+    }
+
+    private inline fun visitWithDeclaration(
+        declaration: FirDeclaration,
+        block: () -> Unit = { declaration.acceptChildren(this, null) }
+    ) {
+        if (!currentAction.lookupForNestedDeclaration) return
+
+        val action = getDeclarationActionOnDeclarationEnter(declaration)
+        withDiagnosticsAction(action) {
+            runComponents(declaration)
+            withDeclaration(declaration) {
+                block()
+            }
+        }
+        onDeclarationExit(declaration)
+    }
+
+    private fun visitWithDeclarationAndReceiver(declaration: FirDeclaration, labelName: Name?, receiverTypeRef: FirTypeRef?) {
+        visitWithDeclaration(declaration) {
+            withLabelAndReceiverType(
+                labelName,
+                declaration,
+                receiverTypeRef?.coneTypeSafe()
+            ) {
+                goToNestedDeclarations(declaration)
+            }
+        }
+    }
+
+    private fun visitWithQualifiedAccess(qualifiedAccess: FirQualifiedAccess) {
+        return withQualifiedAccess(qualifiedAccess) {
+            runComponents(qualifiedAccess)
+            goToNestedDeclarations(qualifiedAccess)
+        }
+    }
+
+    private fun visitWithGetClassCall(getClassCall: FirGetClassCall) {
+        return withGetClassCall(getClassCall) {
+            runComponents(getClassCall)
+            goToNestedDeclarations(getClassCall)
+        }
+    }
+
+    protected open fun getDeclarationActionOnDeclarationEnter(declaration: FirDeclaration): DiagnosticCollectorDeclarationAction =
+        DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
+
+    protected open fun onDeclarationExit(declaration: FirDeclaration) {}
+
+    private inline fun <R> withDiagnosticsAction(action: DiagnosticCollectorDeclarationAction, block: () -> R): R {
+        val oldAction = currentAction
+        currentAction = action
+        return try {
+            block()
+        } finally {
+            currentAction = oldAction
+        }
+    }
 
     @OptIn(PrivateForInline::class)
-    protected inline fun <R> withQualifiedAccess(qualifiedAccess: FirQualifiedAccess, block: () -> R): R {
+    inline fun <R> withQualifiedAccess(qualifiedAccess: FirQualifiedAccess, block: () -> R): R {
         val existingContext = context
         context = context.addQualifiedAccess(qualifiedAccess)
         try {
@@ -36,7 +248,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
 
     @OptIn(PrivateForInline::class)
-    protected inline fun <R> withGetClassCall(getClassCall: FirGetClassCall, block: () -> R): R {
+    inline fun <R> withGetClassCall(getClassCall: FirGetClassCall, block: () -> R): R {
         val existingContext = context
         context = context.addGetClassCall(getClassCall)
         try {
@@ -48,7 +260,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
 
     @OptIn(PrivateForInline::class)
-    protected inline fun <R> withDeclaration(declaration: FirDeclaration, block: () -> R): R {
+    inline fun <R> withDeclaration(declaration: FirDeclaration, block: () -> R): R {
         val existingContext = context
         context = context.addDeclaration(declaration)
         try {
@@ -60,7 +272,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
 
     @OptIn(PrivateForInline::class)
-    protected inline fun <R> withLabelAndReceiverType(
+    inline fun <R> withLabelAndReceiverType(
         labelName: Name?,
         owner: FirDeclaration,
         type: ConeKotlinType?,
@@ -83,7 +295,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
 
     @OptIn(PrivateForInline::class)
-    protected inline fun <R> withSuppressedDiagnostics(annotationContainer: FirAnnotationContainer, block: () -> R): R {
+    inline fun <R> withSuppressedDiagnostics(annotationContainer: FirAnnotationContainer, block: () -> R): R {
         val existingContext = context
         addSuppressedDiagnosticsToContext(annotationContainer)
         return try {
@@ -95,7 +307,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
 
 
     @OptIn(PrivateForInline::class)
-    protected fun addSuppressedDiagnosticsToContext(annotationContainer: FirAnnotationContainer) {
+    fun addSuppressedDiagnosticsToContext(annotationContainer: FirAnnotationContainer) {
         val arguments = AbstractDiagnosticCollector.getDiagnosticsSuppressedForContainer(annotationContainer) ?: return
         context = context.addSuppressedDiagnostics(
             arguments,

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

@@ -5,237 +5,27 @@
 
 package org.jetbrains.kotlin.fir.analysis.collectors
 
-import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
-import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.resolve.defaultType
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
-import org.jetbrains.kotlin.name.Name
 
 open class CheckerRunningDiagnosticCollectorVisitor(
     context: PersistentCheckerContext,
-    private val components: List<AbstractDiagnosticCollectorComponent>
-) : AbstractDiagnosticCollectorVisitor(context) {
-    private var currentAction = DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
+    components: List<AbstractDiagnosticCollectorComponent>
+) : AbstractDiagnosticCollectorVisitor(context, components) {
+    protected val session get() = context.session
 
-    protected open fun beforeRunningAllComponentsOnElement(element: FirElement) {}
-    protected open fun beforeRunningSingleComponentOnElement(element: FirElement) {}
-
-    private fun <T : FirElement> T.runComponents() {
-        if (currentAction.checkInCurrentDeclaration) {
-            beforeRunningAllComponentsOnElement(this)
-            components.forEach {
-                beforeRunningSingleComponentOnElement(this)
-                this.accept(it, context)
-            }
-        }
-    }
-
-    override fun visitElement(element: FirElement, data: Nothing?) {
-        if (element is FirAnnotationContainer) {
-            visitAnnotationContainer(element, data)
-            return
-        }
-        element.runComponents()
+    override fun goToNestedDeclarations(element: FirElement) {
         element.acceptChildren(this, null)
     }
 
-    override fun visitAnnotationContainer(annotationContainer: FirAnnotationContainer, data: Nothing?) {
-        withSuppressedDiagnostics(annotationContainer) {
-            annotationContainer.runComponents()
-            annotationContainer.acceptChildren(this, null)
-        }
-    }
-
-    private fun visitJump(loopJump: FirLoopJump) {
-        withSuppressedDiagnostics(loopJump) {
-            loopJump.runComponents()
-            loopJump.target.labeledElement.takeIf { it is FirErrorLoop }?.accept(this, null)
-        }
-    }
-
-    override fun visitBreakExpression(breakExpression: FirBreakExpression, data: Nothing?) {
-        visitJump(breakExpression)
-    }
-
-    override fun visitContinueExpression(continueExpression: FirContinueExpression, data: Nothing?) {
-        visitJump(continueExpression)
-    }
-
-    private fun visitClassAndChildren(klass: FirClass<*>, type: ConeKotlinType) {
-        val typeRef = buildResolvedTypeRef {
-            this.type = type
-        }
-        visitWithDeclarationAndReceiver(klass, (klass as? FirRegularClass)?.name, typeRef)
-    }
-
-    override fun visitRegularClass(regularClass: FirRegularClass, data: Nothing?) {
-        withSuppressedDiagnostics(regularClass) {
-            visitClassAndChildren(regularClass, regularClass.defaultType())
-        }
-    }
-
-    override fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: Nothing?) {
-        withSuppressedDiagnostics(anonymousObject) {
-            visitClassAndChildren(anonymousObject, anonymousObject.defaultType())
-        }
-    }
-
-    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?) {
-        withSuppressedDiagnostics(simpleFunction) {
-            visitWithDeclarationAndReceiver(simpleFunction, simpleFunction.name, simpleFunction.receiverTypeRef)
-        }
-    }
-
-    override fun visitConstructor(constructor: FirConstructor, data: Nothing?) {
-        withSuppressedDiagnostics(constructor) {
-            visitWithDeclaration(constructor)
-        }
-    }
-
-    override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Nothing?) {
-        withSuppressedDiagnostics(anonymousFunction) {
-            val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }
-            visitWithDeclarationAndReceiver(
-                anonymousFunction,
-                labelName,
-                anonymousFunction.receiverTypeRef
-            )
-        }
-    }
-
-    override fun visitProperty(property: FirProperty, data: Nothing?) {
-        withSuppressedDiagnostics(property) {
-            visitWithDeclaration(property)
-        }
-    }
-
-    override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: Nothing?) {
-        if (propertyAccessor !is FirDefaultPropertyAccessor) {
-            val property = context.containingDeclarations.last() as FirProperty
-            withSuppressedDiagnostics(propertyAccessor) {
-                visitWithDeclarationAndReceiver(propertyAccessor, property.name, property.receiverTypeRef)
+    override fun runComponents(element: FirElement) {
+        if (currentAction.checkInCurrentDeclaration) {
+            beforeRunningAllComponentsOnElement(element)
+            components.forEach {
+                beforeRunningSingleComponentOnElement(element)
+                element.accept(it, context)
             }
         }
     }
-
-    override fun visitValueParameter(valueParameter: FirValueParameter, data: Nothing?) {
-        withSuppressedDiagnostics(valueParameter) {
-            visitWithDeclaration(valueParameter)
-        }
-    }
-
-    override fun visitEnumEntry(enumEntry: FirEnumEntry, data: Nothing?) {
-        withSuppressedDiagnostics(enumEntry) {
-            visitWithDeclaration(enumEntry)
-        }
-    }
-
-    override fun visitFile(file: FirFile, data: Nothing?) {
-        withSuppressedDiagnostics(file) {
-            visitWithDeclaration(file)
-        }
-    }
-
-    override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer, data: Nothing?) {
-        visitWithDeclaration(anonymousInitializer)
-    }
-
-    override fun visitBlock(block: FirBlock, data: Nothing?) {
-        withSuppressedDiagnostics(block) {
-            visitExpression(block, data)
-        }
-    }
-
-    override fun visitTypeRef(typeRef: FirTypeRef, data: Nothing?) {
-        if (typeRef.source != null && typeRef.source?.kind !is FirFakeSourceElementKind) {
-            withSuppressedDiagnostics(typeRef) {
-                typeRef.runComponents()
-                typeRef.acceptChildren(this, data)
-            }
-        }
-    }
-
-    override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: Nothing?) {
-        super.visitResolvedTypeRef(resolvedTypeRef, data)
-        resolvedTypeRef.delegatedTypeRef?.accept(this, data)
-    }
-
-    override fun visitFunctionCall(functionCall: FirFunctionCall, data: Nothing?) {
-        visitWithQualifiedAccess(functionCall)
-    }
-
-    override fun visitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression, data: Nothing?) {
-        visitWithQualifiedAccess(qualifiedAccessExpression)
-    }
-
-    override fun visitGetClassCall(getClassCall: FirGetClassCall, data: Nothing?) {
-        visitWithGetClassCall(getClassCall)
-    }
-
-    private inline fun visitWithDeclaration(
-        declaration: FirDeclaration,
-        block: () -> Unit = { declaration.acceptChildren(this, null) }
-    ) {
-        if (!currentAction.lookupForNestedDeclaration) return
-
-        val action = getDeclarationActionOnDeclarationEnter(declaration)
-        withDiagnosticsAction(action) {
-            declaration.runComponents()
-            withDeclaration(declaration) {
-                block()
-            }
-        }
-        onDeclarationExit(declaration)
-    }
-
-    private fun visitWithDeclarationAndReceiver(declaration: FirDeclaration, labelName: Name?, receiverTypeRef: FirTypeRef?) {
-        visitWithDeclaration(declaration) {
-            withLabelAndReceiverType(
-                labelName,
-                declaration,
-                receiverTypeRef?.coneTypeSafe()
-            ) {
-                declaration.acceptChildren(this, null)
-            }
-        }
-    }
-
-    private fun visitWithQualifiedAccess(qualifiedAccess: FirQualifiedAccess) {
-        return withQualifiedAccess(qualifiedAccess) {
-            qualifiedAccess.runComponents()
-            qualifiedAccess.acceptChildren(this, null)
-        }
-    }
-
-    private fun visitWithGetClassCall(getClassCall: FirGetClassCall) {
-        return withGetClassCall(getClassCall) {
-            getClassCall.runComponents()
-            getClassCall.acceptChildren(this, null)
-        }
-    }
-
-    protected open fun getDeclarationActionOnDeclarationEnter(declaration: FirDeclaration): DiagnosticCollectorDeclarationAction =
-        DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-
-    protected open fun onDeclarationExit(declaration: FirDeclaration) {}
-
-    private inline fun <R> withDiagnosticsAction(action: DiagnosticCollectorDeclarationAction, block: () -> R): R {
-        val oldAction = currentAction
-        currentAction = action
-        return try {
-            block()
-        } finally {
-            currentAction = oldAction
-        }
-    }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt

@@ -36,6 +36,8 @@ interface SessionHolder {
     val scopeSession: ScopeSession
 }
 
+data class SessionHolderImpl(override val session: FirSession, override val scopeSession: ScopeSession): SessionHolder
+
 abstract class BodyResolveComponents : SessionHolder {
     abstract val returnTypeCalculator: ReturnTypeCalculator
     abstract val implicitReceiverStack: ImplicitReceiverStack

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerConte
 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.FirDiagnostic
@@ -34,23 +35,6 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
 ) : AbstractDiagnosticCollector(
     session
 ) {
-    private val beforeElementDiagnosticCollectionHandler: BeforeElementDiagnosticCollectionHandler? =
-        session.beforeElementDiagnosticCollectionHandler
-
-    @Suppress("LeakingThis")
-    override val visitor = run {
-        val returnTypeCalculator = createReturnTypeCalculatorForIDE(
-            session,
-            ScopeSession(),
-            ImplicitBodyResolveComputationSession(),
-            ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-        )
-        CollectingVisitor(
-            PersistentCheckerContext(this, returnTypeCalculator),
-            components
-        )
-    }
-
     init {
         val declarationCheckers = CheckersFactory.createDeclarationCheckers(useExtendedCheckers)
         val expressionCheckers = CheckersFactory.createExpressionCheckers(useExtendedCheckers)
@@ -80,19 +64,6 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
         reporter = Reporter()
     }
 
-    inner class CollectingVisitor(
-        context: PersistentCheckerContext,
-        components: List<AbstractDiagnosticCollectorComponent>
-    ) : DiagnosticCollectingVisitor(context, components) {
-        override fun beforeRunningSingleComponentOnElement(element: FirElement) {
-            checkCanceled()
-        }
-
-        override fun beforeRunningAllComponentsOnElement(element: FirElement) {
-            beforeElementDiagnosticCollectionHandler?.beforeCollectingForElement(element)
-        }
-    }
-
 
     override fun getCollectedDiagnostics(): List<FirDiagnostic<*>> {
         // Not necessary in IDE

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

@@ -5,8 +5,22 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistenceContextCollector
 
 internal abstract class FileStructureElementDiagnosticRetriever {
     abstract fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList
+}
+
+internal class SingleNonLocalDeclarationDiagnosticRetriever(
+    private val declaration: FirDeclaration
+) : FileStructureElementDiagnosticRetriever() {
+    override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
+        val sessionHolder = SessionHolderImpl(firFile.session, ScopeSession())
+        val context = PersistenceContextCollector.collectContext(sessionHolder, firFile, declaration)
+        return collector.collectForStructureElement(declaration, context)
+    }
 }

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

@@ -7,10 +7,12 @@ 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.checkers.context.PersistentCheckerContext
 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.diagnostics.fir.PersistentCheckerContextFactory
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
 
 internal class FileStructureElementDiagnosticsCollector private constructor(private val useExtendedCheckers: Boolean) {
@@ -20,22 +22,27 @@ internal class FileStructureElementDiagnosticsCollector private constructor(priv
     }
 
     fun collectForStructureElement(
-        firFile: FirFile,
+        firDeclaration: FirDeclaration,
+        initialContext: PersistentCheckerContext?,
         onDeclarationExit: (FirDeclaration) -> Unit = {},
-        onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction,
+        onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction = {
+            DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
+        },
     ): FileStructureElementDiagnosticList =
         FirIdeStructureElementDiagnosticsCollector(
-            firFile.session,
+            firDeclaration.session,
+            initialContext,
             useExtendedCheckers,
             onDeclarationEnter,
             onDeclarationExit
         ).let { collector ->
-            collector.collectDiagnostics(firFile)
+            collector.collectDiagnostics(firDeclaration)
             FileStructureElementDiagnosticList(collector.result)
         }
 
     private class FirIdeStructureElementDiagnosticsCollector(
         session: FirSession,
+        initialContext: PersistentCheckerContext?,
         useExtendedCheckers: Boolean,
         private val onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction,
         private val onDeclarationExit: (FirDeclaration) -> Unit
@@ -49,13 +56,18 @@ internal class FileStructureElementDiagnosticsCollector private constructor(priv
             result.addValueFor(diagnostic.psiElement, diagnostic)
         }
 
-        override fun getDeclarationActionOnDeclarationEnter(
-            declaration: FirDeclaration,
-        ): DiagnosticCollectorDeclarationAction =
-            onDeclarationEnter.invoke(declaration)
+        override val visitor = object : FirIdeDiagnosticVisitor(
+            initialContext ?: PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(this),
+            components
+        ) {
+            override fun getDeclarationActionOnDeclarationEnter(
+                declaration: FirDeclaration,
+            ): DiagnosticCollectorDeclarationAction =
+                onDeclarationEnter.invoke(declaration)
 
-        override fun onDeclarationExit(declaration: FirDeclaration) {
-            onDeclarationExit.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/FirIdeDiagnosticVisitor.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.low.level.api.diagnostics
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.collectors.CheckerRunningDiagnosticCollectorVisitor
+import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
+import org.jetbrains.kotlin.fir.resolve.SessionHolder
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistentCheckerContextFactory
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
+
+open class FirIdeDiagnosticVisitor(
+    context: PersistentCheckerContext,
+    components: List<AbstractDiagnosticCollectorComponent>
+) : CheckerRunningDiagnosticCollectorVisitor(context, components) {
+    override fun beforeRunningSingleComponentOnElement(element: FirElement) {
+        checkCanceled()
+    }
+
+    override fun beforeRunningAllComponentsOnElement(element: FirElement) {
+        session.beforeElementDiagnosticCollectionHandler?.beforeCollectingForElement(element)
+    }
+}

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

@@ -0,0 +1,63 @@
+/*
+ * 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.fir
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorVisitor
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.renderWithType
+import org.jetbrains.kotlin.fir.resolve.SessionHolder
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.collectDesignation
+
+private class ContextCollectingDiagnosticCollectorVisitor private constructor(
+    sessionHolder: SessionHolder,
+    private val designation: Iterator<FirDeclaration>,
+) : AbstractDiagnosticCollectorVisitor(
+    PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(sessionHolder),
+    components = emptyList()
+) {
+    private var savedContext: PersistentCheckerContext? = null
+
+    override fun goToNestedDeclarations(element: FirElement) {
+        if (designation.hasNext()) {
+            designation.next().accept(this, null)
+        } else {
+            savedContext = context
+        }
+    }
+
+    override fun runComponents(element: FirElement) {}
+
+    companion object {
+        fun collect(sessionHolder: SessionHolder, firFile: FirFile, designation: List<FirDeclaration>): PersistentCheckerContext {
+            val visitor = ContextCollectingDiagnosticCollectorVisitor(sessionHolder, designation.iterator())
+            firFile.accept(visitor, null)
+            return visitor.savedContext
+                ?: error("Context was not saved")
+        }
+    }
+}
+
+internal object PersistenceContextCollector {
+    fun collectContext(
+        sessionHolder: SessionHolder,
+        firFile: FirFile,
+        declaration: FirDeclaration,
+    ): PersistentCheckerContext {
+        val isLocal = when (declaration) {
+            is FirClassLikeDeclaration<*> -> declaration.symbol.classId.isLocal
+            is FirCallableDeclaration<*> -> declaration.symbol.callableId.isLocal
+            else -> error("Unsupported declaration ${declaration.renderWithType()}")
+        }
+        require(!isLocal) {
+            "Cannot collect context for local declaration ${declaration.renderWithType()}"
+        }
+        val designation = declaration.collectDesignation()
+        return ContextCollectingDiagnosticCollectorVisitor.collect(sessionHolder, firFile, designation)
+    }
+
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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.fir
+
+import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.SessionHolder
+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
+
+internal object PersistentCheckerContextFactory {
+    fun createEmptyPersistenceCheckerContext(sessionHolder: SessionHolder): PersistentCheckerContext {
+        val returnTypeCalculator = createReturnTypeCalculatorForIDE(
+            sessionHolder.session,
+            ScopeSession(),
+            ImplicitBodyResolveComputationSession(),
+            ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+        )
+        return PersistentCheckerContext(sessionHolder, returnTypeCalculator)
+    }
+}

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

@@ -44,7 +44,7 @@ internal class FileStructure(
         val structureElement = structureElements.compute(declaration) { _, structureElement ->
             when {
                 structureElement == null -> createStructureElement(declaration)
-                structureElement is ReanalyzableStructureElement<KtDeclaration> && !structureElement.isUpToDate() -> {
+                structureElement is ReanalyzableStructureElement<KtDeclaration, *> && !structureElement.isUpToDate() -> {
                     structureElement.reanalyze(
                         newKtDeclaration = declaration as KtDeclaration,
                         cache = moduleFileCache,
@@ -90,7 +90,7 @@ internal class FileStructure(
             override fun visitDeclaration(dcl: KtDeclaration) {
                 val structureElement = getStructureElementFor(dcl)
                 structureElements += structureElement
-                if (structureElement !is ReanalyzableStructureElement<*>) {
+                if (structureElement !is ReanalyzableStructureElement<*, *>) {
                     dcl.acceptChildren(this)
                 }
             }

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

@@ -12,10 +12,12 @@ 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.*
 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.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
@@ -31,9 +33,11 @@ internal sealed class FileStructureElement(val firFile: FirFile) {
     abstract val diagnostics: FileStructureElementDiagnostics
 }
 
-internal sealed class ReanalyzableStructureElement<KT : KtDeclaration>(firFile: FirFile) : FileStructureElement(firFile) {
+internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S: AbstractFirBasedSymbol<*>>(
+    firFile: FirFile,
+    protected val firSymbol: S,
+) : FileStructureElement(firFile) {
     abstract override val psi: KtDeclaration
-    abstract val firSymbol: AbstractFirBasedSymbol<*>
     abstract val timestamp: Long
 
     /**
@@ -46,36 +50,11 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration>(firFile:
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         firIdeProvider: FirIdeProvider,
         towerDataContextCollector: FirTowerDataContextCollector,
-    ): ReanalyzableStructureElement<KT>
+    ): ReanalyzableStructureElement<KT, S>
 
     fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
 
-    override val diagnostics = FileStructureElementDiagnostics(firFile, FileStructureElementSingleDeclarationDiagnosticRetriever())
-
-    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_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                        }
-                        inCurrentDeclaration -> DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                        else -> DiagnosticCollectorDeclarationAction.DO_NOT_CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                    }
-                },
-                onDeclarationExit = { firDeclaration ->
-                    if (declaration == firDeclaration) {
-                        inCurrentDeclaration = false
-                    }
-                }
-            )
-        }
-    }
+    override val diagnostics = FileStructureElementDiagnostics(firFile, SingleNonLocalDeclarationDiagnosticRetriever(firSymbol.fir as FirDeclaration))
 
     companion object {
         val recorder = FirElementsRecorder()
@@ -85,9 +64,9 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration>(firFile:
 internal class ReanalyzableFunctionStructureElement(
     firFile: FirFile,
     override val psi: KtNamedFunction,
-    override val firSymbol: FirFunctionSymbol<*>,
+    firSymbol: FirFunctionSymbol<*>,
     override val timestamp: Long
-) : ReanalyzableStructureElement<KtNamedFunction>(firFile) {
+) : ReanalyzableStructureElement<KtNamedFunction, FirFunctionSymbol<*>>(firFile, firSymbol) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
@@ -125,9 +104,9 @@ internal class ReanalyzableFunctionStructureElement(
 internal class ReanalyzablePropertyStructureElement(
     firFile: FirFile,
     override val psi: KtProperty,
-    override val firSymbol: FirPropertySymbol,
+    firSymbol: FirPropertySymbol,
     override val timestamp: Long
-) : ReanalyzableStructureElement<KtProperty>(firFile) {
+) : ReanalyzableStructureElement<KtProperty, FirPropertySymbol>(firFile, firSymbol) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
@@ -164,52 +143,13 @@ internal class ReanalyzablePropertyStructureElement(
 
 internal class NonReanalyzableDeclarationStructureElement(
     firFile: FirFile,
-    private val fir: FirDeclaration,
+    fir: FirDeclaration,
     override val psi: KtDeclaration,
 ) : FileStructureElement(firFile) {
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(fir, recorder)
 
-    override val diagnostics = FileStructureElementDiagnostics(firFile, DiagnosticRetriever())
-
-    private inner class DiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
-        override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
-            var inCurrentDeclaration = false
-            return collector.collectForStructureElement(
-                firFile,
-                onDeclarationEnter = { firDeclaration ->
-                    when {
-                        firDeclaration.isGeneratedDeclaration -> when {
-                            // a generated primary constructor need to be checked. For example it may have a delegated super constructor
-                            // call or some annotations
-                            firDeclaration is FirConstructor && firDeclaration.isPrimary ->
-                                DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                            // Some generated declaration contains structures that we need to check. For example the FIR representation of an
-                            // enum entry initializer, when present, is a generated anonymous object of kind `ENUM_ENTRY`.
-                            else -> DiagnosticCollectorDeclarationAction.DO_NOT_CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                        }
-                        firDeclaration is FirFile -> DiagnosticCollectorDeclarationAction.DO_NOT_CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                        firDeclaration == fir -> {
-                            inCurrentDeclaration = true
-                            DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                        }
-                        FileElementFactory.isReanalyzableContainer(firDeclaration.ktDeclaration) -> {
-                            DiagnosticCollectorDeclarationAction.SKIP
-                        }
-                        inCurrentDeclaration -> {
-                            DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                        }
-                        else -> DiagnosticCollectorDeclarationAction.DO_NOT_CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-                    }
-                },
-                onDeclarationExit = { firDeclaration ->
-                    if (firDeclaration == fir) {
-                        inCurrentDeclaration = false
-                    }
-                },
-            )
-        }
-    }
+    override val diagnostics = FileStructureElementDiagnostics(firFile, SingleNonLocalDeclarationDiagnosticRetriever(fir))
 
 
     companion object {
@@ -243,7 +183,7 @@ internal class RootStructureElement(
 
     private object DiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
         override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
-            return collector.collectForStructureElement(firFile) { firDeclaration ->
+            return collector.collectForStructureElement(firFile, initialContext = null) { firDeclaration ->
                 if (firDeclaration is FirFile) DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_DO_NOT_LOOKUP_FOR_NESTED
                 else DiagnosticCollectorDeclarationAction.SKIP
             }