[FIR IDE] LowLevel api refactoring

2021-05-07 17:08:29 +02:00
parent 445f13d1bf
commit 633e03b710
22 changed files with 303 additions and 207 deletions
@@ -43,7 +43,7 @@ class FirSupertypeResolverProcessor(session: FirSession, scopeSession: ScopeSess
 /**
 * Interceptor needed by IDE to resolve in-air created declarations.
 */
-interface FirProviderInterceptorForSupertypeResolver {
+interface FirProviderInterceptor {
    fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile?
    fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>?
 }
@@ -74,7 +74,7 @@ fun <F : FirClassLikeDeclaration<F>> F.runSupertypeResolvePhaseForLocalClass(
    scopeSession: ScopeSession,
    currentScopeList: List<FirScope>,
    localClassesNavigationInfo: LocalClassesNavigationInfo,
-    firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
+    firProviderInterceptor: FirProviderInterceptor?,
 ): F {
    val supertypeComputationSession = SupertypeComputationSession()
    val supertypeResolverVisitor = FirSupertypeResolverVisitor(
@@ -189,7 +189,7 @@ class FirSupertypeResolverVisitor(
    private val scopeSession: ScopeSession,
    private val scopeForLocalClass: PersistentList<FirScope>? = null,
    private val localClassesNavigationInfo: LocalClassesNavigationInfo? = null,
-    private val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null,
+    private val firProviderInterceptor: FirProviderInterceptor? = null,
 ) : FirDefaultVisitor<Unit, Any?>() {
    private val supertypeGenerationExtensions = session.extensionService.supertypeGenerators

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
@@ -33,7 +33,7 @@ open class FirBodyResolveTransformer(
    val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve(),
    outerBodyResolveContext: BodyResolveContext? = null,
    val firTowerDataContextCollector: FirTowerDataContextCollector? = null,
-    val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null,
+    val firProviderInterceptor: FirProviderInterceptor? = null,
 ) : FirAbstractBodyResolveTransformer(phase) {

    final override val context: BodyResolveContext =
@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.runStatusResolveForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runSupertypeResolvePhaseForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runTypeResolvePhaseForLocalClass
@@ -20,7 +20,7 @@ fun <F : FirClassLikeDeclaration<F>> F.runAllPhasesForLocalClass(
    components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
    resolutionMode: ResolutionMode,
    firTowerDataContextCollector: FirTowerDataContextCollector?,
-    firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
+    firProviderInterceptor: FirProviderInterceptor?,
 ): F {
    if (this.resolvePhase > FirResolvePhase.RAW_FIR) return this
    if (this is FirRegularClass) {
@@ -58,13 +58,13 @@ internal class FirModuleResolveStateDepended(
        )
    }

-    override fun getFirFile(ktFile: KtFile): FirFile =
-        originalState.getFirFile(ktFile)
+    override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
+        originalState.getOrBuildFirFile(ktFile)

-    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
-        originalState.resolvedFirToPhase(declaration, toPhase)
+    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
+        originalState.resolveFirToPhase(declaration, toPhase)

-    override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
+    override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
        val ktFile = declaration.containingKtFileIfAny ?: return null
        cache.getCachedFirFile(ktFile)?.let { return it }
        ktFile.originalKtFile?.let(cache::getCachedFirFile)?.let { return it }
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api

 import com.intellij.openapi.project.Project
-import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
@@ -60,20 +59,18 @@ internal class FirModuleResolveStateImpl(
    override fun getOrBuildFirFor(element: KtElement): FirElement =
        elementBuilder.getOrBuildFirFor(element, firFileBuilder, rootModuleSession.cache, fileStructureCache)

-    override fun getFirFile(ktFile: KtFile): FirFile =
+    override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, lazyBodiesMode = false)

+    override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
+        cache.getContainerFirFile(declaration)
+
    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>> =
        diagnosticsCollector.getDiagnosticsFor(element, filter)

    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
        diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)

-    internal fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile? {
-        val cache = sessionProvider.getModuleCache(ktFile.getModuleInfo() as ModuleSourceInfo)
-        return firFileBuilder.getBuiltFirFileOrNull(ktFile, cache)
-    }
-
    @OptIn(InternalForInline::class)
    override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
        findSourceFirDeclarationByExpression(ktDeclaration)
@@ -117,7 +114,7 @@ internal class FirModuleResolveStateImpl(
            ?: error("FirDeclaration was not found for\n${ktDeclaration.getElementTextInContext()}")
    }

-    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
+    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
        val fileCache = when (val session = declaration.moduleData.session) {
            is FirIdeSourcesSession -> session.cache
            else -> return declaration
@@ -127,11 +124,7 @@ internal class FirModuleResolveStateImpl(
            fileCache,
            toPhase,
            checkPCE = true,
-            towerDataContextCollector = null,
        )
        return declaration
    }
-
-    override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
-        cache.getContainerFirFile(declaration)
 }
@@ -58,11 +58,11 @@ object DeclarationCopyBuilder {
        val originalFunction = rootNonLocalDeclaration.getOrBuildFirOfType<FirSimpleFunction>(state)
        val builtFunction = createCopy(rootNonLocalDeclaration, originalFunction, replacement)

-        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
        val functionBlock = rootNonLocalDeclaration.bodyBlockExpression
        if (functionBlock == null || !PsiTreeUtil.isAncestor(functionBlock, replacement.from, true)) {
            return builtFunction
        }
+        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?

        // right now we can't resolve builtFunction header properly, as it built right in air,
        // without file, which is now required for running stages other then body resolve, so we
@@ -82,11 +82,11 @@ object DeclarationCopyBuilder {
        val originalFirClass = rootNonLocalDeclaration.getOrBuildFirOfType<FirRegularClass>(state)
        val builtClass = createCopy(rootNonLocalDeclaration, originalFirClass, replacement)

-        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
        val classBody = rootNonLocalDeclaration.body
        if (classBody == null || !PsiTreeUtil.isAncestor(classBody, replacement.from, true)) {
            return builtClass
        }
+        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?

        return buildRegularClassCopy(originalFirClass) {
            declarations.clear()
@@ -102,16 +102,8 @@ object DeclarationCopyBuilder {
        state: FirModuleResolveState,
        replacement: RawFirReplacement<T>,
    ): FirTypeAlias {
-
        val originalFirTypeAlias = rootNonLocalDeclaration.getOrBuildFirOfType<FirTypeAlias>(state)
-        val builtTypeAlias = createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
-        return builtTypeAlias
-
-//        return buildTypeAliasCopy(originalFirTypeAlias) {
-//            expandedTypeRef = builtTypeAlias.expandedTypeRef
-//            symbol = builtTypeAlias.symbol
-//            initDeclaration(originalFirTypeAlias, builtTypeAlias, state)
-//        }
+        return createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
    }

    private fun <T : KtElement> createPropertyCopy(
@@ -122,7 +114,6 @@ object DeclarationCopyBuilder {
        val originalProperty = rootNonLocalDeclaration.getOrBuildFirOfType<FirProperty>(state)
        val builtProperty = createCopy(rootNonLocalDeclaration, originalProperty, replacement)

-        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
        val insideGetterBody = rootNonLocalDeclaration.getter?.bodyBlockExpression?.let {
            PsiTreeUtil.isAncestor(it, replacement.from, true)
        } ?: false
@@ -133,6 +124,8 @@ object DeclarationCopyBuilder {

            if (!insideSetterBody) return builtProperty
        }
+        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
+

        val originalSetter = originalProperty.setter
        val builtSetter = builtProperty.setter
@@ -29,9 +29,22 @@ abstract class FirModuleResolveState {

    internal abstract fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession

+    /**
+     * Build fully resolved FIR node for requested element.
+     * This operation could be performance affective, use
+     * @see tryGetCachedFirFile to get [FirFile] in undefined phase
+     */
    internal abstract fun getOrBuildFirFor(element: KtElement): FirElement

-    internal abstract fun getFirFile(ktFile: KtFile): FirFile
+    /**
+     * Get or build or get cached [FirFile] for requested file in undefined phase
+     */
+    internal abstract fun getOrBuildFirFile(ktFile: KtFile): FirFile
+
+    /**
+     * Try get [FirFile] from the cache in undefined phase
+     */
+    internal abstract fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?

    internal abstract fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>>

@@ -41,11 +54,9 @@ abstract class FirModuleResolveState {
        val originalDeclaration = (declaration as? FirCallableDeclaration<*>)?.unwrapFakeOverrides() ?: declaration
        val session = originalDeclaration.moduleData.session
        return when {
-            originalDeclaration.origin == FirDeclarationOrigin.Source
-                    && session is FirIdeSourcesSession
-            -> {
+            originalDeclaration.origin == FirDeclarationOrigin.Source && session is FirIdeSourcesSession -> {
                val cache = session.cache
-                val file = getFirFile(declaration, cache)
+                val file = tryGetCachedFirFile(declaration, cache)
                    ?: error("Fir file was not found for\n${declaration.render()}\n${(declaration.psi as? KtElement)?.getElementTextInContext()}")
                cache.firFileLockProvider.withLock(file, declarationLockType) { action(declaration) }
            }
@@ -63,7 +74,5 @@ abstract class FirModuleResolveState {
        ktDeclaration: KtLambdaExpression,
    ): FirDeclaration

-    internal abstract fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D
-
-    internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?
+    internal abstract fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 }
@@ -8,15 +8,10 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.fir.unwrapFakeOverrides
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
-import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -118,12 +113,14 @@ fun <D : FirDeclaration, R> D.withFirDeclarationInWriteLock(

 /**
 * Returns a list of Diagnostics compiler finds for given [KtElement]
+ * This operation could be performance affective
 */
 fun KtElement.getDiagnostics(resolveState: FirModuleResolveState, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
    resolveState.getDiagnostics(this, filter)

 /**
 * Returns a list of Diagnostics compiler finds for given [KtFile]
+ * This operation could be performance affective
 */
 fun KtFile.collectDiagnosticsForFile(
    resolveState: FirModuleResolveState,
@@ -140,12 +137,13 @@ fun <D : FirDeclaration> D.resolvedFirToPhase(
    phase: FirResolvePhase,
    resolveState: FirModuleResolveState
 ): D =
-    resolveState.resolvedFirToPhase(this, phase)
+    resolveState.resolveFirToPhase(this, phase)


 /**
 * Get a [FirElement] which was created by [KtElement]
 * Returned [FirElement] is guaranteed to be resolved to [FirResolvePhase.BODY_RESOLVE] phase
+ * This operation could be performance affective
 */
 fun KtElement.getOrBuildFir(
    resolveState: FirModuleResolveState,
@@ -154,6 +152,7 @@ fun KtElement.getOrBuildFir(
 /**
 * Get a [FirElement] which was created by [KtElement], but only if it is subtype of [E], `null` otherwise
 * Returned [FirElement] is guaranteed to be resolved to [FirResolvePhase.BODY_RESOLVE] phase
+ * This operation could be performance affective
 */
 inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
    resolveState: FirModuleResolveState,
@@ -162,6 +161,7 @@ inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
 /**
 * Get a [FirElement] which was created by [KtElement], but only if it is subtype of [E], throws [InvalidFirElementTypeException] otherwise
 * Returned [FirElement] is guaranteed to be resolved to [FirResolvePhase.BODY_RESOLVE] phase
+ * This operation could be performance affective
 */
 inline fun <reified E : FirElement> KtElement.getOrBuildFirOfType(
    resolveState: FirModuleResolveState,
@@ -175,8 +175,8 @@ inline fun <reified E : FirElement> KtElement.getOrBuildFirOfType(
 * Get a [FirFile] which was created by [KtElement]
 * Returned [FirFile] can be resolved to any phase from [FirResolvePhase.RAW_FIR] to [FirResolvePhase.BODY_RESOLVE]
 */
-fun KtFile.getFirFile(resolveState: FirModuleResolveState): FirFile =
-    resolveState.getFirFile(this)
+fun KtFile.getOrBuildFirFile(resolveState: FirModuleResolveState): FirFile =
+    resolveState.getOrBuildFirFile(this)

 class InvalidFirElementTypeException(
    ktElement: KtElement,
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.builder.RawFirFragmentForLazyBodiesBuilder
 import org.jetbrains.kotlin.fir.builder.RawFirReplacement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.realPsi
@@ -24,7 +25,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FileTowerProv
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
 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.transformers.FirProviderInterceptorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
@@ -72,8 +72,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
        val declaration = runBodyResolveOnAir(
            state = state,
            replacement = RawFirReplacement(place, elementToResolve),
-            collector = null,
-            useFirProviderInterceptor = true
+            isOnAirResolve = true,
        )

        val expressionLocator = object : FirVisitorVoid() {
@@ -101,8 +100,17 @@ object LowLevelFirApiFacadeForResolveOnAir {
        return if (place is KtFile) {
            FileTowerProvider(place, onAirGetTowerContextForFile(state, place))
        } else {
+            val validPlace = PsiTreeUtil.findFirstParent(place, false) {
+                RawFirFragmentForLazyBodiesBuilder.elementIsApplicable(it as KtElement)
+            } as KtElement
+
            FirTowerDataContextAllElementsCollector().also {
-                runBodyResolveOnAir(state, collector = it, replacement = RawFirReplacement(place, place))
+                runBodyResolveOnAir(
+                    state = state,
+                    collector = it,
+                    replacement = RawFirReplacement(validPlace, validPlace),
+                    isOnAirResolve = false //isOnAirResolve can be little faster because node resolved in it's context
+                )
            }
        }
    }
@@ -154,9 +162,9 @@ object LowLevelFirApiFacadeForResolveOnAir {
        val collector = FirTowerDataContextAllElementsCollector()
        val copiedFirDeclaration = runBodyResolveOnAir(
            originalState,
-            collector = collector,
            replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration),
-            useFirProviderInterceptor = true
+            isOnAirResolve = true,
+            collector = collector,
        )

        val recordedMap = FirElementsRecorder.recordElementsFrom(copiedFirDeclaration, FirElementsRecorder())
@@ -166,8 +174,8 @@ object LowLevelFirApiFacadeForResolveOnAir {
    private fun <T : KtElement> runBodyResolveOnAir(
        state: FirModuleResolveStateImpl,
        replacement: RawFirReplacement<T>,
+        isOnAirResolve: Boolean,
        collector: FirTowerDataContextCollector? = null,
-        useFirProviderInterceptor: Boolean = false
    ): FirDeclaration {

        val nonLocalDeclaration = findEnclosingNonLocalDeclaration(replacement.from)
@@ -179,30 +187,18 @@ object LowLevelFirApiFacadeForResolveOnAir {
            replacement = replacement,
        )

-        val originalFirFile = state.getBuiltFirFileOrNull(replacement.from.containingKtFile)
-            ?: error("Original fir file should be already built")
+        val originalFirFile = state.getOrBuildFirFile(nonLocalDeclaration.containingKtFile)

-        val firProviderInterceptor =
-            if (useFirProviderInterceptor) FirProviderInterceptorForIDE.createForFirElement(
-                session = originalFirFile.declarationSiteSession,
-                firFile = originalFirFile,
-                element = copiedFirDeclaration
-            ) else null
-
-        state.firFileBuilder.runCustomResolveWithPCECheck(originalFirFile, state.rootModuleSession.cache) {
-            state.firLazyDeclarationResolver.runLazyResolveWithoutLock(
-                copiedFirDeclaration,
-                state.rootModuleSession.cache,
-                originalFirFile,
-                originalFirFile.declarationSiteSession.firIdeProvider,
-                fromPhase = copiedFirDeclaration.resolvePhase,
-                toPhase = FirResolvePhase.BODY_RESOLVE,
-                towerDataContextCollector = collector,
-                checkPCE = true,
-                lastNonLazyPhase = FirResolvePhase.IMPORTS,
-                firProviderInterceptor = firProviderInterceptor,
-            )
-        }
+        state.firLazyDeclarationResolver.lazyDesignatedResolveDeclaration(
+            firDeclarationToResolve = copiedFirDeclaration,
+            moduleFileCache = state.rootModuleSession.cache,
+            containerFirFile = originalFirFile,
+            provider = originalFirFile.declarationSiteSession.firIdeProvider,
+            toPhase = FirResolvePhase.BODY_RESOLVE,
+            checkPCE = true,
+            isOnAirResolve = isOnAirResolve,
+            towerDataContextCollector = collector,
+        )

        return copiedFirDeclaration
    }
@@ -21,8 +21,8 @@ import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor

 internal object FirLazyBodiesCalculator {
-    fun calculateLazyBodiesInside(element: FirElement, designation: FirDeclarationDesignation) {
-        element.transform<FirElement, MutableList<FirDeclaration>>(
+    fun calculateLazyBodiesInside(designation: FirDeclarationDesignation) {
+        designation.declaration.transform<FirElement, MutableList<FirDeclaration>>(
            FirLazyBodiesCalculatorTransformer,
            designation.toSequence(includeTarget = true).toMutableList()
        )
@@ -13,7 +13,7 @@ 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
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
@@ -24,34 +24,39 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.*
+import org.jetbrains.kotlin.idea.util.ifTrue
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject

 internal class FirLazyDeclarationResolver(
    private val firFileBuilder: FirFileBuilder
 ) {
+    /**
+     * Fully resolve file annotations (synchronized)
+     * @see resolveFileAnnotationsWithoutLock not synchronized
+     */
    fun resolveFileAnnotations(
        firFile: FirFile,
        moduleFileCache: ModuleFileCache,
        scopeSession: ScopeSession,
-    ) {
-        check(firFile.resolvePhase >= FirResolvePhase.IMPORTS)
-        firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
-            val transformer = FirFileAnnotationsResolveTransformer(firFile.moduleData.session, scopeSession)
-            firFile.accept(transformer, ResolutionMode.ContextDependent)
-        }
+    ) = firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
+        resolveFileAnnotationsWithoutLock(firFile, scopeSession)
    }

-    fun lazyResolveDeclaration(
-        declaration: FirDeclaration,
-        moduleFileCache: ModuleFileCache,
-        toPhase: FirResolvePhase,
-        towerDataContextCollector: FirTowerDataContextCollector? = null,
-        checkPCE: Boolean = false,
-        reresolveFile: Boolean = false,
+    /**
+     * Fully resolve file annotations (not synchronized)
+     * @see resolveFileAnnotations synchronized version
+     */
+    private fun resolveFileAnnotationsWithoutLock(
+        firFile: FirFile,
+        scopeSession: ScopeSession,
    ) {
-        if (declaration.resolvePhase >= toPhase) return
+        check(firFile.resolvePhase >= FirResolvePhase.IMPORTS)
+        val transformer = FirFileAnnotationsResolveTransformer(firFile.moduleData.session, scopeSession)
+        firFile.accept(transformer, ResolutionMode.ContextDependent)
+    }

+    private fun getResolvableDeclaration(declaration: FirDeclaration, moduleFileCache: ModuleFileCache): FirDeclaration {
        if (declaration is FirPropertyAccessor || declaration is FirTypeParameter || declaration is FirValueParameter) {
            val ktContainingResolvableDeclaration = when (val psi = declaration.psi) {
                is KtPropertyAccessor -> psi.property
@@ -66,74 +71,105 @@ internal class FirLazyDeclarationResolver(
                null -> error("Cannot find containing declaration for KtParameter")
                else -> error("Invalid source of property accessor ${psi::class}")
            }
-            val containingProperty = ktContainingResolvableDeclaration
-                .findSourceNonLocalFirDeclaration(firFileBuilder, declaration.moduleData.session.symbolProvider, moduleFileCache)
-            return lazyResolveDeclaration(containingProperty, moduleFileCache, toPhase, towerDataContextCollector)
+            return ktContainingResolvableDeclaration.findSourceNonLocalFirDeclaration(
+                firFileBuilder = firFileBuilder,
+                firSymbolProvider = declaration.moduleData.session.symbolProvider,
+                moduleFileCache = moduleFileCache
+            )
        }

-        val firFile = declaration.getContainingFile()
-            ?: error("FirFile was not found for\n${declaration.render()}")
+        return declaration
+    }
+
+    /**
+     * Run partially designated resolve that resolve declaration into last file-wise resolve and then resolve a designation (synchronized)
+     * @see LAST_NON_LAZY_PHASE is the last file-wise resolve
+     * @see lazyDesignatedResolveDeclaration designated resolve
+     * @see runLazyResolveWithoutLock (not synchronized)
+     */
+    fun lazyResolveDeclaration(
+        declaration: FirDeclaration,
+        moduleFileCache: ModuleFileCache,
+        toPhase: FirResolvePhase,
+        checkPCE: Boolean = false,
+        reresolveFile: Boolean = false,
+    ) {
+        if (declaration.resolvePhase >= toPhase) return
+
+        val resolvableDeclaration = getResolvableDeclaration(declaration, moduleFileCache)
+
+        val firFile = resolvableDeclaration.getContainingFile()
+            ?: error("FirFile was not found for\n${resolvableDeclaration.render()}")
        val provider = firFile.moduleData.session.firIdeProvider
+
        // Lazy since we want to read the resolve phase inside the lock. Otherwise, we may run the same resolve phase multiple times. See
        // KT-45121
        val fromPhase: FirResolvePhase by lazy(LazyThreadSafetyMode.NONE) {
-            if (reresolveFile) declaration.resolvePhase else minOf(firFile.resolvePhase, declaration.resolvePhase)
+            if (reresolveFile) resolvableDeclaration.resolvePhase else minOf(firFile.resolvePhase, resolvableDeclaration.resolvePhase)
        }

        if (checkPCE) {
            firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
                runLazyResolveWithoutLock(
-                    declaration,
-                    moduleFileCache,
-                    firFile,
-                    provider,
-                    fromPhase,
-                    toPhase,
-                    towerDataContextCollector,
-                    checkPCE = true
+                    firDeclarationToResolve = declaration,
+                    moduleFileCache = moduleFileCache,
+                    containerFirFile = firFile,
+                    provider = provider,
+                    fromPhase = fromPhase,
+                    toPhase = toPhase,
+                    checkPCE = true,
                )
            }
        } else {
            firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
                executeWithoutPCE {
                    runLazyResolveWithoutLock(
-                        declaration,
-                        moduleFileCache,
-                        firFile,
-                        provider,
-                        fromPhase,
-                        toPhase,
-                        towerDataContextCollector,
-                        checkPCE = false
+                        firDeclarationToResolve = declaration,
+                        moduleFileCache = moduleFileCache,
+                        containerFirFile = firFile,
+                        provider = provider,
+                        fromPhase = fromPhase,
+                        toPhase = toPhase,
+                        checkPCE = false,
                    )
                }
            }
        }
    }

-    private fun calculateLazyBodies(firDeclaration: FirDeclaration, designation: FirDeclarationDesignation) {
-        FirLazyBodiesCalculator.calculateLazyBodiesInside(firDeclaration, designation)
+    private fun createLazyBodiesCalculator(designation: FirDeclarationDesignation): (FirResolvePhase) -> Unit {
+        var calculated = false
+        return { phase: FirResolvePhase ->
+            if (!calculated && phase >= FirResolvePhase.CONTRACTS) {
+                executeWithoutPCE {
+                    FirLazyBodiesCalculator.calculateLazyBodiesInside(designation)
+                }
+                calculated = true
+            }
+        }
    }

-    fun runLazyResolveWithoutLock(
+    /**
+     * Designated resolve (not synchronized)
+     * @see runLazyDesignatedResolveWithoutLock for designated resolve
+     * @see lazyResolveDeclaration synchronized version
+     */
+    private fun runLazyResolveWithoutLock(
        firDeclarationToResolve: FirDeclaration,
        moduleFileCache: ModuleFileCache,
        containerFirFile: FirFile,
        provider: FirProvider,
        fromPhase: FirResolvePhase,
        toPhase: FirResolvePhase,
-        towerDataContextCollector: FirTowerDataContextCollector? = null,
        checkPCE: Boolean,
-        lastNonLazyPhase: FirResolvePhase = LAST_NON_LAZY_PHASE,
-        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null
    ) {
        if (fromPhase >= toPhase) return
-        val nonLazyPhase = minOf(toPhase, lastNonLazyPhase)
+        val nonLazyPhase = minOf(toPhase, LAST_NON_LAZY_PHASE)

        val scopeSession = ScopeSession()
        if (fromPhase < nonLazyPhase) {
            firFileBuilder.runResolveWithoutLock(
-                containerFirFile,
+                firFile = containerFirFile,
                fromPhase = fromPhase,
                toPhase = nonLazyPhase,
                scopeSession = scopeSession,
@@ -141,27 +177,130 @@ internal class FirLazyDeclarationResolver(
            )
        }
        if (toPhase <= nonLazyPhase) return
-        resolveFileAnnotations(containerFirFile, moduleFileCache, scopeSession)
+        resolveFileAnnotationsWithoutLock(containerFirFile, scopeSession)
+
+        runLazyDesignatedResolveWithoutLock(
+            firDeclarationToResolve = firDeclarationToResolve,
+            moduleFileCache = moduleFileCache,
+            containerFirFile = containerFirFile,
+            provider = provider,
+            fromPhase = LAST_NON_LAZY_PHASE,
+            toPhase = toPhase,
+            checkPCE = checkPCE,
+            isOnAirResolve = false
+        )
+    }
+
+    /**
+     * Run designated resolve only designation with fully resolved path (synchronized).
+     * Suitable for body resolve or/and on-air resolve.
+     * @see lazyResolveDeclaration for ordinary resolve
+     * @param firDeclarationToResolve target non-local declaration
+     * @param isOnAirResolve should be true when node does not belong to it's true designation (OnAir resolve in custom context)
+     */
+    fun lazyDesignatedResolveDeclaration(
+        firDeclarationToResolve: FirDeclaration,
+        moduleFileCache: ModuleFileCache,
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        toPhase: FirResolvePhase,
+        checkPCE: Boolean,
+        isOnAirResolve: Boolean,
+        towerDataContextCollector: FirTowerDataContextCollector? = null,
+    ) {
+        // Lazy since we want to read the resolve phase inside the lock. Otherwise, we may run the same resolve phase multiple times. See
+        // KT-45121
+        val fromPhase: FirResolvePhase by lazy(LazyThreadSafetyMode.NONE) {
+            minOf(containerFirFile.resolvePhase, firDeclarationToResolve.resolvePhase)
+        }
+
+        if (checkPCE) {
+            firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, moduleFileCache) {
+                runLazyDesignatedResolveWithoutLock(
+                    firDeclarationToResolve = firDeclarationToResolve,
+                    moduleFileCache = moduleFileCache,
+                    containerFirFile = containerFirFile,
+                    provider = provider,
+                    fromPhase = fromPhase,
+                    toPhase = toPhase,
+                    checkPCE = checkPCE,
+                    isOnAirResolve = isOnAirResolve,
+                    towerDataContextCollector = towerDataContextCollector,
+                )
+            }
+        } else {
+            firFileBuilder.runCustomResolveUnderLock(containerFirFile, moduleFileCache) {
+                runLazyDesignatedResolveWithoutLock(
+                    firDeclarationToResolve = firDeclarationToResolve,
+                    moduleFileCache = moduleFileCache,
+                    containerFirFile = containerFirFile,
+                    provider = provider,
+                    fromPhase = fromPhase,
+                    toPhase = toPhase,
+                    checkPCE = checkPCE,
+                    isOnAirResolve = isOnAirResolve,
+                    towerDataContextCollector = towerDataContextCollector,
+                )
+            }
+        }
+    }
+
+    /**
+     * Designated resolve (not synchronized)
+     * @see runLazyResolveWithoutLock for ordinary resolve
+     * @see lazyDesignatedResolveDeclaration synchronized version
+     */
+    private fun runLazyDesignatedResolveWithoutLock(
+        firDeclarationToResolve: FirDeclaration,
+        moduleFileCache: ModuleFileCache,
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        fromPhase: FirResolvePhase,
+        toPhase: FirResolvePhase,
+        checkPCE: Boolean,
+        isOnAirResolve: Boolean,
+        towerDataContextCollector: FirTowerDataContextCollector? = null
+    ) {
+        var currentPhase = fromPhase
+        runLazyResolveWithoutLock(
+            firDeclarationToResolve = firDeclarationToResolve,
+            moduleFileCache = moduleFileCache,
+            containerFirFile = containerFirFile,
+            provider = provider,
+            fromPhase = currentPhase,
+            toPhase = FirResolvePhase.IMPORTS,
+            checkPCE = checkPCE
+        )
+        currentPhase = maxOf(fromPhase, FirResolvePhase.IMPORTS)
+        if (currentPhase >= toPhase) return

        val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
        val designation = nonLocalDeclarationToResolve.collectDesignation(containerFirFile)
+        check(!designation.isLocalDesignation) { "Could not resolve local designation" }

-        executeWithoutPCE {
-            calculateLazyBodies(firDeclarationToResolve, designation)
+        val lazyBodiesCalculator = createLazyBodiesCalculator(designation)
+        val scopeSession = ScopeSession()
+
+        //This needed to override standard symbol resolve in supertype transformer with adding on-air created symbols
+        val firProviderInterceptor = isOnAirResolve.ifTrue {
+            FirProviderInterceptorForIDE.createForFirElement(
+                session = designation.firFile.moduleData.session,
+                firFile = designation.firFile,
+                element = designation.declaration
+            )
        }

-        var currentPhase = nonLazyPhase
-
        while (currentPhase < toPhase) {
            currentPhase = currentPhase.next
            if (currentPhase.pluginPhase) continue
            if (checkPCE) checkCanceled()
+            lazyBodiesCalculator(currentPhase)
            runLazyResolvePhase(
-                currentPhase,
-                scopeSession,
-                towerDataContextCollector,
-                firProviderInterceptor,
-                designation,
+                phase = currentPhase,
+                scopeSession = scopeSession,
+                designation = designation,
+                towerDataContextCollector = towerDataContextCollector,
+                firProviderInterceptor = firProviderInterceptor
            )
        }
    }
@@ -169,9 +308,9 @@ internal class FirLazyDeclarationResolver(
    private fun runLazyResolvePhase(
        phase: FirResolvePhase,
        scopeSession: ScopeSession,
+        designation: FirDeclarationDesignationWithFile,
        towerDataContextCollector: FirTowerDataContextCollector?,
-        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
-        designation: FirDeclarationDesignationWithFile
+        firProviderInterceptor: FirProviderInterceptor?
    ) {
        if (designation.declaration.resolvePhase >= phase) return

@@ -191,7 +330,7 @@ internal class FirLazyDeclarationResolver(
        designation: FirDeclarationDesignationWithFile,
        scopeSession: ScopeSession,
        towerDataContextCollector: FirTowerDataContextCollector?,
-        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
+        firProviderInterceptor: FirProviderInterceptor?,
    ): FirLazyTransformerForIDE = when (this) {
        FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
            designation,
@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
@@ -23,7 +23,7 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
    session: FirSession,
    scopeSession: ScopeSession,
    towerDataContextCollector: FirTowerDataContextCollector? = null,
-    firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null,
+    firProviderInterceptor: FirProviderInterceptor? = null,
 ) : FirLazyTransformerForIDE, FirBodyResolveTransformer(
    session,
    phase = FirResolvePhase.BODY_RESOLVE,
@@ -17,42 +17,12 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePathPhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase

-class FirDesignatedStatusResolveTransformerForIDE(
+internal class FirDesignatedStatusResolveTransformerForIDE(
    private val designation: FirDeclarationDesignationWithFile,
    private val session: FirSession,
    private val scopeSession: ScopeSession,
 ) : FirLazyTransformerForIDE {

-    private val firstItemInDesignation = designation.path.firstOrNull() ?: designation.declaration
-
-    private fun resolveClass(targetClass: FirClass<*>) {
-        val transformer = FirDesignatedStatusResolveTransformer(
-            session = session,
-            scopeSession = scopeSession,
-            designation = designation.toSequence(includeTarget = true).iterator(),
-            targetClass = targetClass,
-            statusComputationSession = StatusComputationSession.Regular(),
-            designationMapForLocalClasses = emptyMap(),
-            scopeForLocalClass = null
-        )
-
-        firstItemInDesignation.transformSingle(transformer, null)
-    }
-
-    private fun resolveTypeAlias(targetClass: FirTypeAlias) {
-        val transformer = FirDesignatedStatusResolveTransformer(
-            session = session,
-            scopeSession = scopeSession,
-            designation = designation.toSequence(includeTarget = true).iterator(),
-            targetClass = targetClass,
-            statusComputationSession = StatusComputationSession.Regular(),
-            designationMapForLocalClasses = emptyMap(),
-            scopeForLocalClass = null
-        )
-
-        firstItemInDesignation.transformSingle(transformer, null)
-    }
-
    private fun resolveTopLevelDeclaration(declaration: FirDeclaration) {
        val transformer = FirStatusResolveTransformer(
            session = session,
@@ -62,13 +32,13 @@ class FirDesignatedStatusResolveTransformerForIDE(
        declaration.transformSingle(transformer, null)
    }

-    private fun resolveClassMember(containingClass: FirClass<*>, targetCallable: FirDeclaration) {
+    private fun resolveClassMember(containingClass: FirClass<*>, targetDeclaration: FirDeclaration) {

        val transformer = object : FirDesignatedStatusResolveTransformer(
            session = session,
            scopeSession = scopeSession,
            designation = designation.toSequence(includeTarget = true).iterator(),
-            targetClass = if (targetCallable is FirRegularClass) targetCallable else containingClass,
+            targetClass = if (targetDeclaration is FirRegularClass) targetDeclaration else containingClass,
            statusComputationSession = StatusComputationSession.Regular(),
            designationMapForLocalClasses = emptyMap(),
            scopeForLocalClass = null
@@ -77,12 +47,13 @@ class FirDesignatedStatusResolveTransformerForIDE(
            override fun <F : FirClass<F>> transformClass(klass: FirClass<F>, data: FirResolvedDeclarationStatus?): FirStatement {
                if (klass != containingClass) return super.transformClass(klass, data)
                val result = storeClass(klass) {
-                    targetCallable.transformSingle(this, data)
+                    targetDeclaration.transformSingle(this, data)
                }
                return result as FirStatement
            }
        }

+        val firstItemInDesignation = designation.path.firstOrNull() ?: designation.declaration
        firstItemInDesignation.transformSingle(transformer, null)
    }

@@ -95,7 +66,7 @@ class FirDesignatedStatusResolveTransformerForIDE(
        if (containingClass == null) {
            resolveTopLevelDeclaration(designation.declaration)
        } else {
-            check(containingClass is FirClass<*>) { "Invalid designation - the parent is a class" }
+            check(containingClass is FirClass<*>) { "Invalid designation - the parent is not a class" }
            resolveClassMember(containingClass, designation.declaration)
        }

@@ -10,27 +10,26 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirApplySupertypesTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSupertypeResolverVisitor
 import org.jetbrains.kotlin.fir.resolve.transformers.SupertypeComputationSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePathPhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhaseIfClass

 internal class FirDesignatedSupertypeResolverTransformerForIDE(
    private val designation: FirDeclarationDesignationWithFile,
    private val session: FirSession,
    private val scopeSession: ScopeSession,
-    private val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
+    private val firProviderInterceptor: FirProviderInterceptor?
 ) : FirLazyTransformerForIDE {

-    private val supertypeComputationSession = SupertypeComputationSession()
-
    override fun transformDeclaration() {
        if (designation.declaration.resolvePhase >= FirResolvePhase.SUPER_TYPES) return
        designation.ensurePathPhase(FirResolvePhase.SUPER_TYPES)

+        val supertypeComputationSession = SupertypeComputationSession()
+
        val resolver = FirSupertypeResolverVisitor(
            session = session,
            supertypeComputationSession = supertypeComputationSession,
@@ -17,7 +17,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransform
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhaseIfClass

-class FirDesignatedTypeResolverTransformerForIDE(
+internal class FirDesignatedTypeResolverTransformerForIDE(
    private val designation: FirDeclarationDesignationWithFile,
    session: FirSession,
    scopeSession: ScopeSession,
@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers

 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation

 internal interface FirLazyTransformerForIDE {
    fun transformDeclaration()
@@ -16,7 +16,7 @@ internal interface FirLazyTransformerForIDE {
            override fun transformDeclaration() = Unit
        }

-        fun FirDeclarationDesignationWithFile.ensurePathPhase(firResolvePhase: FirResolvePhase) {
+        fun FirDeclarationDesignation.ensurePathPhase(firResolvePhase: FirResolvePhase) {
            toSequence(includeTarget = false).forEach { firDeclaration ->
                check(firDeclaration.resolvePhase >= firResolvePhase) {
                    "Designation element phase required to be $firResolvePhase but element resolved to ${firDeclaration.resolvePhase}"
@@ -24,19 +24,13 @@ internal interface FirLazyTransformerForIDE {
            }
        }

-        fun FirDeclarationDesignationWithFile.ensureTargetPhase(firResolvePhase: FirResolvePhase) =
+        fun FirDeclarationDesignation.ensureTargetPhase(firResolvePhase: FirResolvePhase) =
            check(declaration.resolvePhase >= firResolvePhase) { "Expected $firResolvePhase but found ${declaration.resolvePhase}" }

-        fun FirDeclarationDesignationWithFile.ensureTargetPhaseIfClass(firResolvePhase: FirResolvePhase) = when (declaration) {
+        fun FirDeclarationDesignation.ensureTargetPhaseIfClass(firResolvePhase: FirResolvePhase) = when (declaration) {
            is FirProperty, is FirSimpleFunction -> Unit
            is FirClass<*>, is FirTypeAlias -> ensureTargetPhase(firResolvePhase)
            else -> error("Unexpected target")
        }
-
-        fun FirDeclarationDesignationWithFile.ensureTargetPhaseIfMember(firResolvePhase: FirResolvePhase) = when (declaration) {
-            is FirProperty, is FirSimpleFunction -> ensureTargetPhase(firResolvePhase)
-            is FirClass<*>, is FirTypeAlias -> Unit
-            else -> error("Unexpected target")
-        }
    }
 }
@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.firProvider
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
@@ -20,7 +20,7 @@ internal class FirProviderInterceptorForIDE private constructor(
    private val session: FirSession,
    private val symbolSet: Set<FirClassLikeSymbol<*>>,
    private val classIdToElementMap: Map<ClassId, FirClassLikeDeclaration<*>>
-) : FirProviderInterceptorForSupertypeResolver {
+) : FirProviderInterceptor {

    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
        if (symbolSet.contains(symbol)) firFile else session.firProvider.getFirClassifierContainerFileIfAny(symbol)
@@ -29,7 +29,7 @@ internal class FirProviderInterceptorForIDE private constructor(
        classIdToElementMap[classId] ?: session.firProvider.getFirClassifierByFqName(classId)

    companion object {
-        fun createForFirElement(session: FirSession, firFile: FirFile, element: FirElement): FirProviderInterceptorForSupertypeResolver {
+        fun createForFirElement(session: FirSession, firFile: FirFile, element: FirElement): FirProviderInterceptor {
            val nodeInfoCollector = object : FirVisitorVoid() {
                val symbolSet = mutableSetOf<FirClassLikeSymbol<*>>()
                val classIdToElementMap = mutableMapOf<ClassId, FirClassLikeDeclaration<*>>()
@@ -10,6 +10,8 @@ import com.intellij.testFramework.LightProjectDescriptor
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
@@ -41,7 +43,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : KotlinLightCodeInsightFix
                firModuleResolveState,
                FirResolvePhase.BODY_RESOLVE
            ) @Suppress("UNUSED_ANONYMOUS_PARAMETER") { firDeclaration ->
-                val firFile = firModuleResolveState.getBuiltFirFileOrNull(ktFile)!!
+                val firFile = firModuleResolveState.getOrBuildFirFile(ktFile)
                firFile.render(FirRenderer.RenderMode.WithResolvePhases)
            }
            val expectedFileName = testDataFile.name.replace(".kt", ".txt")
@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.createResolveStateForNoCaching
 import org.jetbrains.kotlin.idea.project.withLanguageVersionSettings
 import org.jetbrains.kotlin.idea.test.*
@@ -76,7 +76,7 @@ abstract class AbstractCompilerBasedTest : KotlinLightCodeInsightFixtureTestCase
            val resolveState = createResolveState(moduleInfo, module.languageVersionSettings)
            val allFirFiles = files.mapNotNull { (testFile, psiFile) ->
                if (psiFile !is KtFile) return@mapNotNull null
-                testFile to psiFile.getFirFile(resolveState)
+                testFile to psiFile.getOrBuildFirFile(resolveState)
            }.toMap()

            val diagnosticCheckerFilter = if (FirDiagnosticsDirectives.WITH_EXTENDED_CHECKERS in module.directives) {
@@ -17,7 +17,7 @@ import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getDiagnostics
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.createResolveStateForNoCaching
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.BeforeElementDiagnosticCollectionHandler
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistenceContextCollector
@@ -45,7 +45,7 @@ abstract class AbstractFirContextCollectionTest : KotlinLightCodeInsightFixtureT
        val fileStructure = resolveState.fileStructureCache.getFileStructure(ktFile, resolveState.rootModuleSession.cache)
        val allStructureElements = fileStructure.getAllStructureElements()
        handler.elementsToCheckContext = allStructureElements.map { it.getFirDeclaration() }
-        handler.firFile = ktFile.getFirFile(resolveState)
+        handler.firFile = ktFile.getOrBuildFirFile(resolveState)
        ktFile.getDiagnostics(resolveState, DiagnosticCheckerFilter.ONLY_COMMON_CHECKERS)
    }

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForResolveOnAir.getTowerContextProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.ResolutionParameters
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolutionMode
@@ -61,7 +61,7 @@ internal class KtFirCompletionCandidateChecker(
        nameExpression: KtSimpleNameExpression,
        possibleExplicitReceiver: KtExpression?,
    ): Boolean {
-        val file = originalFile.getFirFile(firResolveState)
+        val file = originalFile.getOrBuildFirFile(firResolveState)
        val explicitReceiverExpression = possibleExplicitReceiver?.getOrBuildFirOfType<FirExpression>(firResolveState)
        val resolver = SingleCandidateResolver(firResolveState.rootModuleSession, file)
        val implicitReceivers = getImplicitReceivers(originalFile, file, nameExpression)
@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclarationOfType
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
@@ -41,7 +41,7 @@ internal class KtFirSymbolProvider(
    }

    override fun getFileSymbol(psi: KtFile): KtFileSymbol = withValidityAssertion {
-        firSymbolBuilder.buildFileSymbol(psi.getFirFile(resolveState))
+        firSymbolBuilder.buildFileSymbol(psi.getOrBuildFirFile(resolveState))
    }

    override fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol = withValidityAssertion {