FIR IDE: Make findNonLocalDeclarationForCompiledElement a member of FirModuleResolveState

2021-06-01 12:08:33 +03:00
parent e1f985424f
commit b0bac90a8b
5 changed files with 86 additions and 83 deletions
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api

 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
@@ -81,4 +80,8 @@ internal class FirModuleResolveStateDepended(
    @OptIn(InternalForInline::class)
    override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
        originalState.findSourceFirDeclaration(ktDeclaration)
+
+    @OptIn(InternalForInline::class)
+    override fun findSourceFirCompiledDeclaration(ktDeclaration: KtDeclaration) =
+        originalState.findSourceFirDeclaration(ktDeclaration)
 }
@@ -13,6 +13,8 @@ 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.realPsi
+import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredConstructors
+import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -29,11 +31,11 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarati
 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.FirElementFinder
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject

 internal class FirModuleResolveStateImpl(
    override val project: Project,
@@ -120,6 +122,66 @@ internal class FirModuleResolveStateImpl(
            ?: error("FirDeclaration was not found for\n${ktDeclaration.getElementTextInContext()}")
    }

+    @OptIn(InternalForInline::class)
+    override fun findSourceFirCompiledDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
+        require(ktDeclaration.containingKtFile.isCompiled) {
+            "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
+        }
+
+        return when (ktDeclaration) {
+            is KtClassOrObject -> {
+                require(!ktDeclaration.isLocal)
+                val classId = ktDeclaration.getClassId()
+                    ?: error("Non-local class should have classId. The class is ${ktDeclaration.getElementTextInContext()}")
+
+                val classCandidate = rootModuleSession.symbolProvider.getClassLikeSymbolByFqName(classId)
+                    ?: error("We should be able to find a symbol for $classId")
+
+                classCandidate.fir
+            }
+
+            is KtConstructor<*> -> {
+                val containingClass = ktDeclaration.containingClassOrObject
+                    ?: error("Constructor must have outer class: ${ktDeclaration.getElementTextInContext()}")
+
+                require(!containingClass.isLocal)
+                val classId = containingClass.getClassId()
+                    ?: error("Non-local class should have classId. The class is ${containingClass.getElementTextInContext()}")
+
+                val constructorCandidate =
+                    rootModuleSession.symbolProvider.getClassDeclaredConstructors(classId)
+                        .singleOrNull { representSameConstructor(ktDeclaration, it.fir) }
+                        ?: error("We should be able to find a constructor: ${ktDeclaration.getElementTextInContext()}")
+
+                constructorCandidate.fir
+            }
+
+            is KtNamedFunction -> {
+                require(!ktDeclaration.isLocal)
+
+                val functionCandidate =
+                    ktDeclaration.findFunctionCandidates(rootModuleSession.symbolProvider, ktDeclaration.isTopLevel)
+                        .singleOrNull { representSameFunction(ktDeclaration, it.fir) }
+                        ?: error("We should be able to find a symbol for function ${ktDeclaration.name}: ${ktDeclaration.getElementTextInContext()}")
+
+                functionCandidate.fir
+            }
+
+            is KtProperty -> {
+                require(!ktDeclaration.isLocal)
+
+                val propertyCandidate =
+                    ktDeclaration.findPropertyCandidates(rootModuleSession.symbolProvider, ktDeclaration.isTopLevel)
+                        .singleOrNull()
+                        ?: error("We should be able to find a symbol for property ${ktDeclaration.name}: ${ktDeclaration.getElementTextInContext()}")
+
+                propertyCandidate.fir
+            }
+
+            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
+        }
+    }
+
    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
        val fileCache = when (val session = declaration.moduleData.session) {
            is FirIdeSourcesSession -> session.cache
@@ -6,11 +6,9 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.api

 import com.intellij.openapi.project.Project
-import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 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.file.builder.ModuleFileCache
@@ -72,5 +70,16 @@ abstract class FirModuleResolveState {
        ktDeclaration: KtLambdaExpression,
    ): FirDeclaration

+    /**
+     * 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).
+     */
+    @InternalForInline
+    abstract fun findSourceFirCompiledDeclaration(
+        ktDeclaration: KtDeclaration
+    ): FirDeclaration
+
+
    internal abstract fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 }
