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
2021-01-25 18:55:29 +03:00
parent dd91ebf1ea
commit 75ed167566
11 changed files with 32 additions and 86 deletions
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package testing

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

 fun test() {
@@ -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? {
@@ -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
            )
        }
@@ -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()}")
@@ -32,61 +32,34 @@ object LowLevelFirApiFacadeForCompletion {
        return FirModuleResolveStateForCompletion(originalState.project, originalState)
    }

-    class FirCompletionContext internal constructor(
-        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(
+    fun recordCompletionContextForFunction(
        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, contextCollector)
+        state.lazyResolveDeclarationForCompletion(copyFunction, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE)
        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, contextCollector)
+        state.lazyResolveDeclarationForCompletion(copyProperty, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE)
        state.recordPsiToFirMappingsForCompletionFrom(copyProperty, firFile, fakeElement.containingKtFile)
-
-        return FirCompletionContext(
-            copyProperty.session,
-            contextCollector
-        )
    }

    private fun buildFunctionCopyForCompletion(
@@ -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
@@ -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
@@ -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.getFirFile
-import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getTowerDataContextUnsafe
 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) }
@@ -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)
-    }
-}
@@ -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)