[FIR IDE] LowLevel api refactoring

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -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
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -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 =

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/LocalClassesResolution.kt

@@ -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) {

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

@@ -58,13 +58,13 @@ internal class FirModuleResolveStateDepended(
         )
     }
 
-    override fun getFirFile(ktFile: KtFile): FirFile =
+    override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
-        originalState.getFirFile(ktFile)
+        originalState.getOrBuildFirFile(ktFile)
 
-    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
+    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
-        originalState.resolvedFirToPhase(declaration, toPhase)
+        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 }

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

@@ -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)
 }

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

@@ -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 createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
-        return builtTypeAlias
-
-//        return buildTypeAliasCopy(originalFirTypeAlias) {
-//            expandedTypeRef = builtTypeAlias.expandedTypeRef
-//            symbol = builtTypeAlias.symbol
-//            initDeclaration(originalFirTypeAlias, builtTypeAlias, state)
-//        }
     }
 
     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

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

@@ -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
+            originalDeclaration.origin == FirDeclarationOrigin.Source && session is FirIdeSourcesSession -> {
-                    && 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 <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D
-
-    internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?
 }

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

@@ -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 =
+fun KtFile.getOrBuildFirFile(resolveState: FirModuleResolveState): FirFile =
-    resolveState.getFirFile(this)
+    resolveState.getOrBuildFirFile(this)
 
 class InvalidFirElementTypeException(
     ktElement: KtElement,

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

@@ -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,
+            isOnAirResolve = true,
-            useFirProviderInterceptor = 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)
+        val originalFirFile = state.getOrBuildFirFile(nonLocalDeclaration.containingKtFile)
-            ?: error("Original fir file should be already built")
 
-        val firProviderInterceptor =
+        state.firLazyDeclarationResolver.lazyDesignatedResolveDeclaration(
-            if (useFirProviderInterceptor) FirProviderInterceptorForIDE.createForFirElement(
+            firDeclarationToResolve = copiedFirDeclaration,
-                session = originalFirFile.declarationSiteSession,
+            moduleFileCache = state.rootModuleSession.cache,
-                firFile = originalFirFile,
+            containerFirFile = originalFirFile,
-                element = copiedFirDeclaration
+            provider = originalFirFile.declarationSiteSession.firIdeProvider,
-            ) else null
+            toPhase = FirResolvePhase.BODY_RESOLVE,
-
+            checkPCE = true,
-        state.firFileBuilder.runCustomResolveWithPCECheck(originalFirFile, state.rootModuleSession.cache) {
+            isOnAirResolve = isOnAirResolve,
-            state.firLazyDeclarationResolver.runLazyResolveWithoutLock(
+            towerDataContextCollector = collector,
-                copiedFirDeclaration,
+        )
-                state.rootModuleSession.cache,
-                originalFirFile,
-                originalFirFile.declarationSiteSession.firIdeProvider,
-                fromPhase = copiedFirDeclaration.resolvePhase,
-                toPhase = FirResolvePhase.BODY_RESOLVE,
-                towerDataContextCollector = collector,
-                checkPCE = true,
-                lastNonLazyPhase = FirResolvePhase.IMPORTS,
-                firProviderInterceptor = firProviderInterceptor,
-            )
-        }
 
         return copiedFirDeclaration
     }

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

@@ -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) {
+    fun calculateLazyBodiesInside(designation: FirDeclarationDesignation) {
-        element.transform<FirElement, MutableList<FirDeclaration>>(
+        designation.declaration.transform<FirElement, MutableList<FirDeclaration>>(
             FirLazyBodiesCalculatorTransformer,
             designation.toSequence(includeTarget = true).toMutableList()
         )

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

@@ -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,
-    ) {
+    ) = firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
-        check(firFile.resolvePhase >= FirResolvePhase.IMPORTS)
+        resolveFileAnnotationsWithoutLock(firFile, scopeSession)
-        firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
-            val transformer = FirFileAnnotationsResolveTransformer(firFile.moduleData.session, scopeSession)
-            firFile.accept(transformer, ResolutionMode.ContextDependent)
-        }
     }
 
