FIR IDE: Use FirModuleResolveState instead of TowerDataContextProvider

This way we will have a single place to get information about scopes,
so in the future it would be easier to refactor it

idea/idea-completion/testData/basic/common/BeforeDotInCall.kt

@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package testing
 
 fun testTop() {

idea/idea-completion/testData/basic/common/InLongDotQualifiedExpression.kt

@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 // For KT-2796
 
 fun test() {

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

@@ -16,9 +16,7 @@ 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.annotations.PrivateForInline
 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.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny
@@ -73,9 +71,8 @@ internal class FirModuleResolveStateForCompletion(
         containerFirFile: FirFile,
         firIdeProvider: FirProvider,
         toPhase: FirResolvePhase,
-        towerDataContextCollector: FirTowerDataContextCollector
     ) {
-        originalState.lazyResolveDeclarationForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextCollector)
+        originalState.lazyResolveDeclarationForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase)
     }
 
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {

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

@@ -45,6 +45,15 @@ internal class FirModuleResolveStateImpl(
     val firLazyDeclarationResolver: FirLazyDeclarationResolver,
 ) : FirModuleResolveState() {
     override val rootModuleSession: FirIdeSourcesSession get() = sessionProvider.rootModuleSession
+
+    /**
+     * WARNING! This object contains scopes for all statements and declarations that were ever resolved.
+     * 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 elementBuilder = FirElementBuilder()
@@ -149,7 +158,6 @@ internal class FirModuleResolveStateImpl(
         containerFirFile: FirFile,
         firIdeProvider: FirProvider,
         toPhase: FirResolvePhase,
-        towerDataContextCollector: FirTowerDataContextCollector
     ) {
         firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, rootModuleSession.cache) {
             firLazyDeclarationResolver.runLazyResolveWithoutLock(
@@ -159,7 +167,7 @@ internal class FirModuleResolveStateImpl(
                 firIdeProvider,
                 fromPhase = firFunction.resolvePhase,
                 toPhase,
-                towerDataContextCollector,
+                towerDataContextCollector = collector,
                 checkPCE = true
             )
         }

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

@@ -13,13 +13,8 @@ import org.jetbrains.kotlin.fir.declarations.*
 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.caches.project.getModuleInfo
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
-import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.PrivateForInline
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
-import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
@@ -75,10 +70,12 @@ abstract class FirModuleResolveState {
         containerFirFile: FirFile,
         firIdeProvider: FirProvider,
         toPhase: FirResolvePhase,
-        towerDataContextCollector: FirTowerDataContextCollector
     )
 
     internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?
 
     abstract fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext?
 }
+
+fun FirModuleResolveState.getTowerDataContextUnsafe(element: KtElement): FirTowerDataContext =
+    getTowerDataContextForElement(element) ?: error("No context for ${element.getElementTextInContext()}")

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

@@ -32,61 +32,34 @@ object LowLevelFirApiFacadeForCompletion {
         return FirModuleResolveStateForCompletion(originalState.project, originalState)
     }
 
-    class FirCompletionContext internal constructor(
+    fun recordCompletionContextForFunction(
-        val session: FirSession,
-        private val towerDataContextCollector: FirTowerDataContextCollector,
-    ) {
-        fun getTowerDataContext(element: KtElement): FirTowerDataContext {
-            var current: PsiElement? = element
-            while (current is KtElement) {
-                towerDataContextCollector.getContext(current)?.let { return it }
-                current = current.parent
-            }
-
-            error("No context for ${element.getElementTextInContext()}")
-        }
-    }
-
-    fun buildCompletionContextForFunction(
         firFile: FirFile,
         fakeElement: KtNamedFunction,
         originalElement: KtNamedFunction,
         state: FirModuleResolveState,
-    ): FirCompletionContext {
+    ) {
         val firIdeProvider = firFile.session.firIdeProvider
 
         val originalFunction = state.getOrBuildFirFor(originalElement) as FirSimpleFunction
         val copyFunction = buildFunctionCopyForCompletion(firIdeProvider, fakeElement, originalFunction, state)
 
-        val contextCollector = FirTowerDataContextCollector()
+        state.lazyResolveDeclarationForCompletion(copyFunction, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE)
-        state.lazyResolveDeclarationForCompletion(copyFunction, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE, contextCollector)
         state.recordPsiToFirMappingsForCompletionFrom(copyFunction, firFile, fakeElement.containingKtFile)
-
-        return FirCompletionContext(
-            copyFunction.session,
-            contextCollector
-        )
     }
 
-    fun buildCompletionContextForProperty(
+    fun recordCompletionContextForProperty(
         firFile: FirFile,
         fakeElement: KtProperty,
         originalElement: KtProperty,
         state: FirModuleResolveState,
-    ): FirCompletionContext {
+    ) {
         val firIdeProvider = firFile.session.firIdeProvider
 
         val originalProperty = state.getOrBuildFirFor(originalElement) as FirProperty
         val copyProperty = buildPropertyCopyForCompletion(firIdeProvider, fakeElement, originalProperty, state)
 
-        val contextCollector = FirTowerDataContextCollector()
+        state.lazyResolveDeclarationForCompletion(copyProperty, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE)
-        state.lazyResolveDeclarationForCompletion(copyProperty, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE, contextCollector)
         state.recordPsiToFirMappingsForCompletionFrom(copyProperty, firFile, fakeElement.containingKtFile)
-
-        return FirCompletionContext(
-            copyProperty.session,
-            contextCollector
-        )
     }
 
     private fun buildFunctionCopyForCompletion(

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

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.idea.frontend.api.components.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.components.*
 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.buildCompletionContext
+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.symbols.KtSymbolProvider
 import org.jetbrains.kotlin.psi.KtElement
@@ -38,8 +38,6 @@ private constructor(
         assertIsValidAndAccessible()
     }
 
-    internal val towerDataContextProvider: TowerDataContextProvider = TowerDataContextProvider(this)
-
     override val smartCastProvider: KtSmartCastProvider = KtFirSmartcastProvider(this, token)
     override val expressionTypeProvider: KtExpressionTypeProvider = KtFirExpressionTypeProvider(this, token)
     override val diagnosticProvider: KtDiagnosticProvider = KtFirDiagnosticProvider(this, token)
@@ -110,11 +108,9 @@ internal sealed class KtFirAnalysisSessionContext {
     ) : KtFirAnalysisSessionContext() {
         init {
             require(!fakeContextElement.isPhysical)
-        }
 
-        val completionContext = run {
             val enclosingContext = EnclosingDeclarationContext.detect(originalFile, fakeContextElement)
-            enclosingContext.buildCompletionContext(firFile, fakeModuleResolveState)
+            enclosingContext.recordCompletionContext(firFile, fakeModuleResolveState)
         }
 
         val fakeKtFile = fakeContextElement.containingKtFile

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 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.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
@@ -20,9 +21,6 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFunctionSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirKotlinPropertySymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.buildCompletionContext
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.fakeEnclosingDeclaration
 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
@@ -39,9 +37,6 @@ internal class KtFirCompletionCandidateChecker(
 ) : KtCompletionCandidateChecker(), KtFirAnalysisSessionComponent {
     override val analysisSession: KtFirAnalysisSession by weakRef(analysisSession)
 
-    private val completionContextCache =
-        ConcurrentHashMap<Pair<FirFile, KtCallableDeclaration>, LowLevelFirApiFacadeForCompletion.FirCompletionContext>()
-
     override fun checkExtensionFitsCandidate(
         firSymbolForCandidate: KtCallableSymbol,
         originalFile: KtFile,
@@ -92,7 +87,7 @@ internal class KtFirCompletionCandidateChecker(
         firFile: FirFile,
         fakeNameExpression: KtSimpleNameExpression
     ): Sequence<ImplicitReceiverValue<*>?> {
-        val towerDataContext = analysisSession.towerDataContextProvider.getTowerDataContext(fakeNameExpression)
+        val towerDataContext = analysisSession.firResolveState.getTowerDataContextUnsafe(fakeNameExpression)
 
         return sequence {
             yield(null) // otherwise explicit receiver won't be checked when there are no implicit receivers in completion position

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

@@ -16,9 +16,7 @@ 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.LowLevelFirApiFacadeForCompletion
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getTowerDataContextUnsafe
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
-import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeContext
@@ -31,8 +29,6 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymb
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirEnumEntrySymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFileSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.buildCompletionContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFileSymbol
@@ -147,7 +143,7 @@ internal class KtFirScopeProvider(
         originalFile: KtFile,
         positionInFakeFile: KtElement
     ): KtScopeContext = withValidityAssertion {
-        val towerDataContext = analysisSession.towerDataContextProvider.getTowerDataContext(positionInFakeFile)
+        val towerDataContext = analysisSession.firResolveState.getTowerDataContextUnsafe(positionInFakeFile)
 
         val implicitReceivers = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver }.distinct()
         val implicitReceiversTypes = implicitReceivers.map { builder.buildKtType(it.type) }

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

@@ -1,19 +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.components
-
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
-import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
-import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSessionContext
-import org.jetbrains.kotlin.psi.KtElement
-
-internal class TowerDataContextProvider(private val analysisSession: KtFirAnalysisSession) {
-    fun getTowerDataContext(statement: KtElement): FirTowerDataContext {
-        val fakeContext =  analysisSession.context as? KtFirAnalysisSessionContext.FakeFileContext
-            ?: error("Getting data context for non context-dependent session is not supported yet")
-        return fakeContext.completionContext.getTowerDataContext(statement)
-    }
-}

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

@@ -65,22 +65,23 @@ internal val EnclosingDeclarationContext.fakeEnclosingDeclaration: KtCallableDec
         is PropertyContext -> fakeEnclosingProperty
     }
 
-internal fun EnclosingDeclarationContext.buildCompletionContext(originalFirFile: FirFile, firResolveState: FirModuleResolveState) =
+internal fun EnclosingDeclarationContext.recordCompletionContext(originalFirFile: FirFile, firResolveState: FirModuleResolveState) {
     when (this) {
-        is FunctionContext -> LowLevelFirApiFacadeForCompletion.buildCompletionContextForFunction(
+        is FunctionContext -> LowLevelFirApiFacadeForCompletion.recordCompletionContextForFunction(
             originalFirFile,
             fakeEnclosingFunction,
             originalEnclosingFunction,
             state = firResolveState
         )
 
-        is PropertyContext -> LowLevelFirApiFacadeForCompletion.buildCompletionContextForProperty(
+        is PropertyContext -> LowLevelFirApiFacadeForCompletion.recordCompletionContextForProperty(
             originalFirFile,
             fakeEnclosingProperty,
             originalEnclosingProperty,
             state = firResolveState
         )
     }
+}
 
 private inline fun <reified T : KtElement> KtFile.findDeclarationOfTypeAt(offset: Int): T? =
     findElementAt(offset)