FIR IDE: introduce withFirDeclaration function

2020-09-22 13:52:11 +03:00
parent ebc2ea59cd
commit a516264923
8 changed files with 66 additions and 16 deletions
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 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.structure.FirElementsRecorder
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile

@@ -23,6 +24,7 @@ internal class FirModuleResolveStateForCompletion(
    private val originalState: FirModuleResolveStateImpl
 ) : FirModuleResolveState() {
    override val moduleInfo: IdeaModuleInfo get() = originalState.moduleInfo
+
    override val rootModuleSession get() = originalState.rootModuleSession
    override val firTransformerProvider: FirTransformerProvider get() = originalState.firTransformerProvider
    private val fileStructureCache = originalState.fileStructureCache
@@ -69,4 +71,8 @@ internal class FirModuleResolveStateForCompletion(
    override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic> {
        error("Diagnostics should not be retrieved in completion")
    }
+
+    override fun findNonLocalSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
+        error("Should not be used in completion")
+    }
 }
@@ -13,6 +13,8 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
@@ -20,8 +22,11 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataC
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile

@@ -56,6 +61,14 @@ internal class FirModuleResolveStateImpl(
        error("Should be called only from FirModuleResolveStateForCompletion")
    }

+    override fun findNonLocalSourceFirDeclaration(
+        ktDeclaration: KtDeclaration,
+    ): FirDeclaration = ktDeclaration.findSourceNonLocalFirDeclaration(
+        firFileBuilder,
+        rootModuleSession.firIdeProvider.symbolProvider,
+        sessionProvider.getModuleCache(ktDeclaration.getModuleInfo() as ModuleSourceInfo)
+    )
+
    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
        val fileCache = when (val session = declaration.session) {
            is FirIdeSourcesSession -> session.cache
@@ -12,11 +12,15 @@ 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.providers.FirProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 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.FirTransformerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile

@@ -37,6 +41,10 @@ abstract class FirModuleResolveState {

    internal abstract fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic>

+    abstract fun findNonLocalSourceFirDeclaration(
+        ktDeclaration: KtDeclaration,
+    ): FirDeclaration
+
    // todo temporary, used only in completion
    internal abstract fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile)

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 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.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import kotlin.reflect.KClass
@@ -44,6 +45,25 @@ object LowLevelFirApiFacade {
    inline fun <R> withFirFile(ktFile: KtFile, resolveState: FirModuleResolveState, action: (FirFile) -> R): R =
        action(getFirFile(ktFile, resolveState))

+    /**
+     * Creates [FirDeclaration] by [KtDeclaration] and runs an [action] with it
+     * [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
+     * [FirDeclaration] passed to [action] will be resolved at least to [phase]
+     * Otherwise, some threading problems may arise,
+     *
+     * [ktDeclaration] should be non-local declaration (should have fully qualified name)
+     */
+    inline fun <R> withFirDeclaration(
+        ktDeclaration: KtDeclaration,
+        resolveState: FirModuleResolveState,
+        phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
+        action: (FirDeclaration) -> R
+    ): R {
+        val firDeclaration = resolveState.findNonLocalSourceFirDeclaration(ktDeclaration)
+        resolvedFirToPhase(firDeclaration, phase, resolveState)
+        return action(firDeclaration)
+    }
+
    fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> =
        resolveState.getDiagnostics(element)

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.hasExplicitTypeOrUnit
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.replaceFirst
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
@@ -111,7 +111,7 @@ internal class FileStructure(
    }

    private fun createDeclarationStructure(declaration: KtDeclaration): FileStructureElement {
-        val firDeclaration = declaration.findNonLocalFirDeclaration(firFileBuilder, firIdeProvider, moduleFileCache)
+        val firDeclaration = declaration.findSourceNonLocalFirDeclaration(firFileBuilder, firIdeProvider, moduleFileCache)
        firLazyDeclarationResolver.lazyResolveDeclaration(
            firDeclaration,
            moduleFileCache,
@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE
@@ -19,7 +18,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 import org.jetbrains.kotlin.psi.*

@@ -145,7 +144,7 @@ internal class FirLazyDeclarationResolver(
        if (!KtPsiUtil.isLocal(ktDeclaration)) return this
        val nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
            ?: error("Container for local declaration cannot be null")
-        return nonLocalPsi.findNonLocalFirDeclaration(firFileBuilder, provider, moduleFileCache)
+        return nonLocalPsi.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
    }
 }

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.isLibraryClasses
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.Immutable
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache

@Immutable
 class FirIdeSessionProvider internal constructor(
@@ -23,4 +24,7 @@ class FirIdeSessionProvider internal constructor(
 ) : FirSessionProvider {
    override fun getSession(moduleInfo: ModuleInfo): FirSession? =
        sessions[moduleInfo]
+
+    internal fun getModuleCache(moduleSourceInfo: ModuleSourceInfo): ModuleFileCache =
+        (sessions.getValue(moduleSourceInfo) as FirIdeSourcesSession).cache
 }
@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
@@ -17,18 +17,18 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject

-internal fun KtDeclaration.findNonLocalFirDeclaration(
+internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
    firFileBuilder: FirFileBuilder,
-    provider: FirProvider,
+    firSymbolProvider: FirSymbolProvider,
    moduleFileCache: ModuleFileCache
 ): FirDeclaration {
    require(!KtPsiUtil.isLocal(this))
    return when {
-        this is KtClassOrObject -> findFir(provider)
+        this is KtClassOrObject -> findFir(firSymbolProvider)
        this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
            val containerClass = containingClassOrObject
            val declarations = if (containerClass != null) {
-                val containerClassFir = containerClass.findFir(provider)
+                val containerClassFir = containerClass.findFir(firSymbolProvider)
                containerClassFir.declarations
            } else {
                val ktFile = containingKtFile
@@ -39,11 +39,11 @@ internal fun KtDeclaration.findNonLocalFirDeclaration(
            declarations.first { it.psi === this || it.psi == original }
        }
        this is KtConstructor<*> -> {
-            val containerClassFir = containingClassOrObject?.findFir(provider)
+            val containerClassFir = containingClassOrObject?.findFir(firSymbolProvider)
                ?: error("Container class should be not null for KtConstructor")
            containerClassFir.declarations.first { it.psi === this }
        }
-        this is KtTypeAlias -> findFir(provider)
+        this is KtTypeAlias -> findFir(firSymbolProvider)
        else -> error("Invalid container $this::class")
    }
 }
@@ -53,19 +53,19 @@ val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORI
 var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY)


-private fun KtClassOrObject.findFir(provider: FirProvider): FirRegularClass {
+private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass {
    val classId = classIdIfNonLocal()
        ?: error("Container classId should not be null for non-local declaration")
    return executeWithoutPCE {
-        provider.getFirClassifierByFqName(classId) as? FirRegularClass
+        firSymbolProvider.getClassLikeSymbolByFqName(classId)?.fir as? FirRegularClass
            ?: error("Could not find class $classId")
    }
 }

-private fun KtTypeAlias.findFir(provider: FirProvider): FirTypeAlias {
+private fun KtTypeAlias.findFir(firSymbolProvider: FirSymbolProvider): FirTypeAlias {
    val typeAlias = ClassId(containingKtFile.packageFqName, nameAsSafeName)
    return executeWithoutPCE {
-        provider.getFirClassifierByFqName(typeAlias) as? FirTypeAlias
+        firSymbolProvider.getClassLikeSymbolByFqName(typeAlias)?.fir as? FirTypeAlias
            ?: error("Could not find type alias $typeAlias")
    }
 }