-    fun lazyResolveDeclaration(
+    /**
-        declaration: FirDeclaration,
+     * Fully resolve file annotations (not synchronized)
-        moduleFileCache: ModuleFileCache,
+     * @see resolveFileAnnotations synchronized version
-        toPhase: FirResolvePhase,
+     */
-        towerDataContextCollector: FirTowerDataContextCollector? = null,
+    private fun resolveFileAnnotationsWithoutLock(
-        checkPCE: Boolean = false,
+        firFile: FirFile,
-        reresolveFile: Boolean = false,
+        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
+            return ktContainingResolvableDeclaration.findSourceNonLocalFirDeclaration(
-                .findSourceNonLocalFirDeclaration(firFileBuilder, declaration.moduleData.session.symbolProvider, moduleFileCache)
+                firFileBuilder = firFileBuilder,
-            return lazyResolveDeclaration(containingProperty, moduleFileCache, toPhase, towerDataContextCollector)
+                firSymbolProvider = declaration.moduleData.session.symbolProvider,
+                moduleFileCache = moduleFileCache
+            )
         }
 
-        val firFile = declaration.getContainingFile()
+        return declaration
-            ?: error("FirFile was not found for\n${declaration.render()}")
+    }
+
+    /**
+     * 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,
+                    firDeclarationToResolve = declaration,
-                    moduleFileCache,
+                    moduleFileCache = moduleFileCache,
-                    firFile,
+                    containerFirFile = firFile,
-                    provider,
+                    provider = provider,
-                    fromPhase,
+                    fromPhase = fromPhase,
-                    toPhase,
+                    toPhase = toPhase,
-                    towerDataContextCollector,
+                    checkPCE = true,
-                    checkPCE = true
                 )
             }
         } else {
             firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
                 executeWithoutPCE {
                     runLazyResolveWithoutLock(
-                        declaration,
+                        firDeclarationToResolve = declaration,
-                        moduleFileCache,
+                        moduleFileCache = moduleFileCache,
-                        firFile,
+                        containerFirFile = firFile,
-                        provider,
+                        provider = provider,
-                        fromPhase,
+                        fromPhase = fromPhase,
-                        toPhase,
+                        toPhase = toPhase,
-                        towerDataContextCollector,
+                        checkPCE = false,
-                        checkPCE = false
                     )
                 }
             }
         }
     }
 
-    private fun calculateLazyBodies(firDeclaration: FirDeclaration, designation: FirDeclarationDesignation) {
+    private fun createLazyBodiesCalculator(designation: FirDeclarationDesignation): (FirResolvePhase) -> Unit {
-        FirLazyBodiesCalculator.calculateLazyBodiesInside(firDeclaration, designation)
+        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 {
+        val lazyBodiesCalculator = createLazyBodiesCalculator(designation)
-            calculateLazyBodies(firDeclarationToResolve, 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,
+                phase = currentPhase,
-                scopeSession,
+                scopeSession = scopeSession,
-                towerDataContextCollector,
+                designation = designation,
-                firProviderInterceptor,
+                towerDataContextCollector = towerDataContextCollector,
-                designation,
+                firProviderInterceptor = firProviderInterceptor
             )
         }
     }
@@ -169,9 +308,9 @@ internal class FirLazyDeclarationResolver(
     private fun runLazyResolvePhase(
         phase: FirResolvePhase,
         scopeSession: ScopeSession,
+        designation: FirDeclarationDesignationWithFile,
         towerDataContextCollector: FirTowerDataContextCollector?,
-        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
+        firProviderInterceptor: FirProviderInterceptor?
-        designation: FirDeclarationDesignationWithFile
     ) {
         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,

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

@@ -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,

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

@@ -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)
         }
 

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

@@ -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,

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

@@ -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,

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

@@ -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")
-        }
     }
 }

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

@@ -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<*>>()

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

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

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/compiler/based/AbstractCompilerBasedTest.kt

@@ -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) {

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

@@ -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)
     }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt

@@ -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)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt

@@ -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 {