[FIR IDE] OnAir resolve for annotations with file site

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

@@ -7,21 +7,21 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirNonLocalDeclarationBuilder
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.*
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirNonLocalDeclarationBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.psi.*
 
 internal object DeclarationCopyBuilder {
-    fun <T : KtElement> createDeclarationCopy(
+    fun createDeclarationCopy(
         state: FirModuleResolveState,
         nonLocalDeclaration: KtDeclaration,
-        replacement: RawFirReplacement<T>
+        replacement: RawFirReplacement
     ): FirDeclaration {
 
         return when (nonLocalDeclaration) {
@@ -49,10 +49,10 @@ internal object DeclarationCopyBuilder {
         }
     }
 
-    private fun <T : KtElement> createFunctionCopy(
+    private fun createFunctionCopy(
         rootNonLocalDeclaration: KtNamedFunction,
         state: FirModuleResolveState,
-        replacement: RawFirReplacement<T>,
+        replacement: RawFirReplacement,
     ): FirSimpleFunction {
 
         val originalFunction = rootNonLocalDeclaration.getOrBuildFirOfType<FirSimpleFunction>(state)
@@ -73,10 +73,10 @@ internal object DeclarationCopyBuilder {
         }.apply { reassignAllReturnTargets(builtFunction) }
     }
 
-    private fun <T : KtElement> createClassCopy(
+    private fun createClassCopy(
         rootNonLocalDeclaration: KtClassOrObject,
         state: FirModuleResolveState,
-        replacement: RawFirReplacement<T>,
+        replacement: RawFirReplacement,
     ): FirRegularClass {
         val originalFirClass = rootNonLocalDeclaration.getOrBuildFirOfType<FirRegularClass>(state)
         val builtClass = createCopy(rootNonLocalDeclaration, originalFirClass, replacement)
@@ -95,19 +95,19 @@ internal object DeclarationCopyBuilder {
         }
     }
 
-    private fun <T : KtElement> createTypeAliasCopy(
+    private fun createTypeAliasCopy(
         rootNonLocalDeclaration: KtTypeAlias,
         state: FirModuleResolveState,
-        replacement: RawFirReplacement<T>,
+        replacement: RawFirReplacement,
     ): FirTypeAlias {
         val originalFirTypeAlias = rootNonLocalDeclaration.getOrBuildFirOfType<FirTypeAlias>(state)
         return createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
     }
 
-    private fun <T : KtElement> createPropertyCopy(
+    private fun createPropertyCopy(
         rootNonLocalDeclaration: KtProperty,
         state: FirModuleResolveState,
-        replacement: RawFirReplacement<T>,
+        replacement: RawFirReplacement,
     ): FirProperty {
         val originalProperty = rootNonLocalDeclaration.getOrBuildFirOfType<FirProperty>(state)
         val builtProperty = createCopy(rootNonLocalDeclaration, originalProperty, replacement)
@@ -158,16 +158,16 @@ internal object DeclarationCopyBuilder {
         moduleData = state.rootModuleSession.moduleData
     }
 
-    internal inline fun <reified D : FirDeclaration, T : KtElement> createCopy(
+    internal inline fun <reified D : FirDeclaration> createCopy(
         rootNonLocalDeclaration: KtDeclaration,
         originalFirDeclaration: D,
-        replacement: RawFirReplacement<T>? = null,
+        replacement: RawFirReplacement? = null,
     ): D {
-        return RawFirNonLocalDeclarationBuilder.buildWithReplacement(
+        return RawFirNonLocalDeclarationBuilder.build(
             session = originalFirDeclaration.moduleData.session,
             baseScopeProvider = originalFirDeclaration.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = originalFirDeclaration.collectDesignation(),
-            declarationToBuild = rootNonLocalDeclaration,
+            rootNonLocalDeclaration = rootNonLocalDeclaration,
             replacement = replacement,
         ) as D
     }

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

@@ -6,23 +6,29 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.parentOfType
 import com.intellij.psi.util.parentsOfType
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirNonLocalDeclarationBuilder
+import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.realPsi
-import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.asTowerDataElement
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateDepended
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextAllElementsCollector
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.buildFileFirAnnotation
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
@@ -48,9 +54,9 @@ object LowLevelFirApiFacadeForResolveOnAir {
     private fun recordOriginalDeclaration(targetDeclaration: KtNamedDeclaration, originalDeclaration: KtNamedDeclaration) {
         require(!targetDeclaration.isPhysical)
         require(originalDeclaration.containingKtFile !== targetDeclaration.containingKtFile)
-        val originalDeclrationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
+        val originalDeclarationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
         val fakeDeclarationParents = targetDeclaration.parentsOfType<KtDeclaration>().toList()
-        originalDeclrationParents.zip(fakeDeclarationParents) { original, fake ->
+        originalDeclarationParents.zip(fakeDeclarationParents) { original, fake ->
             fake.originalDeclaration = original
         }
     }
@@ -95,7 +101,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             FileTowerProvider(place, onAirGetTowerContextForFile(state, place))
         } else {
             val validPlace = PsiTreeUtil.findFirstParent(place, false) {
-                RawFirNonLocalDeclarationBuilder.elementIsApplicable(it as KtElement)
+                RawFirReplacement.isApplicableForReplacement(it as KtElement)
             } as KtElement
 
             FirTowerDataContextAllElementsCollector().also {
@@ -165,15 +171,52 @@ object LowLevelFirApiFacadeForResolveOnAir {
         return FirModuleResolveStateDepended(originalState, collector, recordedMap)
     }
 
-    private fun <T : KtElement> runBodyResolveOnAir(
+    private fun tryResolveAsFileAnnotation(
+        annotationEntry: KtAnnotationEntry,
         state: FirModuleResolveStateImpl,
-        replacement: RawFirReplacement<T>,
+        replacement: RawFirReplacement,
+        firFile: FirFile,
+    ): FirAnnotationCall {
+        val annotationCall = buildFileFirAnnotation(
+            firFile.declarationSiteSession,
+            firFile.declarationSiteSession.firIdeProvider.kotlinScopeProvider,
+            annotationEntry,
+            replacement
+        )
+        state.firLazyDeclarationResolver.resolveFileAnnotations(
+            firFile,
+            listOf(annotationCall),
+            state.rootModuleSession.cache,
+            ScopeSession()
+        )
+
+        return annotationCall
+    }
+
+    private fun runBodyResolveOnAir(
+        state: FirModuleResolveStateImpl,
+        replacement: RawFirReplacement,
         isOnAirResolve: Boolean,
         collector: FirTowerDataContextCollector? = null,
-    ): FirDeclaration {
+    ): FirElement {
 
         val nonLocalDeclaration = findEnclosingNonLocalDeclaration(replacement.from)
-            ?: error("Cannot find enclosing declaration for ${replacement.from.getElementTextInContext()}")
+        val originalFirFile = state.getOrBuildFirFile(replacement.from.containingKtFile)
+
+        if (nonLocalDeclaration == null) {
+            //It is possible that it is file annotation is going to resolve
+            val annotationCall = replacement.from.parentOfType<KtAnnotationEntry>(withSelf = true)
+            if (annotationCall != null) {
+                return tryResolveAsFileAnnotation(
+                    annotationEntry = annotationCall,
+                    state = state,
+                    replacement = replacement,
+                    firFile = originalFirFile
+                )
+            } else {
+                error("Cannot find enclosing declaration for ${replacement.from.getElementTextInContext()}")
+            }
+        }
 
         val copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
             state = state,
@@ -181,8 +224,6 @@ object LowLevelFirApiFacadeForResolveOnAir {
             replacement = replacement,
         )
 
-        val originalFirFile = state.getOrBuildFirFile(nonLocalDeclaration.containingKtFile)
-
         state.firLazyDeclarationResolver.lazyDesignatedResolveDeclaration(
             firDeclarationToResolve = copiedFirDeclaration,
             moduleFileCache = state.rootModuleSession.cache,

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

@@ -125,7 +125,7 @@ internal class FileStructure(
                 scopeSession,
                 checkPCE = true
             )
-            firLazyDeclarationResolver.resolveFileAnnotations(firFile, moduleFileCache, scopeSession)
+            firLazyDeclarationResolver.resolveFileAnnotations(firFile, firFile.annotations, moduleFileCache, scopeSession)
             RootStructureElement(
                 firFile,
                 container,

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

@@ -76,7 +76,7 @@ internal class ReanalyzableFunctionStructureElement(
         firIdeProvider: FirIdeProvider,
     ): ReanalyzableFunctionStructureElement {
         val originalFunction = firSymbol.fir as FirSimpleFunction
-        val newFunction = DeclarationCopyBuilder.createCopy<FirSimpleFunction, KtElement>(newKtDeclaration, originalFunction)
+        val newFunction = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalFunction)
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalFunction, newFunction) {
             firLazyDeclarationResolver.lazyResolveDeclaration(
@@ -116,7 +116,7 @@ internal class ReanalyzablePropertyStructureElement(
         firIdeProvider: FirIdeProvider,
     ): ReanalyzablePropertyStructureElement {
         val originalProperty = firSymbol.fir
-        val newProperty = DeclarationCopyBuilder.createCopy<FirProperty, KtElement>(newKtDeclaration, originalProperty)
+        val newProperty = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalProperty)
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalProperty, newProperty) {
             firLazyDeclarationResolver.lazyResolveDeclaration(

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

@@ -98,7 +98,7 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
 
     companion object {
         @OptIn(ExperimentalStdlibApi::class)
-        fun recordElementsFrom(firDeclaration: FirDeclaration, recorder: FirElementsRecorder): Map<KtElement, FirElement> =
+        fun recordElementsFrom(firElement: FirElement, recorder: FirElementsRecorder): Map<KtElement, FirElement> =
-            buildMap { firDeclaration.accept(recorder, this) }
+            buildMap { firElement.accept(recorder, this) }
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -7,9 +7,9 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
@@ -37,10 +37,20 @@ internal class FirLazyDeclarationResolver(
      */
     fun resolveFileAnnotations(
         firFile: FirFile,
+        annotations: List<FirAnnotationCall>,
         moduleFileCache: ModuleFileCache,
         scopeSession: ScopeSession,
-    ) = firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
+    ) {
-        resolveFileAnnotationsWithoutLock(firFile, scopeSession)
+        lazyResolveDeclaration(
+            declaration = firFile,
+            moduleFileCache = moduleFileCache,
+            toPhase = FirResolvePhase.IMPORTS,
+            checkPCE = false,
+            reresolveFile = false
+        )
+        firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
+            resolveFileAnnotationsWithoutLock(firFile, annotations, scopeSession)
+        }
     }
 
     /**
@@ -49,11 +59,15 @@ internal class FirLazyDeclarationResolver(
      */
     private fun resolveFileAnnotationsWithoutLock(
         firFile: FirFile,
+        annotations: List<FirAnnotationCall>,
         scopeSession: ScopeSession,
     ) {
-        check(firFile.resolvePhase >= FirResolvePhase.IMPORTS)
+        FirFileAnnotationsResolveTransformer(
-        val transformer = FirFileAnnotationsResolveTransformer(firFile.moduleData.session, scopeSession)
+            firFile = firFile,
-        firFile.accept(transformer, ResolutionMode.ContextDependent)
+            annotations = annotations,
+            session = firFile.moduleData.session,
+            scopeSession = scopeSession
+        ).transformDeclaration()
     }
 
     private fun getResolvableDeclaration(declaration: FirDeclaration, moduleFileCache: ModuleFileCache): FirDeclaration {
@@ -177,7 +191,7 @@ internal class FirLazyDeclarationResolver(
             )
         }
         if (toPhase <= nonLazyPhase) return
-        resolveFileAnnotationsWithoutLock(containerFirFile, scopeSession)
+        resolveFileAnnotationsWithoutLock(containerFirFile, containerFirFile.annotations, scopeSession)
 
         runLazyDesignatedResolveWithoutLock(
             firDeclarationToResolve = firDeclarationToResolve,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/RawFirFileAnnotationBuilder.kt

@@ -0,0 +1,35 @@
+/*
+ * 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.lazy.resolve
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtElement
+
+internal fun buildFileFirAnnotation(
+    session: FirSession,
+    baseScopeProvider: FirScopeProvider,
+    fileAnnotation: KtAnnotationEntry,
+    replacement: RawFirReplacement? = null
+): FirAnnotationCall {
+
+    val replacementApplier = replacement?.Applier()
+
+    val builder = object : RawFirBuilder(session, baseScopeProvider) {
+        inner class VisitorWithReplacement : Visitor() {
+            override fun convertElement(element: KtElement): FirElement? =
+                super.convertElement(replacementApplier?.tryReplace(element) ?: element)
+        }
+    }
+    builder.context.packageFqName = fileAnnotation.containingKtFile.packageFqName
+    val result = builder.VisitorWithReplacement().convertElement(fileAnnotation) as FirAnnotationCall
+    replacementApplier?.ensureApplied()
+    return result
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -17,87 +17,40 @@ import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignation
 import org.jetbrains.kotlin.psi.*
 
-internal data class RawFirReplacement<T : KtElement>(val from: T, val to: T)
+internal class RawFirNonLocalDeclarationBuilder private constructor(
-
-internal class RawFirNonLocalDeclarationBuilder<T : KtElement> private constructor(
     session: FirSession,
     baseScopeProvider: FirScopeProvider,
     private val declarationToBuild: KtDeclaration,
-    private val replacement: RawFirReplacement<T>? = null
+    private val replacementApplier: RawFirReplacement.Applier? = null
 ) : RawFirBuilder(session, baseScopeProvider, RawFirBuilderMode.NORMAL) {
 
-    private var replacementApplied = false
-
     companion object {
-        fun elementIsApplicable(element: KtElement) = when (element) {
-            is KtFile, is KtClassInitializer, is KtClassOrObject, is KtObjectLiteralExpression, is KtTypeAlias,
-            is KtNamedFunction, is KtLambdaExpression, is KtAnonymousInitializer, is KtProperty, is KtTypeReference,
-            is KtAnnotationEntry, is KtTypeParameter, is KtTypeProjection, is KtParameter, is KtBlockExpression,
-            is KtSimpleNameExpression, is KtConstantExpression, is KtStringTemplateExpression, is KtReturnExpression,
-            is KtTryExpression, is KtIfExpression, is KtWhenExpression, is KtDoWhileExpression, is KtWhileExpression,
-            is KtForExpression, is KtBreakExpression, is KtContinueExpression, is KtBinaryExpression, is KtBinaryExpressionWithTypeRHS,
-            is KtIsExpression, is KtUnaryExpression, is KtCallExpression, is KtArrayAccessExpression, is KtQualifiedExpression,
-            is KtThisExpression, is KtSuperExpression, is KtParenthesizedExpression, is KtLabeledExpression, is KtAnnotatedExpression,
-            is KtThrowExpression, is KtDestructuringDeclaration, is KtClassLiteralExpression, is KtCallableReferenceExpression,
-            is KtCollectionLiteralExpression -> true
-            else -> false
-        }
-
-        fun <T : KtElement> buildWithReplacement(
-            session: FirSession,
-            baseScopeProvider: FirScopeProvider,
-            designation: FirDeclarationUntypedDesignation,
-            declarationToBuild: KtDeclaration,
-            replacement: RawFirReplacement<T>? = null
-        ): FirDeclaration {
-
-            if (replacement != null) {
-                require(elementIsApplicable(replacement.from)) {
-                    "Build with replacement is possible for applicable type but given ${replacement.from::class.simpleName}"
-                }
-                require(replacement.from::class == replacement.to::class) {
-                    "Build with replacement is possible for same type in replacements but given\n${replacement.from::class.simpleName} and ${replacement.to::class.simpleName}"
-                }
-            }
-
-            val builder = RawFirNonLocalDeclarationBuilder(session, baseScopeProvider, declarationToBuild, replacement)
-            builder.context.packageFqName = declarationToBuild.containingKtFile.packageFqName
-
-            val result = builder.moveNext(designation.path.iterator(), containingClass = null)
-            check(replacement == null || builder.replacementApplied) {
-                "Replacement requested but was not applied for ${replacement!!.from::class.simpleName}"
-            }
-            return result
-        }
-
         fun build(
             session: FirSession,
             baseScopeProvider: FirScopeProvider,
             designation: FirDeclarationUntypedDesignation,
-            rootNonLocalDeclaration: KtDeclaration
+            rootNonLocalDeclaration: KtDeclaration,
+            replacement: RawFirReplacement? = null
         ): FirDeclaration {
-            val builder = RawFirNonLocalDeclarationBuilder<KtElement>(session, baseScopeProvider, rootNonLocalDeclaration)
+            val replacementApplier = replacement?.Applier()
+            val builder = RawFirNonLocalDeclarationBuilder(session, baseScopeProvider, rootNonLocalDeclaration, replacementApplier)
             builder.context.packageFqName = rootNonLocalDeclaration.containingKtFile.packageFqName
-            return builder.moveNext(designation.path.iterator(), containingClass = null)
+            return builder.moveNext(designation.path.iterator(), containingClass = null).also {
+                replacementApplier?.ensureApplied()
+            }
         }
     }
 
-    private fun KtElement.replaced(): KtElement {
-        if (replacement == null || replacement.from != this) return this
-        replacementApplied = true
-        return replacement.to
-    }
-
     private inner class VisitorWithReplacement : Visitor() {
         override fun convertElement(element: KtElement): FirElement? =
-            super.convertElement(element.replaced())
+            super.convertElement(replacementApplier?.tryReplace(element) ?: element)
 
         override fun convertProperty(
             property: KtProperty,
             ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
             ownerRegularClassTypeParametersCount: Int?
         ): FirProperty {
-            val replacementProperty = property.replaced()
+            val replacementProperty = replacementApplier?.tryReplace(property) ?: property
             check(replacementProperty is KtProperty)
             return super.convertProperty(
                 property = replacementProperty,
@@ -107,7 +60,7 @@ internal class RawFirNonLocalDeclarationBuilder<T : KtElement> private construct
         }
 
         override fun convertValueParameter(valueParameter: KtParameter, defaultTypeRef: FirTypeRef?): FirValueParameter {
-            val replacementParameter = valueParameter.replaced()
+            val replacementParameter = replacementApplier?.tryReplace(valueParameter) ?: valueParameter
             check(replacementParameter is KtParameter)
             return super.convertValueParameter(
                 valueParameter = replacementParameter,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/RawFirReplacement.kt

@@ -0,0 +1,52 @@
+/*
+ * 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.lazy.resolve
+
+import org.jetbrains.kotlin.psi.*
+
+internal data class RawFirReplacement(val from: KtElement, val to: KtElement) {
+    companion object {
+        fun isApplicableForReplacement(element: KtElement) = when (element) {
+            is KtFile, is KtClassInitializer, is KtClassOrObject, is KtObjectLiteralExpression, is KtTypeAlias,
+            is KtNamedFunction, is KtLambdaExpression, is KtAnonymousInitializer, is KtProperty, is KtTypeReference,
+            is KtAnnotationEntry, is KtTypeParameter, is KtTypeProjection, is KtParameter, is KtBlockExpression,
+            is KtSimpleNameExpression, is KtConstantExpression, is KtStringTemplateExpression, is KtReturnExpression,
+            is KtTryExpression, is KtIfExpression, is KtWhenExpression, is KtDoWhileExpression, is KtWhileExpression,
+            is KtForExpression, is KtBreakExpression, is KtContinueExpression, is KtBinaryExpression, is KtBinaryExpressionWithTypeRHS,
+            is KtIsExpression, is KtUnaryExpression, is KtCallExpression, is KtArrayAccessExpression, is KtQualifiedExpression,
+            is KtThisExpression, is KtSuperExpression, is KtParenthesizedExpression, is KtLabeledExpression, is KtAnnotatedExpression,
+            is KtThrowExpression, is KtDestructuringDeclaration, is KtClassLiteralExpression, is KtCallableReferenceExpression,
+            is KtCollectionLiteralExpression -> true
+            else -> false
+        }
+    }
+
+    inner class Applier {
+        private var replacementApplied = false
+
+        private fun ensureReplacementIsValid() {
+            require(isApplicableForReplacement(from)) {
+                "Replacement is possible for applicable type but given ${from::class.simpleName}"
+            }
+            require(from::class == to::class) {
+                "Replacement is possible for same type in replacements but given\n${from::class.simpleName} and ${to::class.simpleName}"
+            }
+        }
+
+        fun tryReplace(element: KtElement): KtElement {
+            if (from != element) return element
+            ensureReplacementIsValid()
+            replacementApplied = true
+            return to
+        }
+
+        fun ensureApplied() {
+            check(replacementApplied) {
+                "Replacement requested but was not applied for ${from::class.simpleName}"
+            }
+        }
+    }
+}

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

@@ -4,34 +4,52 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationResolveStatus
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 
 internal class FirFileAnnotationsResolveTransformer(
+    private val firFile: FirFile,
+    private val annotations: List<FirAnnotationCall>,
     session: FirSession,
     scopeSession: ScopeSession,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
+    firTowerDataContextCollector: FirTowerDataContextCollector? = null,
 ) : FirBodyResolveTransformer(
-    session,
+    session = session,
-    FirResolvePhase.BODY_RESOLVE,
+    phase = FirResolvePhase.BODY_RESOLVE,
     implicitTypeOnly = false,
-    scopeSession,
+    scopeSession = scopeSession,
-    createReturnTypeCalculatorForIDE(
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(
         session,
         scopeSession,
         implicitBodyResolveComputationSession,
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     ),
-) {
+    firTowerDataContextCollector = firTowerDataContextCollector
-    override fun transformFile(file: FirFile, data: ResolutionMode): FirFile {
+), FirLazyTransformerForIDE {
-        return context.withFile(file, components) {
+
-            file.transformAnnotations(this, data)
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration {
-            file
+        require(declaration is FirFile) { "Unexpected declaration ${declaration::class.simpleName}" }
+        annotations.forEach {
+            if (it.resolveStatus != FirAnnotationResolveStatus.Resolved) {
+                it.visitNoTransform(this, data)
+            }
         }
+        return declaration
+    }
+
+    override fun transformDeclaration() {
+        if (annotations.all { it.resolveStatus == FirAnnotationResolveStatus.Resolved }) return
+        check(firFile.resolvePhase >= FirResolvePhase.IMPORTS) { "Invalid file resolve phase ${firFile.resolvePhase}" }
+        firFile.accept(this, ResolutionMode.ContextDependent)
+        check(annotations.all { it.resolveStatus == FirAnnotationResolveStatus.Resolved }) { "Annotation was not resolved" }
     }
 
     override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = false

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/fileAnnotation.kt

@@ -0,0 +1,5 @@
+/*PLACE*/@file:helloKitty()
+
+fun f() {
+    /*ONAIR*/@file:Deprecated("hello")
+}

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/fileAnnotation.txt

@@ -0,0 +1 @@
+@FILE:R|kotlin/Deprecated|(String(hello))

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirOnAirResolveTest.kt

@@ -6,27 +6,36 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.util.io.FileUtil
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.testFramework.LightProjectDescriptor
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForResolveOnAir
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
+import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtFileAnnotationList
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.util.findElementByCommentPrefix
 import java.io.File
 
-abstract class AbstractFirOnAirResolveTest  : KotlinLightCodeInsightFixtureTestCase() {
+abstract class AbstractFirOnAirResolveTest : KotlinLightCodeInsightFixtureTestCase() {
     override fun isFirPlugin(): Boolean = true
 
     fun doTest(path: String) {
         val testDataFile = File(path)
         val ktFile = myFixture.configureByText(testDataFile.name, FileUtil.loadFile(testDataFile)) as KtFile
 
-        val place = ktFile.findElementByCommentPrefix("/*PLACE*/") as KtElement
+        fun fixUpAnnotations(element: KtElement): KtElement = when (element) {
-        val onAir = ktFile.findElementByCommentPrefix("/*ONAIR*/") as KtElement
+            is KtAnnotated -> element.annotationEntries.firstOrNull() ?: element
+            is KtFileAnnotationList -> element.annotationEntries.first()
+            else -> element
+        }
+
+        val place = (ktFile.findElementByCommentPrefix("/*PLACE*/") as KtElement).let(::fixUpAnnotations)
+        val onAir = (ktFile.findElementByCommentPrefix("/*ONAIR*/") as KtElement).let(::fixUpAnnotations)
 
         check(place::class == onAir::class)
 

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/FirOnAirResolveTestGenerated.java

@@ -34,6 +34,11 @@ public class FirOnAirResolveTestGenerated extends AbstractFirOnAirResolveTest {
         runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/classInClass.kt");
     }
 
+    @TestMetadata("fileAnnotation.kt")
+    public void testFileAnnotation() throws Exception {
+        runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/fileAnnotation.kt");
+    }
+
     @TestMetadata("identifierInContext.kt")
     public void testIdentifierInContext() throws Exception {
         runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/identifierInContext.kt");

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.tree.TokenSet
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
@@ -353,8 +354,7 @@ internal object FirReferenceResolveHelper {
         resolveState: FirModuleResolveState
     ): FirResolvedTypeRef? {
         val expressionUserType = expression.parent as? KtUserType ?: return null
-        val typeReference = PsiTreeUtil.getParentOfType(expressionUserType, KtTypeReference::class.java) ?: return null
+        val typeReference = expressionUserType.parentOfType<KtTypeReference>(withSelf = false) ?: return null
-
         return LowLevelFirApiFacadeForResolveOnAir.onAirResolveElement(
             state = resolveState,
             place = typeReference,