[FIR IDE] Implement on-air resolve of non-local declarations

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt

@@ -17,9 +17,11 @@ import com.intellij.openapi.editor.Document
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.completion.handlers.isTextAt
+import com.intellij.psi.util.PsiUtil
 import org.jetbrains.kotlin.idea.core.asFqNameWithRootPrefixIfNeeded
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.analyse
+import org.jetbrains.kotlin.idea.frontend.api.analyseInDependedAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.addImportToFile
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
@@ -415,7 +417,6 @@ private fun shortenReferences(targetFile: KtFile, textRange: TextRange) {
             collectPossibleReferenceShortenings(targetFile, textRange)
         }
     }
-
     shortenings.invokeShortening()
 }
 

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

@@ -6,22 +6,19 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 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.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getClosestAvailableParentContext
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
-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.util.containingKtFileIfAny
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.psi.KtDeclaration
@@ -30,24 +27,29 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 
 internal class FirModuleResolveStateDepended(
-    dependencyFirDeclaration: FirDeclaration,
-    originalFirFile: FirFile,
     private val originalState: FirModuleResolveStateImpl,
+    val towerProviderBuiltUponElement: FirTowerContextProvider,
+    private val ktToFirMapping: Map<KtElement, FirElement>,
 ) : FirModuleResolveState() {
 
     override val project: Project get() = originalState.project
     override val moduleInfo: IdeaModuleInfo get() = originalState.moduleInfo
     override val rootModuleSession get() = originalState.rootModuleSession
-    private val fileStructureCache = originalState.fileStructureCache
+    private val fileStructureCache get() = originalState.fileStructureCache
-    private val completionMapping = mutableMapOf<KtElement, FirElement>()
-    private val collector = FirTowerDataContextCollector()
 
     override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
         originalState.getSessionFor(moduleInfo)
 
     override fun getOrBuildFirFor(element: KtElement): FirElement {
         val psi = originalState.elementBuilder.getPsiAsFirElementSource(element)
-        synchronized(completionMapping) { completionMapping[psi] }?.let { return it }
+
+        //TODO It return invalid elements for elements with invalid code, but try to return the most closest ones
+        var currentElement: PsiElement = psi
+        while (currentElement !is KtFile) {
+            ktToFirMapping[currentElement]?.let { return it }
+            currentElement = currentElement.parent
+        }
+
         return originalState.elementBuilder.getOrBuildFirFor(
             element,
             originalState.firFileBuilder,
@@ -82,24 +84,4 @@ internal class FirModuleResolveStateDepended(
     @OptIn(InternalForInline::class)
     override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
         originalState.findSourceFirDeclaration(ktDeclaration)
-
-    override fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext? =
-        collector.getClosestAvailableParentContext(element) ?: originalState.getTowerDataContextForElement(element)
-
-    init {
-        originalState.firFileBuilder.runCustomResolveWithPCECheck(originalFirFile, rootModuleSession.cache) {
-            originalState.firLazyDeclarationResolver.runLazyResolveWithoutLock(
-                dependencyFirDeclaration,
-                rootModuleSession.cache,
-                originalFirFile,
-                originalFirFile.session.firIdeProvider,
-                fromPhase = dependencyFirDeclaration.resolvePhase,
-                toPhase = FirResolvePhase.BODY_RESOLVE,
-                towerDataContextCollector = collector,
-                checkPCE = true
-            )
-        }
-
-        synchronized(completionMapping) { dependencyFirDeclaration.accept(FirElementsRecorder(), completionMapping) }
-    }
 }

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

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.realPsi
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -23,8 +22,6 @@ 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.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getClosestAvailableParentContext
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -52,11 +49,8 @@ internal class FirModuleResolveStateImpl(
      * It can grow unbounded if you never edit the files in the opened project.
      *
      * It is a temporary solution until we can retrieve scopes for any fir element without re-resolving it.
-     *
-     * TODO Fix this when refactoring that separates resolving and scopes creation is done
      */
-    private val collector = FirTowerDataContextCollector()
+    val fileStructureCache = FileStructureCache(firFileBuilder, firLazyDeclarationResolver)
-    val fileStructureCache = FileStructureCache(firFileBuilder, firLazyDeclarationResolver, collector)
     val elementBuilder = FirElementBuilder()
     private val diagnosticsCollector = DiagnosticsCollector(fileStructureCache, rootModuleSession.cache)
 
@@ -111,7 +105,6 @@ internal class FirModuleResolveStateImpl(
                 cache,
                 FirResolvePhase.BODY_RESOLVE,
                 checkPCE = false, /*TODO*/
-                towerDataContextCollector = collector,
             )
         }
         val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(nonLocalFirForNamedDeclaration) { firDeclaration ->
@@ -135,15 +128,11 @@ internal class FirModuleResolveStateImpl(
             fileCache,
             toPhase,
             checkPCE = true,
-            towerDataContextCollector = collector,
+            towerDataContextCollector = null,
         )
         return declaration
     }
 
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
         cache.getContainerFirFile(declaration)
-
+}
-    override fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext? {
-        return collector.getClosestAvailableParentContext(element)
-    }
-}

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

@@ -36,6 +36,7 @@ internal class FirPhaseRunner {
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> implicitTypesResolveLock.withLock {
             runPhaseWithoutLock(firFile, phase, scopeSession)
         }
+        //TODO: Make correct sync for TYPES
         else -> {
             runPhaseWithoutLock(firFile, phase, scopeSession)
         }
@@ -51,6 +52,7 @@ internal class FirPhaseRunner {
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> implicitTypesResolveLock.withLock {
             runPhaseWithCustomResolveWithoutLock(resolve)
         }
+        //TODO: Make correct sync for TYPES
         else -> {
             runPhaseWithCustomResolveWithoutLock(resolve)
         }

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

@@ -7,51 +7,58 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 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.declarations.builder.*
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 
 object DeclarationCopyBuilder {
-    fun createDeclarationCopy(
+    fun <T : KtElement> createDeclarationCopy(
-        originalFirDeclaration: FirDeclaration,
-        copiedKtDeclaration: KtDeclaration,
         state: FirModuleResolveState,
+        replacement: RawFirReplacement<T>
     ): FirDeclaration {
-        return when (copiedKtDeclaration) {
+
+        val nonLocalDeclaration = LowLevelFirApiFacadeForDependentCopy.findEnclosingNonLocalDeclaration(replacement.from)
+            ?: error("Cannot find enclosing declaration for ${replacement.from.getElementTextInContext()}")
+
+        return when (nonLocalDeclaration) {
             is KtNamedFunction -> createFunctionCopy(
-                copiedKtDeclaration,
+                nonLocalDeclaration,
-                originalFirDeclaration as FirSimpleFunction,
+                state,
-                state
+                replacement
             )
             is KtProperty -> createPropertyCopy(
-                copiedKtDeclaration,
+                nonLocalDeclaration,
-                originalFirDeclaration as FirProperty,
+                state,
-                state
+                replacement
             )
             is KtClassOrObject -> createClassCopy(
-                copiedKtDeclaration,
+                nonLocalDeclaration,
-                originalFirDeclaration as FirRegularClass,
+                state,
-                state
+                replacement
             )
             is KtTypeAlias -> createTypeAliasCopy(
-                copiedKtDeclaration,
+                nonLocalDeclaration,
-                originalFirDeclaration as FirTypeAlias,
+                state,
-                state
+                replacement
             )
-            else -> error("Unsupported declaration ${copiedKtDeclaration::class.simpleName}")
+            else -> error("Unsupported declaration ${nonLocalDeclaration::class.simpleName}")
         }
     }
 
-    private fun createFunctionCopy(
+    private fun <T : KtElement> createFunctionCopy(
-        element: KtNamedFunction,
+        rootNonLocalDeclaration: KtNamedFunction,
-        originalFunction: FirSimpleFunction,
         state: FirModuleResolveState,
+        replacement: RawFirReplacement<T>,
     ): FirSimpleFunction {
-        val builtFunction = createCopy(element, originalFunction)
+
+        val originalFunction = rootNonLocalDeclaration.getOrBuildFirOfType<FirSimpleFunction>(state)
+        val builtFunction = createCopy(rootNonLocalDeclaration, originalFunction, replacement)
 
         // 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
@@ -63,27 +70,31 @@ object DeclarationCopyBuilder {
         }.apply { reassignAllReturnTargets(builtFunction) }
     }
 
-    private fun createClassCopy(
+    private fun <T : KtElement> createClassCopy(
-        copiedKtClassOrObject: KtClassOrObject,
+        rootNonLocalDeclaration: KtClassOrObject,
-        originalFirClass: FirRegularClass,
         state: FirModuleResolveState,
+        replacement: RawFirReplacement<T>,
     ): FirRegularClass {
-        val builtClass = createCopy(copiedKtClassOrObject, originalFirClass)
+        val originalFirClass = rootNonLocalDeclaration.getOrBuildFirOfType<FirRegularClass>(state)
+        val builtClass = createCopy(rootNonLocalDeclaration, originalFirClass, replacement)
 
         return buildRegularClassCopy(originalFirClass) {
             declarations.clear()
             declarations.addAll(builtClass.declarations)
             symbol = builtClass.symbol
             initDeclaration(originalFirClass, builtClass, state)
+            resolvePhase = minOf(originalFirClass.resolvePhase, FirResolvePhase.IMPORTS) //TODO move into initDeclaration?
         }
     }
 
-    private fun createTypeAliasCopy(
+    private fun <T : KtElement> createTypeAliasCopy(
-        copiedKtTypeAlias: KtTypeAlias,
+        rootNonLocalDeclaration: KtTypeAlias,
-        originalFirTypeAlias: FirTypeAlias,
         state: FirModuleResolveState,
+        replacement: RawFirReplacement<T>,
     ): FirTypeAlias {
-        val builtTypeAlias = createCopy(copiedKtTypeAlias, originalFirTypeAlias)
+
+        val originalFirTypeAlias = rootNonLocalDeclaration.getOrBuildFirOfType<FirTypeAlias>(state)
+        val builtTypeAlias = createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
 
         return buildTypeAliasCopy(originalFirTypeAlias) {
             expandedTypeRef = builtTypeAlias.expandedTypeRef
@@ -92,12 +103,13 @@ object DeclarationCopyBuilder {
         }
     }
 
-    private fun createPropertyCopy(
+    private fun <T : KtElement> createPropertyCopy(
-        element: KtProperty,
+        rootNonLocalDeclaration: KtProperty,
-        originalProperty: FirProperty,
+        state: FirModuleResolveState,
-        state: FirModuleResolveState
+        replacement: RawFirReplacement<T>,
     ): FirProperty {
-        val builtProperty = createCopy(element, originalProperty)
+        val originalProperty = rootNonLocalDeclaration.getOrBuildFirOfType<FirProperty>(state)
+        val builtProperty = createCopy(rootNonLocalDeclaration, originalProperty, replacement)
 
         val originalSetter = originalProperty.setter
         val builtSetter = builtProperty.setter
@@ -134,16 +146,18 @@ object DeclarationCopyBuilder {
         moduleData = state.rootModuleSession.moduleData
     }
 
-    internal inline fun <reified T : FirDeclaration> createCopy(
+    internal inline fun <reified D : FirDeclaration, T : KtElement> createCopy(
-        copiedKtDeclaration: KtDeclaration,
+        rootNonLocalDeclaration: KtDeclaration,
-        originalFirDeclaration: T,
+        originalFirDeclaration: D,
-    ): T {
+        replacement: RawFirReplacement<T>? = null,
-        return RawFirFragmentForLazyBodiesBuilder.build(
+    ): D {
+        return RawFirFragmentForLazyBodiesBuilder.buildWithReplacement(
             session = originalFirDeclaration.moduleData.session,
             baseScopeProvider = originalFirDeclaration.moduleData.session.firIdeProvider.kotlinScopeProvider,
-            designation = originalFirDeclaration.collectDesignation().fullDesignation,
+            designation = originalFirDeclaration.collectDesignation().path,
-            declaration = copiedKtDeclaration
+            declarationToBuild = rootNonLocalDeclaration,
-        ) as T
+            replacement = replacement,
+        ) as D
     }
 
     private fun FirFunction<*>.reassignAllReturnTargets(from: FirFunction<*>) {

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

@@ -10,7 +10,6 @@ import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -67,9 +66,4 @@ abstract class FirModuleResolveState {
     internal abstract fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 
     internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?
-
+}
-    abstract fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext?
-}
-
-fun FirModuleResolveState.getTowerDataContextUnsafe(element: KtElement): FirTowerDataContext =
-    getTowerDataContextForElement(element) ?: error("No context for ${element.getElementTextInContext()}")

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

@@ -7,9 +7,23 @@ 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.RawFirReplacement
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.realPsi
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateDepended
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextAllElementsCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.SingleElementTowerProvider
+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.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
@@ -25,7 +39,7 @@ object LowLevelFirApiFacadeForDependentCopy {
         else -> false
     }
 
-    private fun findEnclosingNonLocalDeclaration(position: KtElement): KtNamedDeclaration? =
+    fun findEnclosingNonLocalDeclaration(position: KtElement): KtNamedDeclaration? =
         position.parentsOfType<KtNamedDeclaration>().firstOrNull { ktDeclaration ->
             ktDeclaration.canBeEnclosingDeclaration()
         }
@@ -46,33 +60,153 @@ object LowLevelFirApiFacadeForDependentCopy {
         }
     }
 
+    private fun <T : KtElement> onAirResolveElement(
+        state: FirModuleResolveState,
+        place: T,
+        elementToResolve: T,
+    ): FirModuleResolveState {
+        require(state is FirModuleResolveStateImpl)
+        require(place.isPhysical)
+        require(!elementToResolve.isPhysical)
+
+        val collector = FirTowerDataContextAllElementsCollector()
+        val declaration = runBodyResolve(state, replacement = RawFirReplacement(place, elementToResolve))
+
+        val expressionLocator = object : FirVisitorVoid() {
+            var result: FirElement? = null
+                private set
+
+            override fun visitElement(element: FirElement) {
+                if (element.realPsi == elementToResolve) result = element
+                if (result != null) return
+                element.acceptChildren(this)
+            }
+        }
+        declaration.accept(expressionLocator)
+
+        val recordedMap = FirElementsRecorder.recordElementsFrom(declaration, FirElementsRecorder())
+        return FirModuleResolveStateDepended(state, collector, recordedMap)
+    }
+
+    fun onAirGetTowerContextProvider(
+        state: FirModuleResolveState,
+        place: KtElement,
+    ): FirTowerContextProvider {
+        require(state is FirModuleResolveStateImpl)
+        require(place.isPhysical)
+
+        return if (place is KtFile) {
+            onAirGetTowerContextForFile(state, place)
+        } else {
+            FirTowerDataContextAllElementsCollector().also {
+                runBodyResolve(state, collector = it, replacement = RawFirReplacement(place, place))
+            }
+        }
+    }
+
+    private fun onAirGetTowerContextForFile(
+        state: FirModuleResolveState,
+        file: KtFile,
+    ): FirTowerContextProvider {
+
+        val firFile = state.getOrBuildFirFor(file) as FirFile
+
+        val fileTransformer = object : FirBodyResolveTransformer(
+            session = firFile.declarationSiteSession,
+            phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+            implicitTypeOnly = true,
+            scopeSession = ScopeSession()
+        ) {
+            var result: FirTowerDataContext? = null
+                private set
+
+            override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration {
+                check(declaration is FirFile)
+                result = context.towerDataContext
+                return declaration
+            }
+        }
+        firFile.transform<FirFile, ResolutionMode>(fileTransformer, ResolutionMode.ContextDependent)
+
+        val fileContext = fileTransformer.result
+        check(fileContext != null) { "File context not found for physical file" }
+
+        return SingleElementTowerProvider(file, fileContext)
+    }
+
     fun getResolveStateForDependedCopy(
         originalState: FirModuleResolveState,
         originalKtFile: KtFile,
         dependencyKtElement: KtElement
     ): FirModuleResolveState {
         require(originalState is FirModuleResolveStateImpl)
+        require(dependencyKtElement !is KtFile) { "KtFile for dependency element not supported" }
+        require(!dependencyKtElement.isPhysical) { "Depended state should be build only for non-physical elements" }
 
         val dependencyNonLocalDeclaration = findEnclosingNonLocalDeclaration(dependencyKtElement)
             ?: error("Cannot find enclosing declaration for ${dependencyKtElement.getElementTextInContext()}")
 
-        val originalNonLocalDeclaration = locateDeclarationInFileByOffset(dependencyNonLocalDeclaration, originalKtFile)
+        val sameDeclarationInOriginalFile = locateDeclarationInFileByOffset(dependencyNonLocalDeclaration, originalKtFile)
             ?: error("Cannot find original function matching to ${dependencyNonLocalDeclaration.getElementTextInContext()} in $originalKtFile")
 
         recordOriginalDeclaration(
             targetDeclaration = dependencyNonLocalDeclaration,
-            originalDeclaration = originalNonLocalDeclaration
+            originalDeclaration = sameDeclarationInOriginalFile
         )
 
-        val originalFirDeclaration = originalNonLocalDeclaration.getOrBuildFirOfType<FirDeclaration>(originalState)
+        val collector = FirTowerDataContextAllElementsCollector()
-        val copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
+        val copiedFirDeclaration = runBodyResolve(
-            originalFirDeclaration = originalFirDeclaration,
+            originalState,
-            copiedKtDeclaration = dependencyNonLocalDeclaration,
+            collector = collector,
-            state = originalState
+            replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration)
         )
 
-        val originalFirFile = originalState.getFirFile(originalKtFile)
+        val recordedMap = FirElementsRecorder.recordElementsFrom(copiedFirDeclaration, FirElementsRecorder())
-
+        return FirModuleResolveStateDepended(originalState, collector, recordedMap)
-        return FirModuleResolveStateDepended(copiedFirDeclaration, originalFirFile, originalState)
     }
+
+    private fun <T : KtElement> runBodyResolve(
+        state: FirModuleResolveStateImpl,
+        replacement: RawFirReplacement<T>,
+        collector: FirTowerDataContextCollector? = null,
+    ): FirDeclaration {
+        val copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
+            state = state,
+            replacement = replacement,
+        )
+
+        val originalFirFile = state.getOrBuildFirFor(replacement.from.containingKtFile) as FirFile
+
+        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
+            )
+        }
+
+        return copiedFirDeclaration
+    }
+
+    private class TowerProviderForElementForState(private val state: FirModuleResolveState) : FirTowerContextProvider {
+        override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? {
+            return if (ktElement.isPhysical) {
+                onAirGetTowerContextProvider(state, ktElement).getClosestAvailableParentContext(ktElement)
+            } else {
+                require(state is FirModuleResolveStateDepended) {
+                    "Invalid resolve state ${this::class.simpleName} but have to be ${FirModuleResolveStateDepended::class.simpleName}"
+                }
+                state.towerProviderBuiltUponElement.getClosestAvailableParentContext(ktElement)
+            }
+        }
+    }
+
+    fun FirModuleResolveState.getTowerContextProvider(): FirTowerContextProvider =
+        TowerProviderForElementForState(this)
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.ThreadSafeMutableState
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirStatement
@@ -15,43 +14,59 @@ import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.psi.*
 
-@ThreadSafeMutableState
+interface FirTowerContextProvider {
-class FirTowerDataContextCollector {
+    fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext?
-    private val state: MutableMap<KtElement, FirTowerDataContext> = hashMapOf()
-
-    fun addFileContext(file: FirFile, context: FirTowerDataContext) {
-        val ktFile = file.psi as? KtFile ?: return
-        state[ktFile] = context
-    }
-
-    fun addStatementContext(statement: FirStatement, context: FirTowerDataContext) {
-        val closestStatementInBlock = statement.psi?.closestBlockLevelOrInitializerExpression() ?: return
-        state[closestStatementInBlock] = context
-    }
-
-    fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext) {
-        (declaration.psi as? KtElement)?.let { state[it] = context }
-    }
-
-    fun getContext(psi: KtElement): FirTowerDataContext? = state[psi]
 }
 
-fun FirTowerDataContextCollector.getClosestAvailableParentContext(element: KtElement): FirTowerDataContext? {
+interface FirTowerDataContextCollector {
-    var current: PsiElement? = element
+    fun addFileContext(file: FirFile, context: FirTowerDataContext)
-    while (current != null) {
+    fun addStatementContext(statement: FirStatement, context: FirTowerDataContext)
-        if (current is KtElement) {
+    fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext)
-            getContext(current)?.let { return it }
+}
-        }
+
-        if (current is KtDeclaration) {
+internal class SingleElementTowerProvider(
-            val originalDeclaration = current.originalDeclaration
+    private val element: KtElement,
-            originalDeclaration?.let { getContext(it) }?.let { return it }
+    private val context: FirTowerDataContext
-        }
+) : FirTowerContextProvider {
-        if (current is KtFile) {
+    override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? =
-            break
+        if (element == ktElement) context else null
-        }
+}
-        current = current.parent
+
+internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextCollector, FirTowerContextProvider {
+    private val elementsToContext: MutableMap<KtElement, FirTowerDataContext> = hashMapOf()
+
+    override fun addFileContext(file: FirFile, context: FirTowerDataContext) {
+        val ktFile = file.psi as? KtFile ?: return
+        elementsToContext[ktFile] = context
+    }
+
+    override fun addStatementContext(statement: FirStatement, context: FirTowerDataContext) {
+        val closestStatementInBlock = statement.psi?.closestBlockLevelOrInitializerExpression() ?: return
+        elementsToContext[closestStatementInBlock] = context
+    }
+
+    override fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext) {
+        val psi = declaration.psi as? KtElement ?: return
+        elementsToContext[psi] = context
+    }
+
+    override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? {
+        var current: PsiElement? = ktElement
+        while (current != null) {
+            if (current is KtElement) {
+                elementsToContext[current]?.let { return it }
+            }
+            if (current is KtDeclaration) {
+                val originalDeclaration = current.originalDeclaration
+                originalDeclaration?.let { elementsToContext[it] }?.let { return it }
+            }
+            if (current is KtFile) {
+                break
+            }
+            current = current.parent
+        }
+        return null
     }
-    return null
 }
 
 private tailrec fun PsiElement.closestBlockLevelOrInitializerExpression(): KtExpression? =

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

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -28,7 +27,6 @@ internal class FileStructure(
     private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
     private val firFileBuilder: FirFileBuilder,
     private val moduleFileCache: ModuleFileCache,
-    private val collector: FirTowerDataContextCollector
 ) {
     private val firIdeProvider = firFile.moduleData.session.firIdeProvider
 
@@ -50,7 +48,6 @@ internal class FileStructure(
                         cache = moduleFileCache,
                         firLazyDeclarationResolver = firLazyDeclarationResolver,
                         firIdeProvider = firIdeProvider,
-                        towerDataContextCollector = collector
                     )
                 }
                 else -> structureElement
@@ -112,7 +109,6 @@ internal class FileStructure(
             moduleFileCache,
             FirResolvePhase.BODY_RESOLVE,
             checkPCE = true,
-            towerDataContextCollector = collector
         )
         return moduleFileCache.firFileLockProvider.withReadLock(firFile) {
             FileElementFactory.createFileStructureElement(firDeclaration, declaration, firFile, moduleFileCache.firFileLockProvider)

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

@@ -18,12 +18,11 @@ import java.util.concurrent.ConcurrentHashMap
 internal class FileStructureCache(
     private val fileBuilder: FirFileBuilder,
     private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
-    private val collector: FirTowerDataContextCollector,
 ) {
     private val cache = ConcurrentHashMap<KtFile, FileStructure>()
 
     fun getFileStructure(ktFile: KtFile, moduleFileCache: ModuleFileCache): FileStructure = cache.computeIfAbsent(ktFile) {
         val firFile = fileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = false)
-        FileStructure(ktFile, firFile, firLazyDeclarationResolver, fileBuilder, moduleFileCache, collector)
+        FileStructure(ktFile, firFile, firLazyDeclarationResolver, fileBuilder, moduleFileCache)
     }
 }

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

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationCopyBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnostics
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.LockProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
@@ -45,7 +44,6 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S : Abstr
         cache: ModuleFileCache,
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         firIdeProvider: FirIdeProvider,
-        towerDataContextCollector: FirTowerDataContextCollector,
     ): ReanalyzableStructureElement<KT, S>
 
     fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
@@ -76,17 +74,15 @@ internal class ReanalyzableFunctionStructureElement(
         cache: ModuleFileCache,
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         firIdeProvider: FirIdeProvider,
-        towerDataContextCollector: FirTowerDataContextCollector,
     ): ReanalyzableFunctionStructureElement {
         val originalFunction = firSymbol.fir as FirSimpleFunction
-        val newFunction = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalFunction)
+        val newFunction = DeclarationCopyBuilder.createCopy<FirSimpleFunction, KtElement>(newKtDeclaration, originalFunction)
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalFunction, newFunction) {
             firLazyDeclarationResolver.lazyResolveDeclaration(
                 newFunction,
                 cache,
                 FirResolvePhase.BODY_RESOLVE,
-                towerDataContextCollector,
                 checkPCE = true,
                 reresolveFile = true,
             )
@@ -118,17 +114,15 @@ internal class ReanalyzablePropertyStructureElement(
         cache: ModuleFileCache,
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         firIdeProvider: FirIdeProvider,
-        towerDataContextCollector: FirTowerDataContextCollector,
     ): ReanalyzablePropertyStructureElement {
         val originalProperty = firSymbol.fir
-        val newProperty = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalProperty)
+        val newProperty = DeclarationCopyBuilder.createCopy<FirProperty, KtElement>(newKtDeclaration, originalProperty)
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalProperty, newProperty) {
             firLazyDeclarationResolver.lazyResolveDeclaration(
                 newProperty,
                 cache,
                 FirResolvePhase.BODY_RESOLVE,
-                towerDataContextCollector,
                 checkPCE = true,
                 reresolveFile = true,
             )

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

@@ -14,14 +14,18 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 
 internal object FirLazyBodiesCalculator {
-    fun calculateLazyBodiesInside(element: FirElement, designation: List<FirDeclaration>) {
+    fun calculateLazyBodiesInside(element: FirElement, designation: FirDeclarationDesignation) {
-        element.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, designation.toMutableList())
+        element.transform<FirElement, MutableList<FirDeclaration>>(
+            FirLazyBodiesCalculatorTransformer,
+            designation.fullDesignation.toMutableList()
+        )
     }
 
     fun calculateLazyBodiesIfPhaseRequires(firFile: FirFile, phase: FirResolvePhase) {
@@ -36,7 +40,7 @@ internal object FirLazyBodiesCalculator {
             session = simpleFunction.moduleData.session,
             baseScopeProvider = simpleFunction.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = designation,
-            declaration = simpleFunction.psi as KtNamedFunction
+            rootNonLocalDeclaration = simpleFunction.psi as KtNamedFunction
         ) as FirSimpleFunction
         simpleFunction.apply {
             replaceBody(newFunction.body)
@@ -52,7 +56,7 @@ internal object FirLazyBodiesCalculator {
             session = secondaryConstructor.moduleData.session,
             baseScopeProvider = secondaryConstructor.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = designation,
-            declaration = secondaryConstructor.psi as KtSecondaryConstructor
+            rootNonLocalDeclaration = secondaryConstructor.psi as KtSecondaryConstructor
         ) as FirSimpleFunction
 
         secondaryConstructor.apply {
@@ -67,7 +71,7 @@ internal object FirLazyBodiesCalculator {
             session = firProperty.moduleData.session,
             baseScopeProvider = firProperty.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = designation,
-            declaration = firProperty.psi as KtProperty
+            rootNonLocalDeclaration = firProperty.psi as KtProperty
         ) as FirProperty
 
         firProperty.getter?.takeIf { it.body is FirLazyBlock }?.let { getter ->

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
-import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
@@ -14,6 +13,8 @@ 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.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -107,7 +108,7 @@ internal class FirLazyDeclarationResolver(
         }
     }
 
-    private fun calculateLazyBodies(firDeclaration: FirDeclaration, designation: List<FirDeclaration>) {
+    private fun calculateLazyBodies(firDeclaration: FirDeclaration, designation: FirDeclarationDesignation) {
         FirLazyBodiesCalculator.calculateLazyBodiesInside(firDeclaration, designation)
     }
 
@@ -120,9 +121,10 @@ internal class FirLazyDeclarationResolver(
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector? = null,
         checkPCE: Boolean,
+        lastNonLazyPhase: FirResolvePhase = LAST_NON_LAZY_PHASE
     ) {
         if (fromPhase >= toPhase) return
-        val nonLazyPhase = minOf(toPhase, LAST_NON_LAZY_PHASE)
+        val nonLazyPhase = minOf(toPhase, lastNonLazyPhase)
 
         val scopeSession = ScopeSession()
         if (fromPhase < nonLazyPhase) {
@@ -138,7 +140,7 @@ internal class FirLazyDeclarationResolver(
         resolveFileAnnotations(containerFirFile, moduleFileCache, scopeSession)
 
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
-        val designation = nonLocalDeclarationToResolve.getDesignation(containerFirFile, provider, moduleFileCache)
+        val designation = nonLocalDeclarationToResolve.collectDesignation()
 
         executeWithoutPCE {
             calculateLazyBodies(firDeclarationToResolve, designation)
@@ -165,9 +167,9 @@ internal class FirLazyDeclarationResolver(
         phase: FirResolvePhase,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?,
-        designation: List<FirDeclaration>
+        designation: FirDeclarationDesignation
     ) {
-        if (designation.all { it.resolvePhase >= phase }) {
+        if (designation.fullDesignation.all { it.resolvePhase >= phase }) {
             return
         }
 
@@ -179,29 +181,49 @@ internal class FirLazyDeclarationResolver(
         )
 
         firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
-            containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextIndependent)
+            transformer.transformDeclaration()
         }
     }
 
     private fun FirResolvePhase.createLazyTransformer(
-        designation: List<FirDeclaration>,
+        designation: FirDeclarationDesignation,
         containerFirFile: FirFile,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?
-    ) = when (this) {
+    ): FirLazyTransformerForIDE = when (this) {
+        FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
+            designation,
+            containerFirFile,
+            containerFirFile.moduleData.session,
+            scopeSession
+        )
+        FirResolvePhase.SEALED_CLASS_INHERITORS -> FirLazyTransformerForIDE.EMPTY
+        FirResolvePhase.TYPES -> FirDesignatedTypeResolverTransformerForIDE(
+            containerFirFile,
+            designation,
+            containerFirFile.moduleData.session,
+            scopeSession,
+        )
+        FirResolvePhase.STATUS -> FirDesignatedStatusResolveTransformerForIDE(
+            designation,
+            containerFirFile.moduleData.session,
+            scopeSession
+        )
         FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
-            FirDesignation(designation),
+            containerFirFile,
+            designation,
             containerFirFile.moduleData.session,
             scopeSession,
         )
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
-            FirDesignation(designation),
+            containerFirFile,
+            designation,
             containerFirFile.moduleData.session,
-            scopeSession,
+            scopeSession
-            towerDataContextCollector,
         )
         FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
-            FirDesignation(designation),
+            containerFirFile,
+            designation,
             containerFirFile.moduleData.session,
             scopeSession,
             towerDataContextCollector
@@ -209,27 +231,6 @@ internal class FirLazyDeclarationResolver(
         else -> error("Non-lazy phase $this")
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
-    private fun FirDeclaration.getDesignation(
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        moduleFileCache: ModuleFileCache
-    ): List<FirDeclaration> = buildList {
-        if (this !is FirFile) {
-            val ktDeclaration = ktDeclaration
-            ktDeclaration.parentsOfType<KtClassOrObject>(withSelf = true)
-                .filter { it !is KtEnumEntry }
-                .map { it.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache, containerFirFile) }
-                .toList()
-                .asReversed()
-                .let(::addAll)
-            if (this@getDesignation is FirCallableDeclaration<*> || this@getDesignation is FirTypeAlias) {
-                add(this@getDesignation)
-            }
-        }
-    }
-
-
     private fun FirDeclaration.getNonLocalDeclarationToResolve(provider: FirProvider, moduleFileCache: ModuleFileCache): FirDeclaration {
         if (this is FirFile) return this
         val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")

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

@@ -77,6 +77,7 @@ private constructor(
         check(mode == AnalysisSessionMode.REGULAR) {
             "Cannot create context-dependent copy of KtAnalysis session from a context dependent one"
         }
+        require(!dependencyKtElement.isPhysical) { "Depended context should be build only for non-physical elements" }
 
         val contextResolveState = LowLevelFirApiFacadeForDependentCopy.getResolveStateForDependedCopy(
             originalState = firResolveState,

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

@@ -8,9 +8,9 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 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.LowLevelFirApiFacadeForDependentCopy.getTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getTowerDataContextUnsafe
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.ResolutionParameters
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolutionMode
 import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolver
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -84,7 +85,9 @@ internal class KtFirCompletionCandidateChecker(
         firFile: FirFile,
         fakeNameExpression: KtSimpleNameExpression
     ): Sequence<ImplicitReceiverValue<*>?> {
-        val towerDataContext = analysisSession.firResolveState.getTowerDataContextUnsafe(fakeNameExpression)
+        val towerDataContext = analysisSession.firResolveState.getTowerContextProvider()
+            .getClosestAvailableParentContext(fakeNameExpression)
+            ?: error("Cannot find enclosing declaration for ${fakeNameExpression.getElementTextInContext()}")
 
         return sequence {
             yield(null) // otherwise explicit receiver won't be checked when there are no implicit receivers in completion position

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

@@ -45,7 +45,10 @@ import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.lowerBoundIfFlexible
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForDependentCopy
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForDependentCopy.getTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.components.KtReferenceShortener
 import org.jetbrains.kotlin.idea.frontend.api.components.ShortenCommand
@@ -72,7 +75,10 @@ internal class KtFirReferenceShortener(
 
         val firDeclaration = declarationToVisit.getOrBuildFir(firResolveState)
 
-        val collector = ElementsToShortenCollector(context)
+        val towerContext =
+            LowLevelFirApiFacadeForDependentCopy.onAirGetTowerContextProvider(firResolveState, declarationToVisit)
+
+        val collector = ElementsToShortenCollector(context, towerContext)
         firDeclaration.accept(collector)
 
         return ShortenCommandImpl(
@@ -139,9 +145,8 @@ private class FirShorteningContext(val firResolveState: FirModuleResolveState) {
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    fun findScopesAtPosition(position: KtElement, newImports: List<FqName>): List<FirScope>? {
+    fun findScopesAtPosition(position: KtElement, newImports: List<FqName>, towerContextProvider: FirTowerContextProvider): List<FirScope>? {
-        val towerDataContext = firResolveState.getTowerDataContextForElement(position) ?: return null
+        val towerDataContext = towerContextProvider.getClosestAvailableParentContext(position) ?: return null
-
         val result = buildList<FirScope> {
             addAll(towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope })
             addIfNotNull(createFakeImportingScope(newImports))
@@ -181,7 +186,8 @@ private sealed class ElementToShorten
 private class ShortenType(val element: KtUserType, val nameToImport: FqName? = null) : ElementToShorten()
 private class ShortenQualifier(val element: KtDotQualifiedExpression, val nameToImport: FqName? = null) : ElementToShorten()
 
-private class ElementsToShortenCollector(private val shorteningContext: FirShorteningContext) : FirVisitorVoid() {
+private class ElementsToShortenCollector(private val shorteningContext: FirShorteningContext, private val towerContextProvider: FirTowerContextProvider) :
+    FirVisitorVoid() {
     val namesToImport: MutableList<FqName> = mutableListOf()
     val typesToShorten: MutableList<KtUserType> = mutableListOf()
     val qualifiersToShorten: MutableList<KtDotQualifiedExpression> = mutableListOf()
@@ -230,7 +236,7 @@ private class ElementsToShortenCollector(private val shorteningContext: FirShort
         val allClassIds = wholeClassifierId.outerClassesWithSelf
         val allTypeElements = wholeTypeElement.qualifiersWithSelf
 
-        val positionScopes = shorteningContext.findScopesAtPosition(wholeTypeElement, namesToImport) ?: return null
+        val positionScopes = shorteningContext.findScopesAtPosition(wholeTypeElement, namesToImport, towerContextProvider) ?: return null
 
         for ((classId, typeElement) in allClassIds.zip(allTypeElements)) {
             // if qualifier is null, then this type have no package and thus cannot be shortened
@@ -280,7 +286,7 @@ private class ElementsToShortenCollector(private val shorteningContext: FirShort
         wholeClassQualifier: ClassId,
         wholeQualifierElement: KtDotQualifiedExpression
     ): ShortenQualifier? {
-        val positionScopes = shorteningContext.findScopesAtPosition(wholeQualifierElement, namesToImport) ?: return null
+        val positionScopes = shorteningContext.findScopesAtPosition(wholeQualifierElement, namesToImport, towerContextProvider) ?: return null
 
         val allClassIds = wholeClassQualifier.outerClassesWithSelf
         val allQualifiers = wholeQualifierElement.qualifiersWithSelf
@@ -313,7 +319,7 @@ private class ElementsToShortenCollector(private val shorteningContext: FirShort
 
         val propertyId = (resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*>)?.callableId ?: return
 
-        val scopes = shorteningContext.findScopesAtPosition(qualifiedProperty, namesToImport) ?: return
+        val scopes = shorteningContext.findScopesAtPosition(qualifiedProperty, namesToImport, towerContextProvider) ?: return
         val singleAvailableProperty = shorteningContext.findPropertiesInScopes(scopes, propertyId.callableName)
 
         val propertyToShorten = when {
@@ -334,7 +340,7 @@ private class ElementsToShortenCollector(private val shorteningContext: FirShort
         val calleeReference = functionCall.calleeReference
         val callableId = findUnambiguousReferencedCallableId(calleeReference) ?: return
 
-        val scopes = shorteningContext.findScopesAtPosition(callExpression, namesToImport) ?: return
+        val scopes = shorteningContext.findScopesAtPosition(callExpression, namesToImport, towerContextProvider) ?: return
         val availableCallables = shorteningContext.findFunctionsInScopes(scopes, callableId.callableName)
 
         val callToShorten = when {

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

@@ -17,7 +17,8 @@ import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.scopes.*
 import org.jetbrains.kotlin.fir.scopes.impl.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getTowerDataContextUnsafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForDependentCopy.getTowerContextProvider
+import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.components.KtImplicitReceiver
 import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeContext
@@ -33,9 +34,9 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFileSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithMembers
-import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import java.util.*
@@ -142,7 +143,10 @@ internal class KtFirScopeProvider(
         originalFile: KtFile,
         positionInFakeFile: KtElement
     ): KtScopeContext = withValidityAssertion {
-        val towerDataContext = analysisSession.firResolveState.getTowerDataContextUnsafe(positionInFakeFile)
+
+        val towerDataContext =
+            analysisSession.firResolveState.getTowerContextProvider().getClosestAvailableParentContext(positionInFakeFile)
+                ?: error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
 
         val implicitReceivers = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver }.distinct()
         val implicitKtReceivers = implicitReceivers.map { receiver ->