FIR IDE: fix duplicated diagnostic collection

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

@@ -32,6 +32,9 @@ abstract class AbstractDiagnosticCollectorVisitor(
     protected open fun beforeRunningAllComponentsOnElement(element: FirElement) {}
     protected open fun beforeRunningSingleComponentOnElement(element: FirElement) {}
 
+    protected open fun shouldVisitDeclaration(declaration: FirDeclaration) = true
+    protected open fun onDeclarationExit(declaration: FirDeclaration) {}
+
     protected abstract fun goToNestedDeclarations(element: FirElement)
     protected abstract fun runComponents(element: FirElement)
 
@@ -114,6 +117,12 @@ abstract class AbstractDiagnosticCollectorVisitor(
         }
     }
 
+    override fun visitTypeAlias(typeAlias: FirTypeAlias, data: Nothing?) {
+        withSuppressedDiagnostics(typeAlias) {
+            visitWithDeclaration(typeAlias)
+        }
+    }
+
     override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: Nothing?) {
         if (propertyAccessor !is FirDefaultPropertyAccessor) {
             val property = context.containingDeclarations.last() as FirProperty
@@ -177,13 +186,16 @@ abstract class AbstractDiagnosticCollectorVisitor(
         visitWithGetClassCall(getClassCall)
     }
 
-    private inline fun visitWithDeclaration(
+    protected inline fun visitWithDeclaration(
         declaration: FirDeclaration,
-        block: () -> Unit = { declaration.acceptChildren(this, null) }
+        block: () -> Unit = { goToNestedDeclarations(declaration) }
     ) {
-        runComponents(declaration)
+        if (shouldVisitDeclaration(declaration)) {
-        withDeclaration(declaration) {
+            runComponents(declaration)
-            block()
+            withDeclaration(declaration) {
+                block()
+            }
+            onDeclarationExit(declaration)
         }
     }
 

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

@@ -13,7 +13,6 @@ open class CheckerRunningDiagnosticCollectorVisitor(
     context: PersistentCheckerContext,
     components: List<AbstractDiagnosticCollectorComponent>
 ) : AbstractDiagnosticCollectorVisitor(context, components) {
-    protected val session get() = context.session
 
     override fun goToNestedDeclarations(element: FirElement) {
         element.acceptChildren(this, null)

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

@@ -0,0 +1,77 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+
+abstract class ContextByDesignationCollector<C : Any>(private val designation: FirDeclarationDesignation, var firFile: FirFile) {
+    private var context: C? = null
+    private val designationState = FirDesignationState(designation, firFile)
+
+    protected abstract fun getCurrentContext(): C
+    protected abstract fun goToNestedDeclaration(declaration: FirDeclaration)
+
+    fun getCollectedContext(): C {
+        return context
+            ?: error("Context is not collected yet")
+    }
+
+    fun nextStep() {
+        if (designationState.canGoNext()) {
+            designationState.goNext()
+            if (designationState.currentDeclarationIfPresent == designation.declaration) {
+                check(context == null)
+                context = getCurrentContext()
+            }
+            goToNestedDeclaration(designationState.currentDeclaration)
+        } else {
+            if (designationState.currentDeclarationIfPresent == designation.declaration) {
+                designationState.goToInnerDeclaration()
+            }
+        }
+    }
+}
+
+private class FirDesignationState(val designation: FirDeclarationDesignation, val firFile: FirFile) {
+    /**
+     * Holds current declaration index
+     * if `currentIndex in [0, designation.path.lastIndex]` then current declaration is in path
+     * if `currentIndex == `designation.path.lastIndex + 1` then current declaration is our target declaration
+     * if `currentIndex > designation.path.lastIndex + 1` then we are inside current declaration
+     */
+    private var currentIndex = -1
+
+    fun canGoNext(): Boolean = currentIndex < designation.path.size
+
+    val currentDeclarationIfPresent: FirDeclaration?
+        get() = when (currentIndex) {
+            in designation.path.indices -> designation.path[currentIndex]
+            designation.path.size -> designation.declaration
+            else -> null
+        }
+
+    val currentDeclaration: FirDeclaration
+        get() = currentDeclarationIfPresent
+            ?: error("Went inside target declaration")
+
+    fun goNext() {
+        if (canGoNext()) {
+            currentIndex++
+        } else {
+            throw IndexOutOfBoundsException()
+        }
+    }
+
+    fun goToInnerDeclaration() {
+        if (currentIndex == designation.path.size) {
+            currentIndex++
+        } else {
+            throw IndexOutOfBoundsException()
+        }
+    }
+}

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

@@ -6,9 +6,10 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
 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.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirFile
 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
@@ -22,8 +23,8 @@ internal abstract class FileStructureElementDiagnosticRetriever {
     ): FileStructureElementDiagnosticList
 }
 
-internal class SingleNonLocalDeclarationDiagnosticRetriever (
+internal class SingleNonLocalDeclarationDiagnosticRetriever(
-    private val declaration: FirDeclaration
+    private val structureElementDeclaration: FirDeclaration
 ) : FileStructureElementDiagnosticRetriever() {
     override fun retrieve(
         firFile: FirFile,
@@ -32,10 +33,60 @@ internal class SingleNonLocalDeclarationDiagnosticRetriever (
     ): FileStructureElementDiagnosticList {
         val sessionHolder = SessionHolderImpl.createWithEmptyScopeSession(firFile.session)
         val context = lockProvider.withReadLock(firFile) {
-            PersistenceContextCollector.collectContext(sessionHolder, firFile, declaration)
+            PersistenceContextCollector.collectContext(sessionHolder, firFile, structureElementDeclaration)
         }
-        return collector.collectForStructureElement(declaration) { components ->
+        return collector.collectForStructureElement(structureElementDeclaration) { components ->
-            FirIdeDiagnosticVisitor(context, components)
+            Visitor(structureElementDeclaration, context, components)
+        }
+    }
+
+    private class Visitor(
+        private val structureElementDeclaration: FirDeclaration,
+        context: PersistentCheckerContext,
+        components: List<AbstractDiagnosticCollectorComponent>
+    ) : FirIdeDiagnosticVisitor(context, components) {
+        private var insideAlwaysVisitableDeclarations = 0
+
+        override fun shouldVisitDeclaration(declaration: FirDeclaration): Boolean {
+            if (declaration.shouldVisitWithNestedDeclarations()) {
+                insideAlwaysVisitableDeclarations++
+            }
+
+            if (insideAlwaysVisitableDeclarations > 0) {
+                return true
+            }
+
+            @Suppress("IntroduceWhenSubject")
+            return when {
+                structureElementDeclaration !is FirRegularClass -> true
+                structureElementDeclaration == declaration -> true
+                else -> false
+            }
+        }
+
+        private fun FirDeclaration.shouldVisitWithNestedDeclarations(): Boolean {
+            if (shouldDiagnosticsAlwaysBeCheckedOn(this)) return true
+            return when (this) {
+                is FirAnonymousInitializer -> true
+                is FirEnumEntry -> true
+                is FirValueParameter -> true
+                is FirConstructor -> isPrimary
+                else -> false
+            }
+        }
+
+        override fun onDeclarationExit(declaration: FirDeclaration) {
+            if (declaration.shouldVisitWithNestedDeclarations()) {
+                insideAlwaysVisitableDeclarations--
+            }
+        }
+    }
+
+    companion object {
+        fun shouldDiagnosticsAlwaysBeCheckedOn(firElement: FirElement) = when (firElement.source?.kind) {
+            FirFakeSourceElementKind.PropertyFromParameter -> true
+            FirFakeSourceElementKind.ImplicitConstructor -> true
+            else -> false
         }
     }
 }
@@ -55,16 +106,15 @@ internal object FileDiagnosticRetriever : FileStructureElementDiagnosticRetrieve
         components: List<AbstractDiagnosticCollectorComponent>
     ) : FirIdeDiagnosticVisitor(
         PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(SessionHolderImpl.createWithEmptyScopeSession(firFile.session)),
-        components
+        components,
     ) {
-        override fun goToNestedDeclarations(element: FirElement) {
+        override fun visitFile(file: FirFile, data: Nothing?) {
-            val goNested = when (element) {
+            withSuppressedDiagnostics(file) {
-                is FirFile -> true
+                visitWithDeclaration(file) {
-                is FirDeclaration -> false
+                    file.annotations.forEach { it.accept(this, data) }
-                else -> true
+                    file.imports.forEach { it.accept(this, data) }
-            }
+                    // do not visit declarations here
-            if (goNested) {
+                }
-                super.goToNestedDeclarations(element)
             }
         }
     }

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

@@ -14,20 +14,22 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 
 internal open class FirIdeDiagnosticVisitor(
     context: PersistentCheckerContext,
-    components: List<AbstractDiagnosticCollectorComponent>
+    components: List<AbstractDiagnosticCollectorComponent>,
 ) : CheckerRunningDiagnosticCollectorVisitor(context, components) {
+    private val beforeElementDiagnosticCollectionHandler = context.session.beforeElementDiagnosticCollectionHandler
+
     override fun beforeRunningSingleComponentOnElement(element: FirElement) {
         checkCanceled()
     }
 
-    override fun goToNestedElements(element: FirElement) {
+    override fun goToNestedDeclarations(element: FirElement) {
         if (element is FirDeclaration) {
-            session.beforeElementDiagnosticCollectionHandler?.beforeGoingNestedDeclaration(element, context)
+            beforeElementDiagnosticCollectionHandler?.beforeGoingNestedDeclaration(element, context)
         }
-        super.goToNestedElements(element)
+        super.goToNestedDeclarations(element)
     }
 
     override fun beforeRunningAllComponentsOnElement(element: FirElement) {
-        session.beforeElementDiagnosticCollectionHandler?.beforeCollectingForElement(element)
+        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

@@ -11,33 +11,41 @@ import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorV
 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
+import org.jetbrains.kotlin.idea.fir.low.level.api.ContextByDesignationCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
 
 private class ContextCollectingDiagnosticCollectorVisitor private constructor(
     sessionHolder: SessionHolder,
-    private val designation: Iterator<FirDeclaration>,
+    designation: FirDeclarationDesignation,
+    firFile: FirFile,
 ) : AbstractDiagnosticCollectorVisitor(
     PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(sessionHolder),
     components = emptyList()
 ) {
-    private var savedContext: PersistentCheckerContext? = null
+    private val contextCollector = object : ContextByDesignationCollector<PersistentCheckerContext>(designation, firFile) {
+        override fun getCurrentContext(): PersistentCheckerContext = context
+
+        override fun goToNestedDeclaration(declaration: FirDeclaration) {
+            declaration.accept(this@ContextCollectingDiagnosticCollectorVisitor, null)
+        }
+    }
 
     override fun goToNestedDeclarations(element: FirElement) {
-        if (designation.hasNext()) {
+        if (element is FirDeclaration) {
-            designation.next().accept(this, null)
+            contextCollector.nextStep()
         } else {
-            savedContext = context
+            element.accept(this, null)
         }
     }
 
     override fun runComponents(element: FirElement) {}
 
     companion object {
-        fun collect(sessionHolder: SessionHolder, firFile: FirFile, designation: List<FirDeclaration>): PersistentCheckerContext {
+        fun collect(sessionHolder: SessionHolder, firFile: FirFile, designation: FirDeclarationDesignation): PersistentCheckerContext {
-            val visitor = ContextCollectingDiagnosticCollectorVisitor(sessionHolder, designation.iterator())
+            val visitor = ContextCollectingDiagnosticCollectorVisitor(sessionHolder, designation, firFile)
             firFile.accept(visitor, null)
-            return visitor.savedContext
+            return visitor.contextCollector.getCollectedContext()
-                ?: error("Context was not saved")
         }
     }
 }
@@ -57,7 +65,9 @@ internal object PersistenceContextCollector {
             "Cannot collect context for local declaration ${declaration.renderWithType()}"
         }
         val designation = declaration.collectDesignation()
+        check(!designation.isLocalDesignation) {
+            "Designation should not local for ${declaration.renderWithType()}"
+        }
         return ContextCollectingDiagnosticCollectorVisitor.collect(sessionHolder, firFile, designation)
     }
-
 }