[FIR IDE] Refactor low level resolve state for completion

Rename refactorings
Remove redundant methods and implementations
Increase code-logic locality
Removed redundant entities

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.completion
 
 import com.intellij.codeInsight.completion.*
-import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PsiJavaPatterns
 import com.intellij.util.ProcessingContext
@@ -26,7 +25,6 @@ import org.jetbrains.kotlin.idea.frontend.api.scopes.KtCompositeScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
-import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithVisibility
 import org.jetbrains.kotlin.idea.frontend.api.types.KtClassType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/context/FirPositionCompletionContext.kt

@@ -8,16 +8,13 @@ package org.jetbrains.kotlin.idea.completion.context
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.analyse
-import org.jetbrains.kotlin.idea.frontend.api.analyseInFakeAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.analyseInDependedAnalysisSession
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtLabelReferenceExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
-import kotlin.contracts.ExperimentalContracts
-import kotlin.contracts.InvocationKind
-import kotlin.contracts.contract
 
 internal sealed class FirPositionCompletionContext {
     abstract val position: PsiElement
@@ -51,7 +48,7 @@ internal object FirPositionCompletionContextDetector {
         action: KtAnalysisSession.() -> Unit
     ) {
         return when (positionContext) {
-            is FirNameReferencePositionContext -> analyseInFakeAnalysisSession(
+            is FirNameReferencePositionContext -> analyseInDependedAnalysisSession(
                 basicContext.originalKtFile,
                 positionContext.nameExpression,
                 action

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -47,7 +47,7 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
 
     override val analysisSession: KtAnalysisSession get() = this
 
-    abstract fun createContextDependentCopy(originalKtFile: KtFile, fakeKtElement: KtElement): KtAnalysisSession
+    abstract fun createContextDependentCopy(originalKtFile: KtFile, dependencyKtElement: KtElement): KtAnalysisSession
 
     internal val smartCastProvider: KtSmartCastProvider get() = smartCastProviderImpl
     protected abstract val smartCastProviderImpl: KtSmartCastProvider

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt

@@ -36,10 +36,14 @@ abstract class KtAnalysisSessionProvider : Disposable {
     abstract fun getAnalysisSession(contextElement: KtElement, factory: ValidityTokenFactory): KtAnalysisSession
 
     @InvalidWayOfUsingAnalysisSession
-    inline fun <R> analyseInFakeAnalysisSession(originalFile: KtFile, fakeExpresion: KtElement, action: KtAnalysisSession.() -> R): R {
-        val fakeAnalysisSession = getAnalysisSession(originalFile, ReadActionConfinementValidityTokenFactory)
-            .createContextDependentCopy(originalFile, fakeExpresion)
-        return analyse(fakeAnalysisSession, ReadActionConfinementValidityTokenFactory, action)
+    inline fun <R> analyseInDependedAnalysisSession(
+        originalFile: KtFile,
+        dependencyExpression: KtElement,
+        action: KtAnalysisSession.() -> R
+    ): R {
+        val dependedAnalysisSession = getAnalysisSession(originalFile, ReadActionConfinementValidityTokenFactory)
+            .createContextDependentCopy(originalFile, dependencyExpression)
+        return analyse(dependedAnalysisSession, ReadActionConfinementValidityTokenFactory, action)
     }
 
     @InvalidWayOfUsingAnalysisSession
@@ -89,8 +93,12 @@ inline fun <R> analyseWithCustomToken(
     contextElement.project.service<KtAnalysisSessionProvider>().analyse(contextElement, tokenFactory, action)
 
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
-inline fun <R> analyseInFakeAnalysisSession(originalFile: KtFile, fakeExpresion: KtElement, action: KtAnalysisSession.() -> R): R =
-    originalFile.project.service<KtAnalysisSessionProvider>().analyseInFakeAnalysisSession(originalFile, fakeExpresion, action)
+inline fun <R> analyseInDependedAnalysisSession(
+    originalFile: KtFile,
+    dependencyExpression: KtElement,
+    action: KtAnalysisSession.() -> R
+): R =
+    originalFile.project.service<KtAnalysisSessionProvider>().analyseInDependedAnalysisSession(originalFile, dependencyExpression, action)
 
 /**
  * Execute given [action] in [KtAnalysisSession] context like [analyse] does but execute it in read action

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
@@ -14,31 +13,34 @@ 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.fir.resolve.providers.FirProvider
 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.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.idea.frontend.api.components.KtDiagnosticCheckerFilter
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 
-internal class FirModuleResolveStateForCompletion(
-    override val project: Project,
-    private val originalState: FirModuleResolveStateImpl
+internal class FirModuleResolveStateDepended(
+    dependencyFirDeclaration: FirDeclaration,
+    originalFirFile: FirFile,
+    private val originalState: FirModuleResolveStateImpl,
 ) : FirModuleResolveState() {
-    override val moduleInfo: IdeaModuleInfo get() = originalState.moduleInfo
 
+    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()
 
     override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
         originalState.getSessionFor(moduleInfo)
@@ -57,24 +59,8 @@ internal class FirModuleResolveStateForCompletion(
     override fun getFirFile(ktFile: KtFile): FirFile =
         originalState.getFirFile(ktFile)
 
-    override fun isFirFileBuilt(ktFile: KtFile): Boolean {
-        error("Should not be called in in completion")
-    }
-
-    override fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
-        synchronized(completionMapping) { fir.accept(FirElementsRecorder(), completionMapping) }
-    }
-
-    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
-        return originalState.resolvedFirToPhase(declaration, toPhase)
-    }
-
-    override fun lazyResolveDeclarationForCompletion(
-        firDeclaration: FirDeclaration,
-        containerFirFile: FirFile,
-    ) {
-        originalState.lazyResolveDeclarationForCompletion(firDeclaration, containerFirFile)
-    }
+    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
+        originalState.resolvedFirToPhase(declaration, toPhase)
 
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
         val ktFile = declaration.containingKtFileIfAny ?: return null
@@ -83,34 +69,37 @@ internal class FirModuleResolveStateForCompletion(
         return null
     }
 
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>> =
+        TODO("Diagnostics are not implemented for depended state")
 
-    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>> {
-        error("Diagnostics should not be retrieved in completion")
-    }
-
-    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> {
-        error("Diagnostics should not be retrieved in completion")
-    }
+    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
+        TODO("Diagnostics are not implemented for depended state")
 
     @OptIn(InternalForInline::class)
-    override fun findNonLocalSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
-        error("Should not be used in completion")
-    }
+    override fun findSourceFirDeclaration(ktDeclaration: KtLambdaExpression): FirDeclaration =
+        originalState.findSourceFirDeclaration(ktDeclaration)
 
     @OptIn(InternalForInline::class)
-    override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
-        return originalState.findSourceFirDeclaration(ktDeclaration)
-    }
-
-    @OptIn(InternalForInline::class)
-    override fun findSourceFirDeclaration(ktDeclaration: KtLambdaExpression): FirDeclaration {
-        error("Should not be used in completion")
-    }
-
-    override fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile? {
-        error("Should not be used in completion")
-    }
+    override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
+        originalState.findSourceFirDeclaration(ktDeclaration)
 
     override fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext? =
-        originalState.getTowerDataContextForElement(element)
+        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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
@@ -74,24 +75,12 @@ internal class FirModuleResolveStateImpl(
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>> =
         diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
 
-    override fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile? {
+    @TestOnly
+    internal fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile? {
         val cache = sessionProvider.getModuleCache(ktFile.getModuleInfo() as ModuleSourceInfo)
         return firFileBuilder.getBuiltFirFileOrNull(ktFile, cache)
     }
 
-    override fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
-        error("Should be called only from FirModuleResolveStateForCompletion")
-    }
-
-    @OptIn(InternalForInline::class)
-    override fun findNonLocalSourceFirDeclaration(
-        ktDeclaration: KtDeclaration,
-    ): FirDeclaration = ktDeclaration.findSourceNonLocalFirDeclaration(
-        firFileBuilder,
-        rootModuleSession.firIdeProvider.symbolProvider,
-        sessionProvider.getModuleCache(ktDeclaration.getModuleInfo() as ModuleSourceInfo)
-    )
-
     @OptIn(InternalForInline::class)
     override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
         findSourceFirDeclarationByExpression(ktDeclaration)
@@ -104,25 +93,28 @@ internal class FirModuleResolveStateImpl(
      * [ktDeclaration] should be either [KtDeclaration] or [KtLambdaExpression]
      */
     private fun findSourceFirDeclarationByExpression(ktDeclaration: KtExpression): FirDeclaration {
-        val nonLocalFirDeclaration = ktDeclaration.getNonLocalContainingOrThisDeclaration()
+        val nonLocalNamedDeclaration = ktDeclaration.getNonLocalContainingOrThisDeclaration()
             ?: error("Declaration should have non-local container${ktDeclaration.getElementTextInContext()}")
-        if (ktDeclaration == nonLocalFirDeclaration) return findNonLocalSourceFirDeclaration(ktDeclaration as KtDeclaration)
-        val container = nonLocalFirDeclaration.findSourceNonLocalFirDeclaration(
+
+        val nonLocalFirForNamedDeclaration = nonLocalNamedDeclaration.findSourceNonLocalFirDeclaration(
             firFileBuilder,
             rootModuleSession.firIdeProvider.symbolProvider,
             sessionProvider.getModuleCache(ktDeclaration.getModuleInfo() as ModuleSourceInfo)
         )
-        if (container.resolvePhase < FirResolvePhase.BODY_RESOLVE) {
-            val cache = (container.moduleData.session as FirIdeSourcesSession).cache
+
+        if (ktDeclaration == nonLocalNamedDeclaration) return nonLocalFirForNamedDeclaration
+
+        if (nonLocalFirForNamedDeclaration.resolvePhase < FirResolvePhase.BODY_RESOLVE) {
+            val cache = (nonLocalFirForNamedDeclaration.moduleData.session as FirIdeSourcesSession).cache
             firLazyDeclarationResolver.lazyResolveDeclaration(
-                container,
+                nonLocalFirForNamedDeclaration,
                 cache,
                 FirResolvePhase.BODY_RESOLVE,
                 checkPCE = false, /*TODO*/
                 towerDataContextCollector = collector,
             )
         }
-        val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(container) { firDeclaration ->
+        val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(nonLocalFirForNamedDeclaration) { firDeclaration ->
             when (val realPsi = firDeclaration.realPsi) {
                 is KtObjectLiteralExpression -> realPsi.objectDeclaration == ktDeclaration
                 is KtFunctionLiteral -> realPsi.parent == ktDeclaration
@@ -133,12 +125,6 @@ internal class FirModuleResolveStateImpl(
             ?: error("FirDeclaration was not found for\n${ktDeclaration.getElementTextInContext()}")
     }
 
-    override fun isFirFileBuilt(ktFile: KtFile): Boolean {
-        val moduleSourceInfo = ktFile.getModuleInfo() as? ModuleSourceInfo ?: return true
-        val cache = sessionProvider.getModuleCache(moduleSourceInfo)
-        return firFileBuilder.isFirFileBuilt(ktFile, cache)
-    }
-
     override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
         val fileCache = when (val session = declaration.moduleData.session) {
             is FirIdeSourcesSession -> session.cache
@@ -154,24 +140,6 @@ internal class FirModuleResolveStateImpl(
         return declaration
     }
 
-    override fun lazyResolveDeclarationForCompletion(
-        firDeclaration: FirDeclaration,
-        containerFirFile: FirFile,
-    ) {
-        firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, rootModuleSession.cache) {
-            firLazyDeclarationResolver.runLazyResolveWithoutLock(
-                firDeclaration,
-                rootModuleSession.cache,
-                containerFirFile,
-                containerFirFile.moduleData.session.firIdeProvider,
-                fromPhase = firDeclaration.resolvePhase,
-                toPhase = FirResolvePhase.BODY_RESOLVE,
-                towerDataContextCollector = collector,
-                checkPCE = true
-            )
-        }
-    }
-
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
         cache.getContainerFirFile(declaration)
 

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

@@ -18,31 +18,31 @@ import org.jetbrains.kotlin.psi.*
 object DeclarationCopyBuilder {
     fun createDeclarationCopy(
         originalFirDeclaration: FirDeclaration,
-        fakeKtDeclaration: KtDeclaration,
+        copiedKtDeclaration: KtDeclaration,
         state: FirModuleResolveState,
     ): FirDeclaration {
-        return when (fakeKtDeclaration) {
+        return when (copiedKtDeclaration) {
             is KtNamedFunction -> createFunctionCopy(
-                fakeKtDeclaration,
+                copiedKtDeclaration,
                 originalFirDeclaration as FirSimpleFunction,
                 state
             )
             is KtProperty -> createPropertyCopy(
-                fakeKtDeclaration,
+                copiedKtDeclaration,
                 originalFirDeclaration as FirProperty,
                 state
             )
             is KtClassOrObject -> createClassCopy(
-                fakeKtDeclaration,
+                copiedKtDeclaration,
                 originalFirDeclaration as FirRegularClass,
                 state
             )
             is KtTypeAlias -> createTypeAliasCopy(
-                fakeKtDeclaration,
+                copiedKtDeclaration,
                 originalFirDeclaration as FirTypeAlias,
                 state
             )
-            else -> error("Unsupported declaration ${fakeKtDeclaration::class.simpleName}")
+            else -> error("Unsupported declaration ${copiedKtDeclaration::class.simpleName}")
         }
     }
 
@@ -64,11 +64,11 @@ object DeclarationCopyBuilder {
     }
 
     private fun createClassCopy(
-        fakeKtClassOrObject: KtClassOrObject,
+        copiedKtClassOrObject: KtClassOrObject,
         originalFirClass: FirRegularClass,
         state: FirModuleResolveState,
     ): FirRegularClass {
-        val builtClass = createCopy(fakeKtClassOrObject, originalFirClass)
+        val builtClass = createCopy(copiedKtClassOrObject, originalFirClass)
 
         return buildRegularClassCopy(originalFirClass) {
             declarations.clear()
@@ -79,11 +79,11 @@ object DeclarationCopyBuilder {
     }
 
     private fun createTypeAliasCopy(
-        fakeKtTypeAlias: KtTypeAlias,
+        copiedKtTypeAlias: KtTypeAlias,
         originalFirTypeAlias: FirTypeAlias,
         state: FirModuleResolveState,
     ): FirTypeAlias {
-        val builtTypeAlias = createCopy(fakeKtTypeAlias, originalFirTypeAlias)
+        val builtTypeAlias = createCopy(copiedKtTypeAlias, originalFirTypeAlias)
 
         return buildTypeAliasCopy(originalFirTypeAlias) {
             expandedTypeRef = builtTypeAlias.expandedTypeRef
@@ -135,14 +135,14 @@ object DeclarationCopyBuilder {
     }
 
     internal inline fun <reified T : FirDeclaration> createCopy(
-        fakeKtDeclaration: KtDeclaration,
+        copiedKtDeclaration: KtDeclaration,
         originalFirDeclaration: T,
     ): T {
         return RawFirFragmentForLazyBodiesBuilder.build(
             session = originalFirDeclaration.moduleData.session,
             baseScopeProvider = originalFirDeclaration.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = originalFirDeclaration.collectDesignation().fullDesignation,
-            declaration = fakeKtDeclaration
+            declaration = copiedKtDeclaration
         ) as T
     }
 

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

@@ -34,8 +34,6 @@ abstract class FirModuleResolveState {
 
     internal abstract fun getFirFile(ktFile: KtFile): FirFile
 
-    internal abstract fun isFirFileBuilt(ktFile: KtFile): Boolean
-
     internal abstract fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<FirPsiDiagnostic<*>>
 
     internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>>
@@ -56,14 +54,6 @@ abstract class FirModuleResolveState {
         }
     }
 
-    @TestOnly
-    internal abstract fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile?
-
-    @InternalForInline
-    abstract fun findNonLocalSourceFirDeclaration(
-        ktDeclaration: KtDeclaration,
-    ): FirDeclaration
-
     @InternalForInline
     abstract fun findSourceFirDeclaration(
         ktDeclaration: KtDeclaration,
@@ -74,18 +64,8 @@ abstract class FirModuleResolveState {
         ktDeclaration: KtLambdaExpression,
     ): FirDeclaration
 
-
-    // todo temporary, used only in completion
-    internal abstract fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile)
-
     internal abstract fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 
-    // todo temporary, used only in completion
-    internal abstract fun lazyResolveDeclarationForCompletion(
-        firDeclaration: FirDeclaration,
-        containerFirFile: FirFile,
-    )
-
     internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?
 
     abstract fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext?

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

@@ -1,33 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.idea.fir.low.level.api.api
-
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateForCompletion
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtFile
-
-object LowLevelFirApiFacadeForCompletion {
-    fun getResolveStateForCompletion(originalState: FirModuleResolveState): FirModuleResolveState {
-        check(originalState is FirModuleResolveStateImpl)
-        return FirModuleResolveStateForCompletion(originalState.project, originalState)
-    }
-
-    fun recordCompletionContextForDeclaration(
-        firFile: FirFile,
-        fakeKtDeclaration: KtDeclaration,
-        originalKtDeclaration: KtDeclaration,
-        fakeContainingFile: KtFile,
-        state: FirModuleResolveState,
-    ) {
-        val originalFirDeclaration = originalKtDeclaration.getOrBuildFirOfType<FirDeclaration>(state)
-        val fakeFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(originalFirDeclaration, fakeKtDeclaration, state)
-
-        state.lazyResolveDeclarationForCompletion(fakeFirDeclaration, firFile)
-        state.recordPsiToFirMappingsForCompletionFrom(fakeFirDeclaration, firFile, fakeContainingFile)
-    }
-}

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

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.api
+
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.parentsOfType
+import org.jetbrains.kotlin.fir.declarations.*
+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.util.originalDeclaration
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+
+object LowLevelFirApiFacadeForDependentCopy {
+
+    private fun KtDeclaration.canBeEnclosingDeclaration(): Boolean = when (this) {
+        is KtNamedFunction -> isTopLevel || containingClassOrObject?.isLocal == false
+        is KtProperty -> isTopLevel || containingClassOrObject?.isLocal == false
+        is KtClassOrObject -> !isLocal
+        is KtTypeAlias -> isTopLevel() || containingClassOrObject?.isLocal == false
+        else -> false
+    }
+
+    private fun findEnclosingNonLocalDeclaration(position: KtElement): KtNamedDeclaration? =
+        position.parentsOfType<KtNamedDeclaration>().firstOrNull { ktDeclaration ->
+            ktDeclaration.canBeEnclosingDeclaration()
+        }
+
+    private fun <T : KtElement> locateDeclarationInFileByOffset(offsetElement: T, file: KtFile): T? {
+        val elementOffset = offsetElement.textOffset
+        val elementAtOffset = file.findElementAt(elementOffset) ?: return null
+        return PsiTreeUtil.getParentOfType(elementAtOffset, offsetElement::class.java, false)?.takeIf { it.textOffset == elementOffset }
+    }
+
+    private fun recordOriginalDeclaration(targetDeclaration: KtNamedDeclaration, originalDeclaration: KtNamedDeclaration) {
+        require(!targetDeclaration.isPhysical)
+        require(originalDeclaration.containingKtFile !== targetDeclaration.containingKtFile)
+        val originalDeclrationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
+        val fakeDeclarationParents = targetDeclaration.parentsOfType<KtDeclaration>().toList()
+        originalDeclrationParents.zip(fakeDeclarationParents) { original, fake ->
+            fake.originalDeclaration = original
+        }
+    }
+
+    fun getResolveStateForDependedCopy(
+        originalState: FirModuleResolveState,
+        originalKtFile: KtFile,
+        dependencyKtElement: KtElement
+    ): FirModuleResolveState {
+        require(originalState is FirModuleResolveStateImpl)
+
+        val dependencyNonLocalDeclaration = findEnclosingNonLocalDeclaration(dependencyKtElement)
+            ?: error("Cannot find enclosing declaration for ${dependencyKtElement.getElementTextInContext()}")
+
+        val originalNonLocalDeclaration = locateDeclarationInFileByOffset(dependencyNonLocalDeclaration, originalKtFile)
+            ?: error("Cannot find original function matching to ${dependencyNonLocalDeclaration.getElementTextInContext()} in $originalKtFile")
+
+        recordOriginalDeclaration(
+            targetDeclaration = dependencyNonLocalDeclaration,
+            originalDeclaration = originalNonLocalDeclaration
+        )
+
+        val originalFirDeclaration = originalNonLocalDeclaration.getOrBuildFirOfType<FirDeclaration>(originalState)
+        val copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
+            originalFirDeclaration = originalFirDeclaration,
+            copiedKtDeclaration = dependencyNonLocalDeclaration,
+            state = originalState
+        )
+
+        val originalFirFile = originalState.getFirFile(originalKtFile)
+
+        return FirModuleResolveStateDepended(copiedFirDeclaration, originalFirFile, originalState)
+    }
+}

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

@@ -36,6 +36,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : KotlinLightCodeInsightFix
         val declarationToResolve = lazyDeclarations.firstOrNull { it.name?.lowercase() == "resolveme" }
             ?: error("declaration with name `resolveMe` was not found")
         resolveWithClearCaches(ktFile) { firModuleResolveState ->
+            check(firModuleResolveState is FirModuleResolveStateImpl)
             val rendered = declarationToResolve.withFirDeclaration(
                 firModuleResolveState,
                 FirResolvePhase.BODY_RESOLVE

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

@@ -7,12 +7,9 @@ package org.jetbrains.kotlin.idea.frontend.api.fir
 
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.*
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.components.KtVisibilityChecker
@@ -20,8 +17,6 @@ import org.jetbrains.kotlin.idea.frontend.api.components.KtSymbolDeclarationRend
 import org.jetbrains.kotlin.idea.frontend.api.fir.components.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirOverrideInfoProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbolProvider
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.recordCompletionContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.threadLocal
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.psi.KtElement
@@ -33,9 +28,14 @@ private constructor(
     val firResolveState: FirModuleResolveState,
     internal val firSymbolBuilder: KtSymbolByFirBuilder,
     token: ValidityToken,
-    val context: KtFirAnalysisSessionContext,
+    private val mode: AnalysisSessionMode,
 ) : KtAnalysisSession(token) {
 
+    private enum class AnalysisSessionMode {
+        REGULAR,
+        DEPENDENT_COPY
+    }
+
     override val smartCastProviderImpl = KtFirSmartcastProvider(this, token)
 
     override val expressionTypeProviderImpl = KtFirExpressionTypeProvider(this, token)
@@ -73,19 +73,23 @@ private constructor(
 
     override val subtypingComponentImpl = KtFirSubtypingComponent(this, token)
 
-    override fun createContextDependentCopy(originalKtFile: KtFile, fakeKtElement: KtElement): KtAnalysisSession {
-        check(context == KtFirAnalysisSessionContext.DefaultContext) {
+    override fun createContextDependentCopy(originalKtFile: KtFile, dependencyKtElement: KtElement): KtAnalysisSession {
+        check(mode == AnalysisSessionMode.REGULAR) {
             "Cannot create context-dependent copy of KtAnalysis session from a context dependent one"
         }
-        val contextResolveState = LowLevelFirApiFacadeForCompletion.getResolveStateForCompletion(firResolveState)
-        val originalFirFile = originalKtFile.getFirFile(firResolveState)
-        val context = KtFirAnalysisSessionContext.FakeFileContext(originalKtFile, originalFirFile, fakeKtElement, contextResolveState)
+
+        val contextResolveState = LowLevelFirApiFacadeForDependentCopy.getResolveStateForDependedCopy(
+            originalState = firResolveState,
+            originalKtFile = originalKtFile,
+            dependencyKtElement = dependencyKtElement
+        )
+
         return KtFirAnalysisSession(
             project,
             contextResolveState,
             firSymbolBuilder.createReadOnlyCopy(contextResolveState),
             token,
-            context
+            AnalysisSessionMode.DEPENDENT_COPY
         )
     }
 
@@ -110,28 +114,8 @@ private constructor(
                 firResolveState,
                 firSymbolBuilder,
                 token,
-                KtFirAnalysisSessionContext.DefaultContext
+                AnalysisSessionMode.REGULAR
             )
         }
     }
-}
-
-internal sealed class KtFirAnalysisSessionContext {
-    object DefaultContext : KtFirAnalysisSessionContext()
-
-    class FakeFileContext(
-        originalFile: KtFile,
-        firFile: FirFile,
-        fakeContextElement: KtElement,
-        fakeModuleResolveState: FirModuleResolveState
-    ) : KtFirAnalysisSessionContext() {
-        init {
-            require(!fakeContextElement.isPhysical)
-
-            val enclosingContext = EnclosingDeclarationContext.detect(originalFile, fakeContextElement)
-            enclosingContext.recordCompletionContext(firFile, fakeModuleResolveState)
-        }
-
-        val fakeKtFile = fakeContextElement.containingKtFile
-    }
-}
+}

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

@@ -1,79 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.idea.frontend.api.fir.utils
-
-import com.intellij.psi.util.PsiTreeUtil
-import com.intellij.psi.util.parentsOfType
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
-import kotlin.reflect.KClass
-
-internal sealed class EnclosingDeclarationContext {
-    abstract val fakeKtEnclosingDeclaration: KtDeclaration
-    abstract val originalKtEnclosingDeclaration: KtDeclaration
-
-    companion object {
-        private fun KtDeclaration.canBeEnclosingDeclaration(): Boolean = when (this) {
-            is KtNamedFunction -> isTopLevel || containingClassOrObject?.isLocal == false
-            is KtProperty -> isTopLevel || containingClassOrObject?.isLocal == false
-            is KtClassOrObject -> !isLocal
-            is KtTypeAlias -> isTopLevel() || containingClassOrObject?.isLocal == false
-            else -> false
-        }
-
-        fun detect(originalFile: KtFile, positionInFakeFile: KtElement): EnclosingDeclarationContext {
-            val fakeKtDeclaration = positionInFakeFile.parentsOfType<KtNamedDeclaration>().firstOrNull { ktDeclaration ->
-                ktDeclaration.canBeEnclosingDeclaration()
-            }
-            if (fakeKtDeclaration != null) {
-                val originalDeclaration = findMatchingElementInCopy(fakeKtDeclaration, originalFile)
-                    ?: error("Cannot find original declaration matching to ${fakeKtDeclaration.getElementTextInContext()} in $originalFile")
-                recordOriginalDeclaration(originalDeclaration, fakeKtDeclaration)
-                return EnclosingDeclarationContextImpl(fakeKtDeclaration, originalDeclaration)
-            }
-
-            error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
-        }
-
-        private fun recordOriginalDeclaration(originalDeclaration: KtNamedDeclaration, fakeDeclaration: KtNamedDeclaration) {
-            require(!fakeDeclaration.isPhysical)
-            require(originalDeclaration.containingKtFile !== fakeDeclaration.containingKtFile)
-            val originalDeclrationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
-            val fakeDeclarationParents = fakeDeclaration.parentsOfType<KtDeclaration>().toList()
-
-            originalDeclrationParents.zip(fakeDeclarationParents) { original, fake ->
-                fake.originalDeclaration = original
-            }
-        }
-    }
-}
-
-internal class EnclosingDeclarationContextImpl(
-    override val fakeKtEnclosingDeclaration: KtDeclaration,
-    override val originalKtEnclosingDeclaration: KtDeclaration
-) : EnclosingDeclarationContext()
-
-
-internal fun EnclosingDeclarationContext.recordCompletionContext(originalFirFile: FirFile, firResolveState: FirModuleResolveState) {
-    LowLevelFirApiFacadeForCompletion.recordCompletionContextForDeclaration(
-        originalFirFile,
-        fakeKtEnclosingDeclaration,
-        originalKtEnclosingDeclaration,
-        state = firResolveState,
-        fakeContainingFile = fakeKtEnclosingDeclaration.containingKtFile
-    )
-}
-
-fun <T : KtElement> findMatchingElementInCopy(element: T, copy: KtFile): T? {
-    val elementOffset = element.textOffset
-    val elementAtOffset = copy.findElementAt(elementOffset) ?: return null
-    return PsiTreeUtil.getParentOfType(elementAtOffset, element::class.java, false)?.takeIf { it.textOffset == elementOffset }
-}