LL API: simplify FirModuleResolveState interface

2022-01-21 16:56:00 +01:00
parent 05d65275bf
commit d9200affe4
12 changed files with 73 additions and 88 deletions
@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSessionProviderStorage
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
@@ -26,10 +27,10 @@ internal class FirIdeResolveStateService(project: Project) {
        project.createProjectWideOutOfBlockModificationTracker(),
        ProjectRootModificationTracker.getInstance(project),
    ) {
-        ConcurrentHashMap<KtModule, FirModuleResolveStateImpl>()
+        ConcurrentHashMap<KtModule, FirSourceModuleResolveState>()
    }

-    fun getResolveState(module: KtModule): FirModuleResolveStateImpl =
+    fun getResolveState(module: KtModule): FirSourceModuleResolveState =
        stateCache.computeIfAbsent(module) { createResolveStateFor(module, sessionProviderStorage) }

    companion object {
@@ -40,13 +41,13 @@ internal class FirIdeResolveStateService(project: Project) {
            module: KtModule,
            sessionProviderStorage: FirIdeSessionProviderStorage,
            configureSession: (FirIdeSession.() -> Unit)? = null,
-        ): FirModuleResolveStateImpl {
+        ): FirSourceModuleResolveState {
            if (module !is KtSourceModule) {
                error("Creating FirModuleResolveState is not yet supported for $module")
            }
            val sessionProvider = sessionProviderStorage.getSessionProvider(module, configureSession)
            val firFileBuilder = sessionProvider.rootModuleSession.firFileBuilder
-            return FirModuleResolveStateImpl(
+            return FirSourceModuleResolveState(
                sessionProviderStorage.project,
                module,
                sessionProvider,
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingKtFileIfAny
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -27,7 +28,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtLambdaExpression

 internal class FirModuleResolveStateDepended(
-    val originalState: FirModuleResolveStateImpl,
+    val originalState: FirSourceModuleResolveState,
    val towerProviderBuiltUponElement: FirTowerContextProvider,
    private val ktToFirMapping: KtToFirMapping?,
 ) : FirModuleResolveState() {
@@ -78,15 +79,7 @@ internal class FirModuleResolveStateDepended(
    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
        TODO("Diagnostics are not implemented for depended state")

-    override fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
-        return originalState.findSourceFirSymbol(ktDeclaration)
-    }
-
-    override fun findSourceFirSymbol(ktDeclaration: KtLambdaExpression): FirBasedSymbol<*> {
-        return originalState.findSourceFirSymbol(ktDeclaration)
-    }
-
-    override fun findSourceFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
-        return originalState.findSourceFirCompiledSymbol(ktDeclaration)
+    override fun resolveToFirSymbol(ktDeclaration: KtDeclaration, phase: FirResolvePhase): FirBasedSymbol<*> {
+        return originalState.resolveToFirSymbol(ktDeclaration, phase)
    }
 }
@@ -53,37 +53,7 @@ abstract class FirModuleResolveState {

    internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic>

-    abstract fun findSourceFirSymbol(
-        ktDeclaration: KtDeclaration,
-    ): FirBasedSymbol<*>
-
-    abstract fun findSourceFirSymbol(
-        ktDeclaration: KtLambdaExpression,
-    ): FirBasedSymbol<*>
-
-    /**
-     * Looks for compiled non-local [ktDeclaration] declaration by querying its classId/callableId from the SymbolProvider.
-     *
-     * Works only if [ktDeclaration] is compiled (i.e. comes from .class file).
-     */
-    abstract fun findSourceFirCompiledSymbol(
-        ktDeclaration: KtDeclaration
-    ): FirBasedSymbol<*>
-
-
-    open fun resolveToFirSymbol(
-        ktDeclaration: KtDeclaration,
-        resolveState: FirModuleResolveState,
-        phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
-    ): FirBasedSymbol<*> {
-        return when (val module = ktDeclaration.getKtModule()) {
-            is KtSourceModule -> resolveState.findSourceFirSymbol(ktDeclaration).also { resolveState.resolveFirToPhase(it.fir, phase) }
-            is KtLibraryModule -> resolveState.findSourceFirCompiledSymbol(ktDeclaration)
-            else -> error("unsupported module $module")
-        }
-    }
+    abstract fun resolveToFirSymbol(ktDeclaration: KtDeclaration, phase: FirResolvePhase): FirBasedSymbol<*>

    internal abstract fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase)
-
-
 }
@@ -43,7 +43,7 @@ fun KtDeclaration.resolveToFirSymbol(
    resolveState: FirModuleResolveState,
    phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
 ): FirBasedSymbol<*> {
-    return resolveState.resolveToFirSymbol(this, resolveState, phase)
+    return resolveState.resolveToFirSymbol(this, phase)
 }

 /**
@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.analysis.api.impl.barebone.parentOfType
 import org.jetbrains.kotlin.analysis.api.impl.barebone.parentsOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.DeclarationCopyBuilder.withBodyFrom
 import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateDepended
-import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerDataContextAllElementsCollector
@@ -67,7 +67,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
        place: T,
        elementToResolve: T,
    ): FirElement {
-        require(state is FirModuleResolveStateImpl)
+        require(state is FirSourceModuleResolveState)
        require(place.isPhysical)

        val declaration = runBodyResolveOnAir(
@@ -95,7 +95,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
        state: FirModuleResolveState,
        place: KtElement,
    ): FirTowerContextProvider {
-        require(state is FirModuleResolveStateImpl)
+        require(state is FirSourceModuleResolveState)

        return if (place is KtFile) {
            FileTowerProvider(place, onAirGetTowerContextForFile(state, place))
@@ -116,7 +116,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
    }

    private fun onAirGetTowerContextForFile(
-        state: FirModuleResolveStateImpl,
+        state: FirSourceModuleResolveState,
        file: KtFile,
    ): FirTowerDataContext {
        require(file.isPhysical)
@@ -144,7 +144,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
        originalKtFile: KtFile,
        elementToAnalyze: KtElement
    ): FirModuleResolveState {
-        require(originalState is FirModuleResolveStateImpl)
+        require(originalState is FirSourceModuleResolveState)
        require(elementToAnalyze !is KtFile) { "KtFile for dependency element not supported" }
        require(!elementToAnalyze.isPhysical) { "Depended state should be build only for non-physical elements" }

@@ -178,7 +178,7 @@ object LowLevelFirApiFacadeForResolveOnAir {

    private fun tryResolveAsFileAnnotation(
        annotationEntry: KtAnnotationEntry,
-        state: FirModuleResolveStateImpl,
+        state: FirSourceModuleResolveState,
        replacement: RawFirReplacement,
        firFile: FirFile,
        collector: FirTowerDataContextCollector? = null,
@@ -202,7 +202,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
    }

    private fun runBodyResolveOnAir(
-        state: FirModuleResolveStateImpl,
+        state: FirSourceModuleResolveState,
        replacement: RawFirReplacement,
        onAirCreatedDeclaration: Boolean,
        collector: FirTowerDataContextCollector? = null,
@@ -0,0 +1,10 @@
+/*
+ * Copyright 2010-2022 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.analysis.low.level.api.fir.state
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
+
+abstract class FirPhysicalModuleResolveState : FirModuleResolveState()
@@ -1,13 +1,12 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2022 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.analysis.low.level.api.fir
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state

 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirElementBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
@@ -22,6 +21,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirDeclarationForCom
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getElementTextInContext
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
+import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
@@ -38,13 +38,13 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.*

-internal class FirModuleResolveStateImpl(
+internal class FirSourceModuleResolveState(
    override val project: Project,
-    override val module: KtModule,
+    override val module: KtSourceModule,
    private val sessionProvider: FirIdeSessionProvider,
    val firFileBuilder: FirFileBuilder,
    val firLazyDeclarationResolver: FirLazyDeclarationResolver,
-) : FirModuleResolveState() {
+) : FirPhysicalModuleResolveState() {
    override val rootModuleSession: FirIdeSourcesSession get() = sessionProvider.rootModuleSession

    /**
@@ -82,30 +82,28 @@ internal class FirModuleResolveStateImpl(
    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
        diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)

-    override fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+    override fun resolveToFirSymbol(
+        ktDeclaration: KtDeclaration,
+        phase: FirResolvePhase
+    ): FirBasedSymbol<*> {
+        return when (val module = ktDeclaration.getKtModule()) {
+            is KtSourceModule -> findSourceFirSymbol(ktDeclaration).also { resolveFirToPhase(it.fir, phase) }
+            is KtLibraryModule -> findFirCompiledSymbol(ktDeclaration)
+            else -> error("unsupported module $module")
+        }
+    }
+
+    private fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
        return findSourceFirDeclarationByExpression(ktDeclaration.originalDeclaration ?: ktDeclaration).symbol
    }

-    override fun findSourceFirSymbol(ktDeclaration: KtLambdaExpression): FirBasedSymbol<*> {
-        return findSourceFirDeclarationByExpression(ktDeclaration).symbol
-    }
-
-    override fun findSourceFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+    private fun findFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
        require(ktDeclaration.containingKtFile.isCompiled) {
            "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
        }

        val searcher = FirDeclarationForCompiledElementSearcher(rootModuleSession.symbolProvider)
-
-        val firDeclaration = when (ktDeclaration) {
-            is KtEnumEntry -> searcher.findNonLocalEnumEntry(ktDeclaration)
-            is KtClassOrObject -> searcher.findNonLocalClass(ktDeclaration)
-            is KtConstructor<*> -> searcher.findConstructorOfNonLocalClass(ktDeclaration)
-            is KtNamedFunction -> searcher.findNonLocalFunction(ktDeclaration)
-            is KtProperty -> searcher.findNonLocalProperty(ktDeclaration)
-
-            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
-        }
+        val firDeclaration = searcher.findNonLocalDeclaration(ktDeclaration)
        return firDeclaration.symbol
    }

@@ -21,8 +21,19 @@ import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 * Allows to search for FIR declarations by compiled [KtDeclaration]s.
 */
 internal class FirDeclarationForCompiledElementSearcher(private val symbolProvider: FirSymbolProvider) {
+    fun findNonLocalDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
+        return when (ktDeclaration) {
+            is KtEnumEntry -> findNonLocalEnumEntry(ktDeclaration)
+            is KtClassOrObject -> findNonLocalClass(ktDeclaration)
+            is KtConstructor<*> -> findConstructorOfNonLocalClass(ktDeclaration)
+            is KtNamedFunction -> findNonLocalFunction(ktDeclaration)
+            is KtProperty -> findNonLocalProperty(ktDeclaration)

-    fun findNonLocalEnumEntry(declaration: KtEnumEntry): FirEnumEntry {
+            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
+        }
+    }
+
+    private fun findNonLocalEnumEntry(declaration: KtEnumEntry): FirEnumEntry {
        require(!declaration.isLocal)
        val classId = declaration.containingClassOrObject?.getClassId()
            ?: error("Non-local class should have classId. The class is ${declaration.getElementTextInContext()}")
@@ -35,7 +46,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
        } as FirEnumEntry
    }

-    fun findNonLocalClass(declaration: KtClassOrObject): FirClassLikeDeclaration {
+    private fun findNonLocalClass(declaration: KtClassOrObject): FirClassLikeDeclaration {
        require(!declaration.isLocal)
        val classId = declaration.getClassId()
            ?: error("Non-local class should have classId. The class is ${declaration.getElementTextInContext()}")
@@ -46,7 +57,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
        return classCandidate.fir
    }

-    fun findConstructorOfNonLocalClass(declaration: KtConstructor<*>): FirConstructor {
+    private fun findConstructorOfNonLocalClass(declaration: KtConstructor<*>): FirConstructor {
        val containingClass = declaration.containingClassOrObject
            ?: error("Constructor must have outer class: ${declaration.getElementTextInContext()}")

@@ -62,7 +73,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
        return constructorCandidate.fir
    }

-    fun findNonLocalFunction(declaration: KtNamedFunction): FirFunction {
+    private fun findNonLocalFunction(declaration: KtNamedFunction): FirFunction {
        require(!declaration.isLocal)

        val functionCandidate =
@@ -73,7 +84,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
        return functionCandidate.fir
    }

-    fun findNonLocalProperty(declaration: KtProperty): FirProperty {
+    private fun findNonLocalProperty(declaration: KtProperty): FirProperty {
        require(!declaration.isLocal)

        val propertyCandidate =
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir

+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
@@ -53,7 +54,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractLowLevelApiSingle
        val rendererOption = FirRenderer.RenderMode.WithDeclarationAttributes.copy(renderDeclarationResolvePhase = true)
        val resultBuilder = StringBuilder()
        resolveWithClearCaches(ktFile) { firModuleResolveState ->
-            check(firModuleResolveState is FirModuleResolveStateImpl)
+            check(firModuleResolveState is FirSourceModuleResolveState)
            val declarationToResolve = firModuleResolveState
                .getOrBuildFirFile(ktFile)
                .findResolveMe()
@@ -67,7 +68,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractLowLevelApiSingle
        }

        resolveWithClearCaches(ktFile) { firModuleResolveState ->
-            check(firModuleResolveState is FirModuleResolveStateImpl)
+            check(firModuleResolveState is FirSourceModuleResolveState)
            val firFile = firModuleResolveState.getOrBuildFirFile(ktFile)
            firFile.ensureResolved(FirResolvePhase.BODY_RESOLVE)
            resultBuilder.append("\nFILE RAW TO BODY:\n")
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir

 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.render
@@ -32,7 +33,7 @@ abstract class AbstractFirOnAirResolveTest : AbstractLowLevelApiSingleFileTest()
        check(place::class == onAir::class)

        resolveWithClearCaches(ktFile) { resolveState ->
-            check(resolveState is FirModuleResolveStateImpl)
+            check(resolveState is FirSourceModuleResolveState)
            val firElement = LowLevelFirApiFacadeForResolveOnAir.onAirResolveElement(resolveState, place, onAir)
            val rendered = firElement.render(FirRenderer.RenderMode.WithResolvePhases)
            testServices.assertions.assertEqualsToTestDataFileSibling(rendered)
@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
-import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getDiagnostics
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
@@ -41,7 +41,7 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
                register(BeforeElementDiagnosticCollectionHandler::class, handler)
            }
        ) { resolveState ->
-            check(resolveState is FirModuleResolveStateImpl)
+            check(resolveState is FirSourceModuleResolveState)

            val fileStructure = resolveState.fileStructureCache.getFileStructure(ktFile, resolveState.rootModuleSession.cache)
            val allStructureElements = fileStructure.getAllStructureElements()
@@ -9,7 +9,7 @@ import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.impl.source.tree.LeafPsiElement
-import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.psi.*
@@ -85,7 +85,7 @@ abstract class AbstractFileStructureTest : AbstractLowLevelApiSingleFileTest() {
    }

    private fun KtFile.getFileStructure(): FileStructure {
-        val moduleResolveState = getResolveState() as FirModuleResolveStateImpl
+        val moduleResolveState = getResolveState() as FirSourceModuleResolveState
        return moduleResolveState.fileStructureCache.getFileStructure(
            ktFile = this,
            moduleFileCache = moduleResolveState.rootModuleSession.cache