From 38a3010ac84166758f303ef931c4cd98a0744b23 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Tue, 13 Feb 2024 15:30:28 +0100
Subject: [PATCH] [LL FIR] FirElementBuilder: search file annotations without
 body resolution

^KT-65780
---
 .../fir/element/builder/FirElementBuilder.kt  | 22 ++++++++++++-------
 .../annotations/fileAnnotation.txt            |  4 ++--
 .../annotations/fileAnnotationScript.txt      |  4 ++--
 .../annotations/useSite/file.txt              |  4 ++--
 .../annotations/useSite/fileScript.txt        |  4 ++--
 5 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
index 0b75451ed53..4bbc18c7291 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
@@ -116,16 +116,21 @@ internal class FirElementBuilder(
         resolveAndFindFirForAnchor: (FirAnnotationContainer, T) -> FirElement?,
     ): FirElement? {
         val anchorElement = anchorElementProvider(element) ?: return null
-        val declaration = annotatedElementProvider(anchorElement) ?: return null
-        val nonLocalDeclaration = declaration.getNonLocalContainingOrThisDeclaration()
-        if (declaration != nonLocalDeclaration) return null
+        val annotationElement = annotatedElementProvider(anchorElement) ?: return null
 
-        val firDeclaration = nonLocalDeclaration.findSourceNonLocalFirDeclaration(
-            firFileBuilder = moduleComponents.firFileBuilder,
-            provider = moduleComponents.session.firProvider,
-        )
+        val firAnnotationContainer = if (annotationElement is KtFile) {
+            moduleComponents.firFileBuilder.buildRawFirFileWithCaching(annotationElement)
+        } else {
+            val nonLocalDeclaration = annotationElement.getNonLocalContainingOrThisDeclaration()
+            if (annotationElement != nonLocalDeclaration) return null
 
-        val anchorFir = resolveAndFindFirForAnchor(firDeclaration, anchorElement) ?: return null
+            nonLocalDeclaration.findSourceNonLocalFirDeclaration(
+                firFileBuilder = moduleComponents.firFileBuilder,
+                provider = moduleComponents.session.firProvider,
+            )
+        }
+
+        val anchorFir = resolveAndFindFirForAnchor(firAnnotationContainer, anchorElement) ?: return null
         // We use identity comparison here intentionally to check that it is exactly the object we want to find
         if (element === anchorElement) return anchorFir
 
@@ -136,6 +141,7 @@ internal class FirElementBuilder(
         val modifierList = when (val parent = parent) {
             is KtModifierList -> parent
             is KtAnnotation -> parent.parent as? KtModifierList
+            is KtFileAnnotationList -> return parent.parent as? KtFile
             else -> null
         }
 
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotation.txt b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotation.txt
index 387054528f4..3539865b5dc 100644
--- a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotation.txt
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotation.txt
@@ -6,6 +6,6 @@ FIR element rendered:
 @FILE:R|kotlin/Suppress|[Types](names = vararg(String()))
 
 FIR FILE:
-FILE: [ResolvedTo(BODY_RESOLVE)] fileAnnotation.kt
+FILE: [ResolvedTo(ANNOTATION_ARGUMENTS)] fileAnnotation.kt
     @FILE:R|kotlin/Suppress|[Types](names = vararg(String()))
-    [ResolvedTo(BODY_RESOLVE)] annotations container
+    [ResolvedTo(ANNOTATION_ARGUMENTS)] annotations container
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotationScript.txt b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotationScript.txt
index 354cc6f2d7c..ca635e1c183 100644
--- a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotationScript.txt
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/fileAnnotationScript.txt
@@ -6,9 +6,9 @@ FIR element rendered:
 @FILE:R|kotlin/Suppress|[Types](names = vararg(String()))
 
 FIR FILE:
-FILE: [ResolvedTo(BODY_RESOLVE)] fileAnnotationScript.kts
+FILE: [ResolvedTo(ANNOTATION_ARGUMENTS)] fileAnnotationScript.kts
     @FILE:R|kotlin/Suppress|[Types](names = vararg(String()))
-    [ResolvedTo(BODY_RESOLVE)] annotations container
+    [ResolvedTo(ANNOTATION_ARGUMENTS)] annotations container
     context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
     SCRIPT: [ResolvedTo(RAW_FIR)] <script-fileAnnotationScript.kts>
         [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/file.txt b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/file.txt
index 4dcb3a49c88..5d898e7ad2e 100644
--- a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/file.txt
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/file.txt
@@ -6,9 +6,9 @@ FIR element rendered:
 @FILE:R|Ann|[Types]()
 
 FIR FILE:
-FILE: [ResolvedTo(BODY_RESOLVE)] file.kt
+FILE: [ResolvedTo(ANNOTATION_ARGUMENTS)] file.kt
     @FILE:R|Ann|[Types]()
-    [ResolvedTo(BODY_RESOLVE)] annotations container
+    [ResolvedTo(ANNOTATION_ARGUMENTS)] annotations container
     public final [ResolvedTo(STATUS)] annotation class Ann : R|kotlin/Annotation| {
         public [ResolvedTo(STATUS)] [ContainingClassKey=Ann] constructor(): R|Ann| {
             LAZY_super<R|kotlin/Any|>
diff --git a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/fileScript.txt b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/fileScript.txt
index cbd48ca1b7d..bd59cfc7c00 100644
--- a/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/fileScript.txt
+++ b/analysis/low-level-api-fir/testData/getOrBuildFir/annotations/useSite/fileScript.txt
@@ -6,9 +6,9 @@ FIR element rendered:
 @FILE:R|Ann|[Types]()
 
 FIR FILE:
-FILE: [ResolvedTo(BODY_RESOLVE)] fileScript.kts
+FILE: [ResolvedTo(ANNOTATION_ARGUMENTS)] fileScript.kts
     @FILE:R|Ann|[Types]()
-    [ResolvedTo(BODY_RESOLVE)] annotations container
+    [ResolvedTo(ANNOTATION_ARGUMENTS)] annotations container
     context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
     SCRIPT: [ResolvedTo(RAW_FIR)] <script-fileScript.kts>
         [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|