FIR IDE: get rid of DiagnosticCollectorDeclarationAction

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

@@ -67,12 +67,6 @@ abstract class AbstractDiagnosticCollector(
     }
 }
 
-enum class DiagnosticCollectorDeclarationAction(val checkInCurrentDeclaration: Boolean, val lookupForNestedDeclaration: Boolean) {
-    CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED(checkInCurrentDeclaration = true, lookupForNestedDeclaration = true),
-    CHECK_IN_CURRENT_DECLARATION_AND_DO_NOT_LOOKUP_FOR_NESTED(checkInCurrentDeclaration = true, lookupForNestedDeclaration = false),
-    DO_NOT_CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED(checkInCurrentDeclaration = false, lookupForNestedDeclaration = true),
-    SKIP(checkInCurrentDeclaration = false, lookupForNestedDeclaration = false),
-}
 
 fun AbstractDiagnosticCollector.registerAllComponents() {
     initializeComponents(

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

@@ -28,7 +28,6 @@ abstract class AbstractDiagnosticCollectorVisitor(
     @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) {}
@@ -182,16 +181,10 @@ abstract class AbstractDiagnosticCollectorVisitor(
         declaration: FirDeclaration,
         block: () -> Unit = { declaration.acceptChildren(this, null) }
     ) {
-        if (!currentAction.lookupForNestedDeclaration) return
+        runComponents(declaration)
-
+        withDeclaration(declaration) {
-        val action = getDeclarationActionOnDeclarationEnter(declaration)
+            block()
-        withDiagnosticsAction(action) {
-            runComponents(declaration)
-            withDeclaration(declaration) {
-                block()
-            }
         }
-        onDeclarationExit(declaration)
     }
 
     private fun visitWithDeclarationAndReceiver(declaration: FirDeclaration, labelName: Name?, receiverTypeRef: FirTypeRef?) {
@@ -220,21 +213,6 @@ abstract class AbstractDiagnosticCollectorVisitor(
         }
     }
 
-    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)
     inline fun <R> withQualifiedAccess(qualifiedAccess: FirQualifiedAccess, block: () -> R): R {
         val existingContext = context

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

@@ -20,12 +20,10 @@ open class CheckerRunningDiagnosticCollectorVisitor(
     }
 
     override fun runComponents(element: FirElement) {
-        if (currentAction.checkInCurrentDeclaration) {
+        beforeRunningAllComponentsOnElement(element)
-            beforeRunningAllComponentsOnElement(element)
+        components.forEach {
-            components.forEach {
+            beforeRunningSingleComponentOnElement(element)
-                beforeRunningSingleComponentOnElement(element)
+            element.accept(it, context)
-                element.accept(it, context)
-            }
         }
     }
 }

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

@@ -36,7 +36,11 @@ interface SessionHolder {
     val scopeSession: ScopeSession
 }
 
-data class SessionHolderImpl(override val session: FirSession, override val scopeSession: ScopeSession): SessionHolder
+data class SessionHolderImpl(override val session: FirSession, override val scopeSession: ScopeSession) : SessionHolder {
+    companion object {
+        fun createWithEmptyScopeSession(session: FirSession): SessionHolderImpl = SessionHolderImpl(session, ScopeSession())
+    }
+}
 
 abstract class BodyResolveComponents : SessionHolder {
     abstract val returnTypeCalculator: ReturnTypeCalculator

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

@@ -5,11 +5,13 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
 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
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistentCheckerContextFactory
 
 internal abstract class FileStructureElementDiagnosticRetriever {
     abstract fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList
@@ -19,8 +21,36 @@ internal class SingleNonLocalDeclarationDiagnosticRetriever(
     private val declaration: FirDeclaration
 ) : FileStructureElementDiagnosticRetriever() {
     override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
-        val sessionHolder = SessionHolderImpl(firFile.session, ScopeSession())
+        val sessionHolder = SessionHolderImpl.createWithEmptyScopeSession(firFile.session)
         val context = PersistenceContextCollector.collectContext(sessionHolder, firFile, declaration)
-        return collector.collectForStructureElement(declaration, context)
+        return collector.collectForStructureElement(declaration) { components ->
+            FirIdeDiagnosticVisitor(context, components)
+        }
+    }
+}
+
+internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
+    override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList =
+        collector.collectForStructureElement(firFile) { components ->
+            Visitor(firFile, components)
+        }
+
+    private class Visitor(
+        firFile: FirFile,
+        components: List<AbstractDiagnosticCollectorComponent>
+    ) : FirIdeDiagnosticVisitor(
+        PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(SessionHolderImpl.createWithEmptyScopeSession(firFile.session)),
+        components
+    ) {
+        override fun goToNestedDeclarations(element: FirElement) {
+            val goNested = when (element) {
+                is FirFile -> true
+                is FirDeclaration -> false
+                else -> true
+            }
+            if (goNested) {
+                super.goToNestedDeclarations(element)
+            }
+        }
     }
 }

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

@@ -5,15 +5,10 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
-import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.analysis.collectors.CheckerRunningDiagnosticCollectorVisitor
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
-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.FirIdeStructureElementDiagnosticsCollector
-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) {
     companion object {
@@ -23,51 +18,14 @@ internal class FileStructureElementDiagnosticsCollector private constructor(priv
 
     fun collectForStructureElement(
         firDeclaration: FirDeclaration,
-        initialContext: PersistentCheckerContext?,
+        createVisitor: (components: List<AbstractDiagnosticCollectorComponent>) -> CheckerRunningDiagnosticCollectorVisitor,
-        onDeclarationExit: (FirDeclaration) -> Unit = {},
-        onDeclarationEnter: (FirDeclaration) -> DiagnosticCollectorDeclarationAction = {
-            DiagnosticCollectorDeclarationAction.CHECK_IN_CURRENT_DECLARATION_AND_LOOKUP_FOR_NESTED
-        },
     ): FileStructureElementDiagnosticList =
         FirIdeStructureElementDiagnosticsCollector(
             firDeclaration.session,
-            initialContext,
+            createVisitor,
             useExtendedCheckers,
-            onDeclarationEnter,
-            onDeclarationExit
         ).let { collector ->
             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
-    ) : AbstractFirIdeDiagnosticsCollector(
-        session,
-        useExtendedCheckers,
-    ) {
-        val result = mutableMapOf<PsiElement, MutableList<FirPsiDiagnostic<*>>>()
-
-        override fun onDiagnostic(diagnostic: FirPsiDiagnostic<*>) {
-            result.addValueFor(diagnostic.psiElement, diagnostic)
-        }
-
-        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)
-            }
-        }
-    }
 }

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

@@ -13,7 +13,7 @@ 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(
+internal open class FirIdeDiagnosticVisitor(
     context: PersistentCheckerContext,
     components: List<AbstractDiagnosticCollectorComponent>
 ) : CheckerRunningDiagnosticCollectorVisitor(context, components) {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/fir/FirIdeStructureElementDiagnosticsCollector.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.low.level.api.diagnostics.fir
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.collectors.CheckerRunningDiagnosticCollectorVisitor
+import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.AbstractFirIdeDiagnosticsCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
+
+internal class FirIdeStructureElementDiagnosticsCollector(
+    session: FirSession,
+    createVisitor: (components: List<AbstractDiagnosticCollectorComponent>) -> CheckerRunningDiagnosticCollectorVisitor,
+    useExtendedCheckers: Boolean,
+) : AbstractFirIdeDiagnosticsCollector(
+    session,
+    useExtendedCheckers,
+) {
+    override val visitor: CheckerRunningDiagnosticCollectorVisitor = createVisitor(components)
+
+    val result = mutableMapOf<PsiElement, MutableList<FirPsiDiagnostic<*>>>()
+
+    override fun onDiagnostic(diagnostic: FirPsiDiagnostic<*>) {
+        result.addValueFor(diagnostic.psiElement, diagnostic)
+    }
+}

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

@@ -6,25 +6,19 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
 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.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileDiagnosticRetriever
-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
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.hasFqName
-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 sealed class FileStructureElement(val firFile: FirFile) {
@@ -179,16 +173,7 @@ internal class RootStructureElement(
     override val mappings: Map<KtElement, FirElement> =
         FirElementsRecorder.recordElementsFrom(firFile, recorder)
 
-    override val diagnostics = FileStructureElementDiagnostics(firFile, DiagnosticRetriever)
+    override val diagnostics = FileStructureElementDiagnostics(firFile, FileDiagnosticRetriever)
-
-    private object DiagnosticRetriever : FileStructureElementDiagnosticRetriever() {
-        override fun retrieve(firFile: FirFile, collector: FileStructureElementDiagnosticsCollector): FileStructureElementDiagnosticList {
-            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
-            }
-        }
-    }
 
     companion object {
         private val recorder = object : FirElementsRecorder() {