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

2021-05-04 17:03:46 +02:00
parent 779a06e86f
commit 261098a122
17 changed files with 348 additions and 209 deletions
@@ -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()
 }

@@ -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.getClosestAvailableParentContext
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 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 completionMapping = mutableMapOf<KtElement, FirElement>()
-    private val collector = FirTowerDataContextCollector()
+    private val fileStructureCache get() = originalState.fileStructureCache

    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) }
-    }
 }
@@ -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, collector)
+    val fileStructureCache = FileStructureCache(firFileBuilder, firLazyDeclarationResolver)
    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)
-    }
-}
+}
@@ -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)
        }
@@ -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(
-        originalFirDeclaration: FirDeclaration,
-        copiedKtDeclaration: KtDeclaration,
+    fun <T : KtElement> createDeclarationCopy(
        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,
-                originalFirDeclaration as FirSimpleFunction,
-                state
+                nonLocalDeclaration,
+                state,
+                replacement
            )
            is KtProperty -> createPropertyCopy(
-                copiedKtDeclaration,
-                originalFirDeclaration as FirProperty,
-                state
+                nonLocalDeclaration,
+                state,
+                replacement
            )
            is KtClassOrObject -> createClassCopy(
-                copiedKtDeclaration,
-                originalFirDeclaration as FirRegularClass,
-                state
+                nonLocalDeclaration,
+                state,
+                replacement
            )
            is KtTypeAlias -> createTypeAliasCopy(
-                copiedKtDeclaration,
-                originalFirDeclaration as FirTypeAlias,
-                state
+                nonLocalDeclaration,
+                state,
+                replacement
            )
-            else -> error("Unsupported declaration ${copiedKtDeclaration::class.simpleName}")
+            else -> error("Unsupported declaration ${nonLocalDeclaration::class.simpleName}")
        }
    }

-    private fun createFunctionCopy(
-        element: KtNamedFunction,
-        originalFunction: FirSimpleFunction,
+    private fun <T : KtElement> createFunctionCopy(
+        rootNonLocalDeclaration: KtNamedFunction,
        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(
-        copiedKtClassOrObject: KtClassOrObject,
-        originalFirClass: FirRegularClass,
+    private fun <T : KtElement> createClassCopy(
+        rootNonLocalDeclaration: KtClassOrObject,
        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(
-        copiedKtTypeAlias: KtTypeAlias,
-        originalFirTypeAlias: FirTypeAlias,
+    private fun <T : KtElement> createTypeAliasCopy(
+        rootNonLocalDeclaration: KtTypeAlias,
        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(
-        element: KtProperty,
-        originalProperty: FirProperty,
-        state: FirModuleResolveState
+    private fun <T : KtElement> createPropertyCopy(
+        rootNonLocalDeclaration: KtProperty,
+        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(
-        copiedKtDeclaration: KtDeclaration,
-        originalFirDeclaration: T,
-    ): T {
-        return RawFirFragmentForLazyBodiesBuilder.build(
+    internal inline fun <reified D : FirDeclaration, T : KtElement> createCopy(
+        rootNonLocalDeclaration: KtDeclaration,
+        originalFirDeclaration: D,
+        replacement: RawFirReplacement<T>? = null,
+    ): D {
+        return RawFirFragmentForLazyBodiesBuilder.buildWithReplacement(
            session = originalFirDeclaration.moduleData.session,
            baseScopeProvider = originalFirDeclaration.moduleData.session.firIdeProvider.kotlinScopeProvider,
-            designation = originalFirDeclaration.collectDesignation().fullDesignation,
-            declaration = copiedKtDeclaration
-        ) as T
+            designation = originalFirDeclaration.collectDesignation().path,
+            declarationToBuild = rootNonLocalDeclaration,
+            replacement = replacement,
+        ) as D
    }

    private fun FirFunction<*>.reassignAllReturnTargets(from: FirFunction<*>) {
@@ -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()}")
+}
@@ -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 copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
-            originalFirDeclaration = originalFirDeclaration,
-            copiedKtDeclaration = dependencyNonLocalDeclaration,
-            state = originalState
+        val collector = FirTowerDataContextAllElementsCollector()
+        val copiedFirDeclaration = runBodyResolve(
+            originalState,
+            collector = collector,
+            replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration)
        )

-        val originalFirFile = originalState.getFirFile(originalKtFile)
-
-        return FirModuleResolveStateDepended(copiedFirDeclaration, originalFirFile, originalState)
+        val recordedMap = FirElementsRecorder.recordElementsFrom(copiedFirDeclaration, FirElementsRecorder())
+        return FirModuleResolveStateDepended(originalState, collector, recordedMap)
    }
+
+    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)
 }
@@ -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
-class FirTowerDataContextCollector {
-    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]
+interface FirTowerContextProvider {
+    fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext?
 }

-fun FirTowerDataContextCollector.getClosestAvailableParentContext(element: KtElement): FirTowerDataContext? {
-    var current: PsiElement? = element
-    while (current != null) {
-        if (current is KtElement) {
-            getContext(current)?.let { return it }
-        }
-        if (current is KtDeclaration) {
-            val originalDeclaration = current.originalDeclaration
-            originalDeclaration?.let { getContext(it) }?.let { return it }
-        }
-        if (current is KtFile) {
-            break
-        }
-        current = current.parent
+interface FirTowerDataContextCollector {
+    fun addFileContext(file: FirFile, context: FirTowerDataContext)
+    fun addStatementContext(statement: FirStatement, context: FirTowerDataContext)
+    fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext)
+}
+
+internal class SingleElementTowerProvider(
+    private val element: KtElement,
+    private val context: FirTowerDataContext
+) : FirTowerContextProvider {
+    override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? =
+        if (element == ktElement) context else null
+}
+
+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? =
@@ -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)
@@ -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)
    }
 }
@@ -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,
            )
@@ -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>) {
-        element.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, designation.toMutableList())
+    fun calculateLazyBodiesInside(element: FirElement, designation: FirDeclarationDesignation) {
+        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 ->
@@ -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,
-            towerDataContextCollector,
+            scopeSession
        )
        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")
@@ -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,
@@ -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
@@ -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>? {
-        val towerDataContext = firResolveState.getTowerDataContextForElement(position) ?: return null
-
+    fun findScopesAtPosition(position: KtElement, newImports: List<FqName>, towerContextProvider: FirTowerContextProvider): List<FirScope>? {
+        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 {
@@ -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 ->