FIR IDE: Find FIR declarations for compiled PSI elements separately

When PSI declaration comes from the library with classfiles,
its `moduleInfo` is represented by `LibrarySourceInfo` class.
In this case we have to resort to other ways of looking for
corresponding FIR declaration

It is easy to do for classes (by `classId`) and for the properties
(by `classId` of the containing class and property's name)

It is harder for callables, so we use
`KtDeclarationAndFirDeclarationEqualityChecker` to do that

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt

@@ -8,8 +8,12 @@ package org.jetbrains.kotlin.fir.resolve.providers
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.resolve.getSymbolByLookupTag
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.getDeclaredConstructors
+import org.jetbrains.kotlin.fir.scopes.getFunctions
+import org.jetbrains.kotlin.fir.scopes.getProperties
 import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
-import org.jetbrains.kotlin.fir.symbols.*
+import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -51,13 +55,24 @@ abstract class FirSymbolProvider(val session: FirSession) : FirSessionComponent
     abstract fun getPackage(fqName: FqName): FqName? // TODO: Replace to symbol sometime
 }
 
-fun FirSymbolProvider.getClassDeclaredPropertySymbols(classId: ClassId, name: Name): List<FirVariableSymbol<*>> {
+private fun FirSymbolProvider.getClassDeclaredMemberScope(classId: ClassId): FirScope? {
-    val classSymbol = getClassLikeSymbolByFqName(classId) as? FirRegularClassSymbol ?: return emptyList()
+    val classSymbol = getClassLikeSymbolByFqName(classId) as? FirRegularClassSymbol ?: return null
-    val declaredMemberScope = session.declaredMemberScope(classSymbol.fir)
+    return session.declaredMemberScope(classSymbol.fir)
-    val result = mutableListOf<FirVariableSymbol<*>>()
+}
-    declaredMemberScope.processPropertiesByName(name, result::add)
 
-    return result
+fun FirSymbolProvider.getClassDeclaredConstructors(classId: ClassId): List<FirConstructorSymbol> {
+    val classMemberScope = getClassDeclaredMemberScope(classId)
+    return classMemberScope?.getDeclaredConstructors().orEmpty()
+}
+
+fun FirSymbolProvider.getClassDeclaredFunctionSymbols(classId: ClassId, name: Name): List<FirNamedFunctionSymbol> {
+    val classMemberScope = getClassDeclaredMemberScope(classId)
+    return classMemberScope?.getFunctions(name).orEmpty()
+}
+
+fun FirSymbolProvider.getClassDeclaredPropertySymbols(classId: ClassId, name: Name): List<FirVariableSymbol<*>> {
+    val classMemberScope = getClassDeclaredMemberScope(classId)
+    return classMemberScope?.getProperties(name).orEmpty()
 }
 
 inline fun <reified T : AbstractFirBasedSymbol<*>> FirSymbolProvider.getSymbolByTypeRef(typeRef: FirTypeRef): T? {

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

@@ -8,10 +8,12 @@ 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
@@ -44,7 +46,11 @@ inline fun <R> KtDeclaration.withFirDeclaration(
     phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
     action: (FirDeclaration) -> R
 ): R {
-    val firDeclaration = resolveState.findSourceFirDeclaration(this)
+    val firDeclaration = if (containingKtFile.isCompiled) {
+        findNonLocalDeclarationForCompiledElement(resolveState.rootModuleSession.symbolProvider)
+    } else {
+        resolveState.findSourceFirDeclaration(this)
+    }
 
     val resolvedDeclaration = if (firDeclaration.resolvePhase < phase) {
         firDeclaration.resolvedFirToPhase(phase, resolveState)

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

@@ -7,22 +7,26 @@ 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.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
+import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredConstructors
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredFunctionSymbols
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+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.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.api.collectDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 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.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
+/**
+ * 'Non-local' stands for not local classes/functions/etc.
+ */
 internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
     firFileBuilder: FirFileBuilder,
     firSymbolProvider: FirSymbolProvider,
@@ -61,6 +65,111 @@ 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<*>) =
+    KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiFunction, it)
+
+private fun representSameConstructor(
+    psiConstructor: KtConstructor<*>,
+    firConstructor: FirConstructor
+): Boolean {
+    if (firConstructor.isPrimary) {
+        return psiConstructor is KtPrimaryConstructor
+    }
+
+    return KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiConstructor, firConstructor)
+}
+
+private fun KtCallableDeclaration.findFunctionCandidates(
+    symbolProvider: FirSymbolProvider,
+    isTopLevel: Boolean
+): List<FirFunctionSymbol<*>> =
+    findCallableCandidates(symbolProvider, isTopLevel).filterIsInstance<FirFunctionSymbol<*>>()
+
+private fun KtCallableDeclaration.findPropertyCandidates(
+    symbolProvider: FirSymbolProvider,
+    isTopLevel: Boolean
+): List<FirPropertySymbol> =
+    findCallableCandidates(symbolProvider, isTopLevel).filterIsInstance<FirPropertySymbol>()
+
+private fun KtCallableDeclaration.findCallableCandidates(
+    symbolProvider: FirSymbolProvider,
+    isTopLevel: Boolean
+): List<FirCallableSymbol<*>> {
+    if (isTopLevel) {
+        return symbolProvider.getTopLevelCallableSymbols(containingKtFile.packageFqName, nameAsSafeName)
+    }
+
+    val containerClassId = containingClassOrObject?.getClassId()
+        ?: error("No containing non-local declaration found for ${getElementTextInContext()}")
+
+    return symbolProvider.getClassDeclaredFunctionSymbols(containerClassId, nameAsSafeName) +
+            symbolProvider.getClassDeclaredPropertySymbols(containerClassId, nameAsSafeName)
+}
+
 private fun KtDeclaration.findSourceOfNonLocalFirDeclarationByTraversingWholeTree(
     firFileBuilder: FirFileBuilder,
     moduleFileCache: ModuleFileCache,