[LL FIR] ContextCollector: optimize collector for file elements

We shouldn't trigger declaration processing if we want to check just
annotations/package/imports.
But currently, we still will iterate over the entire file
if there is an element somewhere inside package/import
(see `fileImportInside.kt` test case)

^KT-65344
^KT-65560 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt

@@ -335,18 +335,34 @@ private class ContextCollectorVisitor(
         else -> true
     }
 
-    override fun visitScript(script: FirScript) {
+    override fun visitScript(script: FirScript) = withProcessor(script) {
-        context.withScript(script, holder) {
+        dumpContext(script, ContextKind.SELF)
-            withInterceptor {
+
-                super.visitScript(script)
+        processSignatureAnnotations(script)
+
+        onActiveBody {
+            context.withScript(script, holder) {
+                dumpContext(script, ContextKind.BODY)
+
+                onActive {
+                    withInterceptor {
+                        processChildren(script)
+                    }
+                }
             }
         }
     }
 
-    override fun visitFile(file: FirFile) {
+    override fun visitFile(file: FirFile) = withProcessor(file) {
         context.withFile(file, holder) {
-            withInterceptor {
+            dumpContext(file, ContextKind.SELF)
-                super.visitFile(file)
+
+            processFileHeader(file)
+
+            onActive {
+                withInterceptor {
+                    processChildren(file)
+                }
             }
         }
     }
@@ -417,6 +433,13 @@ private class ContextCollectorVisitor(
         }
     }
 
+    @OptIn(PrivateForInline::class)
+    private fun Processor.processFileHeader(file: FirFile) {
+        process(file.packageDirective)
+        processList(file.imports)
+        process(file.annotationsContainer)
+    }
+
     /**
      * Same as [processClassHeader], but for anonymous objects.
      *

analysis/low-level-api-fir/testData/contextCollector/fileAnnotation.copy.txt

@@ -1,23 +0,0 @@
-Tower Data Context:
-    Element 0
-        Scope: FirDefaultStarImportingScope
-    Element 1
-        Scope: FirExplicitStarImportingScope
-    Element 2
-        Scope: FirDefaultSimpleImportingScope
-    Element 3
-        Scope: FirDefaultSimpleImportingScope
-    Element 4
-        Scope: FirPackageMemberScope
-    Element 5
-        Scope: FirExplicitSimpleImportingScope
-
-FILE: [ResolvedTo(IMPORTS)] fileAnnotation.kt
-    @FILE:R|Anno|[Types]()
-    [ResolvedTo(STATUS)] annotations container
-    @R|kotlin/annotation/Target|[Types](FILE#) public final [ResolvedTo(STATUS)] annotation class Anno : R|kotlin/Annotation| {
-        public? [ResolvedTo(RAW_FIR)] constructor(): R|Anno| {
-            LAZY_super<R|kotlin/Any|>
-        }
-
-    }

analysis/low-level-api-fir/testData/contextCollector/fileAnnotation.txt

@@ -13,10 +13,10 @@ Tower Data Context:
         Scope: FirExplicitSimpleImportingScope
 
 FILE: [ResolvedTo(IMPORTS)] fileAnnotation.kt
-    @FILE:R|Anno|[Types]()
+    @FILE:Anno[Unresolved]()
-    [ResolvedTo(STATUS)] annotations container
+    [ResolvedTo(RAW_FIR)] annotations container
-    @R|kotlin/annotation/Target|[Types](FILE#) public final [ResolvedTo(STATUS)] annotation class Anno : R|kotlin/Annotation| {
+    @Target[Unresolved](LAZY_EXPRESSION) public final? [ResolvedTo(RAW_FIR)] annotation class Anno : R|kotlin/Annotation| {
-        public [ResolvedTo(STATUS)] constructor(): R|Anno| {
+        public? [ResolvedTo(RAW_FIR)] constructor(): R|Anno| {
             LAZY_super<R|kotlin/Any|>
         }
 

analysis/low-level-api-fir/testData/contextCollector/fileImport.copy.txt

@@ -1,24 +0,0 @@
-Tower Data Context:
-    Element 0
-        Scope: FirDefaultStarImportingScope
-    Element 1
-        Scope: FirExplicitStarImportingScope
-    Element 2
-        Scope: FirDefaultSimpleImportingScope
-    Element 3
-        Scope: FirDefaultSimpleImportingScope
-    Element 4
-        Scope: FirPackageMemberScope
-    Element 5
-        Scope: FirExplicitSimpleImportingScope
-            Import name:Anno
-                Qualified name: foo.Anno
-                Is all under: false
-
-FILE: [ResolvedTo(IMPORTS)] fileImport.kt
-    @R|kotlin/annotation/Target|[Types](FILE#) public final [ResolvedTo(STATUS)] annotation class Anno : R|kotlin/Annotation| {
-        public? [ResolvedTo(RAW_FIR)] constructor(): R|foo/Anno| {
-            LAZY_super<R|kotlin/Any|>
-        }
-
-    }

analysis/low-level-api-fir/testData/contextCollector/fileImport.txt

@@ -16,8 +16,8 @@ Tower Data Context:
                 Is all under: false
 
 FILE: [ResolvedTo(IMPORTS)] fileImport.kt
-    @R|kotlin/annotation/Target|[Types](FILE#) public final [ResolvedTo(STATUS)] annotation class Anno : R|kotlin/Annotation| {
+    @Target[Unresolved](LAZY_EXPRESSION) public final? [ResolvedTo(RAW_FIR)] annotation class Anno : R|kotlin/Annotation| {
-        public [ResolvedTo(STATUS)] constructor(): R|foo/Anno| {
+        public? [ResolvedTo(RAW_FIR)] constructor(): R|foo/Anno| {
             LAZY_super<R|kotlin/Any|>
         }
 

analysis/low-level-api-fir/testData/contextCollector/filePackage.copy.txt

@@ -1,21 +0,0 @@
-Tower Data Context:
-    Element 0
-        Scope: FirDefaultStarImportingScope
-    Element 1
-        Scope: FirExplicitStarImportingScope
-    Element 2
-        Scope: FirDefaultSimpleImportingScope
-    Element 3
-        Scope: FirDefaultSimpleImportingScope
-    Element 4
-        Scope: FirPackageMemberScope
-    Element 5
-        Scope: FirExplicitSimpleImportingScope
-
-FILE: [ResolvedTo(IMPORTS)] filePackage.kt
-    @R|kotlin/annotation/Target|[Types](FILE#) public final [ResolvedTo(STATUS)] annotation class Anno : R|kotlin/Annotation| {
-        public? [ResolvedTo(RAW_FIR)] constructor(): R|foo/Anno| {
-            LAZY_super<R|kotlin/Any|>
-        }
-
-    }

analysis/low-level-api-fir/testData/contextCollector/filePackage.txt

@@ -13,9 +13,9 @@ Tower Data Context:
         Scope: FirExplicitSimpleImportingScope
 
 FILE: [ResolvedTo(IMPORTS)] filePackage.kt
-    @R|kotlin/annotation/Target|[Types](FILE#) public final [ResolvedTo(STATUS)] annotation class Anno : R|kotlin/Annotation| {
+    @Target[Unresolved](LAZY_EXPRESSION) public final? [ResolvedTo(RAW_FIR)] annotation class Anno : R|kotlin/Annotation| {
-        public [ResolvedTo(STATUS)] constructor(): R|foo/Anno| {
+        public? [ResolvedTo(RAW_FIR)] constructor(): R|foo/Anno| {
             LAZY_super<R|kotlin/Any|>
         }
 
-    }
+    }