@@ -8,12 +8,10 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.symbolProvider
 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.util.findNonLocalDeclarationForCompiledElement
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -47,7 +45,7 @@ inline fun <R> KtDeclaration.withFirDeclaration(
    action: (FirDeclaration) -> R
 ): R {
    val firDeclaration = if (containingKtFile.isCompiled) {
-        findNonLocalDeclarationForCompiledElement(resolveState.rootModuleSession.symbolProvider)
+        resolveState.findSourceFirCompiledDeclaration(this)
    } else {
        resolveState.findSourceFirDeclaration(this)
    }
@@ -8,14 +8,12 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.util
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredConstructors
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredFunctionSymbols
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredPropertySymbols
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.KtDeclarationAndFirDeclarationEqualityChecker
-import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.InvalidFirElementTypeException
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -65,91 +63,24 @@ internal inline fun <reified F : FirDeclaration> KtDeclaration.findFirDeclaratio
    return fir
 }

-/**
- * Looks for compiled non-local [this] declaration by querying its classId/callableId from the [symbolProvider].
- *
- * Works only if [this] is compiled (i.e. comes from .class file).
- */
-@InternalForInline
-fun KtDeclaration.findNonLocalDeclarationForCompiledElement(symbolProvider: FirSymbolProvider): FirDeclaration {
-    require(containingKtFile.isCompiled) {
-        "This method will only work on compiled declarations, but this declaration is not compiled: ${getElementTextInContext()}"
-    }
-
-    return when (this) {
-        is KtClassOrObject -> {
-            require(!isLocal)
-            val classId = getClassId()
-                ?: error("Non-local class should have classId. The class is ${getElementTextInContext()}")
-
-            val classCandidate = symbolProvider.getClassLikeSymbolByFqName(classId)
-                ?: error("We should be able to find a symbol for $classId")
-
-            classCandidate.fir
-        }
-
-        is KtConstructor<*> -> {
-            val containingClass = containingClassOrObject ?: error("Constructor must have outer class: ${getElementTextInContext()}")
-
-            require(!containingClass.isLocal)
-            val classId = containingClass.getClassId()
-                ?: error("Non-local class should have classId. The class is ${containingClass.getElementTextInContext()}")
-
-            val constructorCandidate =
-                symbolProvider.getClassDeclaredConstructors(classId)
-                    .singleOrNull { representSameConstructor(this, it.fir) }
-                    ?: error("We should be able to find a constructor: ${getElementTextInContext()}")
-
-            constructorCandidate.fir
-        }
-
-        is KtNamedFunction -> {
-            require(!isLocal)
-
-            val functionCandidate =
-                findFunctionCandidates(symbolProvider, isTopLevel)
-                    .singleOrNull { representSameFunction(this, it.fir) }
-                    ?: error("We should be able to find a symbol for function ${name}: ${getElementTextInContext()}")
-
-            functionCandidate.fir
-        }
-
-        is KtProperty -> {
-            require(!isLocal)
-
-            val propertyCandidate =
-                findPropertyCandidates(symbolProvider, isTopLevel)
-                    .singleOrNull()
-                    ?: error("We should be able to find a symbol for property ${name}: ${getElementTextInContext()}")
-
-            propertyCandidate.fir
-        }
-
-        else -> error("Unsupported compiled declaration of type ${this::class}: ${getElementTextInContext()}")
-    }
-}
-
-private fun representSameFunction(psiFunction: KtNamedFunction, it: FirFunction<*>) =
+internal fun representSameFunction(psiFunction: KtNamedFunction, it: FirFunction<*>): Boolean =
    KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiFunction, it)

-private fun representSameConstructor(
-    psiConstructor: KtConstructor<*>,
-    firConstructor: FirConstructor
-): Boolean {
-    if (firConstructor.isPrimary) {
-        return psiConstructor is KtPrimaryConstructor
+internal fun representSameConstructor(psiConstructor: KtConstructor<*>, firConstructor: FirConstructor): Boolean {
+    if ((firConstructor.isPrimary) != (psiConstructor is KtPrimaryConstructor)) {
+        return false
    }

    return KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiConstructor, firConstructor)
 }

-private fun KtCallableDeclaration.findFunctionCandidates(
+internal fun KtCallableDeclaration.findFunctionCandidates(
    symbolProvider: FirSymbolProvider,
    isTopLevel: Boolean
 ): List<FirFunctionSymbol<*>> =
    findCallableCandidates(symbolProvider, isTopLevel).filterIsInstance<FirFunctionSymbol<*>>()

-private fun KtCallableDeclaration.findPropertyCandidates(
+internal fun KtCallableDeclaration.findPropertyCandidates(
    symbolProvider: FirSymbolProvider,
    isTopLevel: Boolean
 ): List<FirPropertySymbol> =