[FIR IDE] OnAir resolve for annotations with file site

2021-05-11 13:56:17 +02:00
parent dedcaf1eab
commit a5f701bac7
15 changed files with 248 additions and 115 deletions
@@ -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
    }
@@ -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.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
-import org.jetbrains.kotlin.fir.declarations.*
+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.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,
@@ -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,
@@ -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(
@@ -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> =
-            buildMap { firDeclaration.accept(recorder, this) }
+        fun recordElementsFrom(firElement: FirElement, recorder: FirElementsRecorder): Map<KtElement, FirElement> =
+            buildMap { firElement.accept(recorder, this) }
    }
 }
@@ -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)
-        val transformer = FirFileAnnotationsResolveTransformer(firFile.moduleData.session, scopeSession)
-        firFile.accept(transformer, ResolutionMode.ContextDependent)
+        FirFileAnnotationsResolveTransformer(
+            firFile = firFile,
+            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,
@@ -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
+}
@@ -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<T : KtElement> private constructor(
+internal class RawFirNonLocalDeclarationBuilder 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,
@@ -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}"
+            }
+        }
+    }
+}
@@ -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,
-    FirResolvePhase.BODY_RESOLVE,
+    session = session,
+    phase = FirResolvePhase.BODY_RESOLVE,
    implicitTypeOnly = false,
-    scopeSession,
-    createReturnTypeCalculatorForIDE(
+    scopeSession = scopeSession,
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(
        session,
        scopeSession,
        implicitBodyResolveComputationSession,
        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
    ),
-) {
-    override fun transformFile(file: FirFile, data: ResolutionMode): FirFile {
-        return context.withFile(file, components) {
-            file.transformAnnotations(this, data)
-            file
+    firTowerDataContextCollector = firTowerDataContextCollector
+), FirLazyTransformerForIDE {
+
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration {
+        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
@@ -0,0 +1,5 @@
+/*PLACE*/@file:helloKitty()
+
+fun f() {
+    /*ONAIR*/@file:Deprecated("hello")
+}
@@ -0,0 +1 @@
+@FILE:R|kotlin/Deprecated|(String(hello))
@@ -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
-        val onAir = ktFile.findElementByCommentPrefix("/*ONAIR*/") as KtElement
+        fun fixUpAnnotations(element: KtElement): KtElement = when (element) {
+            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)

@@ -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");
@@ -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,