FIR IDE: fix collecting diagnostics for raanalysable non-toplevel declarations

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

@@ -5,6 +5,7 @@
 
 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.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
@@ -17,10 +18,8 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarati
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.psi.KtAnnotated
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
 import java.util.concurrent.ConcurrentHashMap
 
 
@@ -58,16 +57,31 @@ internal class FileStructure(
 
     @OptIn(ExperimentalStdlibApi::class)
     fun getAllDiagnosticsForFile(): Collection<FirPsiDiagnostic<*>> {
-        val containersForStructureElement = buildList {
+        val structureElements = getAllStructureElements()
-            add(ktFile)
-            addAll(ktFile.declarations)
-        }
-        val structureElements = containersForStructureElement.map(::getStructureElementFor)
         return buildSet {
             structureElements.forEach { it.diagnostics.forEach { diagnostics -> addAll(diagnostics) } }
         }
     }
 
+    private fun getAllStructureElements(): Collection<FileStructureElement> {
+        val structureElements = mutableSetOf(getStructureElementFor(ktFile))
+        ktFile.accept(object : KtVisitorVoid() {
+            override fun visitElement(element: PsiElement) {
+                element.acceptChildren(this)
+            }
+
+            override fun visitDeclaration(dcl: KtDeclaration) {
+                val structureElement = getStructureElementFor(dcl)
+                structureElements += structureElement
+                if (structureElement !is ReanalyzableStructureElement<*>) {
+                    dcl.acceptChildren(this)
+                }
+            }
+        })
+
+        return structureElements
+    }
+
 
     private fun createDeclarationStructure(declaration: KtDeclaration): FileStructureElement {
         val firDeclaration = declaration.findSourceNonLocalFirDeclaration(

idea/testData/quickfix/modifiers/removeIncompatibleModifier.kt

@@ -2,3 +2,4 @@
 class Foo {
     public<caret> private fun bar() { }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeIncompatibleModifier.kt.after

@@ -2,3 +2,4 @@
 class Foo {
     private fun bar() { }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeRedundantModifier1.kt

@@ -6,3 +6,5 @@ abstract class B() {
 abstract class A() : B() {
     <caret>open abstract override fun foo()
 }
+
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeRedundantModifier1.kt.after

@@ -6,3 +6,5 @@ abstract class B() {
 abstract class A() : B() {
     <caret>abstract override fun foo()
 }
+
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeRedundantModifier2.kt

@@ -6,3 +6,5 @@ abstract class B() {
 abstract class A() : B() {
     abstract <caret>open override fun foo()
 }
+
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeRedundantModifier2.kt.after

@@ -6,3 +6,5 @@ abstract class B() {
 abstract class A() : B() {
     abstract <caret>override fun foo()
 }
+
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeRedundantModifier3.kt

@@ -6,3 +6,5 @@ abstract class B() {
 abstract class A() : B() {
     abstract override <caret>open fun foo()
 }
+
+/* FIR_COMPARISON */

idea/testData/quickfix/modifiers/removeRedundantModifier3.kt.after

@@ -6,3 +6,5 @@ abstract class B() {
 abstract class A() : B() {
     abstract override<caret> fun foo()
 }
+
+/* FIR_COMPARISON */