[LL API] Fix declaration mapping for library dependencies

As there is no accurate way to figure out exact dependencies for each
library JAR, all project libraries but the current one are added as
dependencies. However, for performance reasons, each library is not
treated as an individual 'LLFirLibrarySession', instead a scope-based
'JvmClassFileBasedSymbolProvider' is used.

Although code analysis is correct with this setup, navigation between
libraries was broken. The reason is that source declarations were only
queried inside the library itself.

After the fix, declarations are first queried in the library,
and if there are no good candidates, all other libraries are scanned.
The same way things work in the Java plugin.

^KTIJ-23073 Fixed

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirDeserializedDeclarationSourceProvider.kt

@@ -8,9 +8,13 @@ package org.jetbrains.kotlin.analysis.api.fir
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.ProjectScope
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.KtDeclarationAndFirDeclarationEqualityChecker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.project.structure.KtBuiltinsModule
+import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
+import org.jetbrains.kotlin.analysis.project.structure.KtLibrarySourceModule
+import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.fir.FirElement
@@ -26,17 +30,19 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFunction
-import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.filterIsInstanceWithChecker
 
 //todo introduce LibraryModificationTracker based cache?
 internal object FirDeserializedDeclarationSourceProvider {
     fun findPsi(fir: FirElement, project: Project): PsiElement? {
         return when (fir) {
-            is FirSimpleFunction -> provideSourceForFunction(fir, project)
+            is FirSimpleFunction -> provideSourceForFunction(fir.delegatedWrapperData?.wrapped ?: fir, project)
             is FirProperty -> provideSourceForProperty(fir, project)
             is FirClass -> provideSourceForClass(fir, project)
             is FirTypeAlias -> provideSourceForTypeAlias(fir, project)
@@ -74,56 +80,81 @@ internal object FirDeserializedDeclarationSourceProvider {
         }
     }
 
-    private fun provideSourceForFunction(
+    private fun provideSourceForFunction(function: FirSimpleFunction, project: Project): PsiElement? {
-        function: FirSimpleFunction,
+        return provideSourceForCallable(
-        project: Project
+            function,
-    ): PsiElement? {
+            project,
-        val baseFunction = function.delegatedWrapperData?.wrapped ?: function
+            topLevelSearcher = { declarationProvider -> declarationProvider.getTopLevelFunctions(function.symbol.callableId) },
-        val candidates = if (function.isTopLevel) {
+            predicate = { it is KtFunction && it.name == function.name.asString() }
-            project.createDeclarationProvider(function.scope(project)).getTopLevelFunctions(function.symbol.callableId)
+        )
-                .filter(KtNamedFunction::isCompiled)
-        } else {
-            baseFunction.containingKtClass(project)?.body?.functions
-                ?.filter { it.name == baseFunction.name.asString() && it.isCompiled() }
-                .orEmpty()
-        }
-
-        return baseFunction.unwrapFakeOverrides().chooseCorrespondingPsi(candidates)
     }
 
     private fun provideSourceForProperty(property: FirProperty, project: Project): PsiElement? {
-        val candidates = if (property.isTopLevel) {
+        return provideSourceForCallable(
-            project.createDeclarationProvider(property.scope(project)).getTopLevelProperties(property.symbol.callableId)
+            property,
-        } else {
+            project,
-            property.containingKtClass(project)?.declarations
+            topLevelSearcher = { declarationProvider -> declarationProvider.getTopLevelProperties(property.symbol.callableId) },
-                ?.filter { it.name == property.name.asString() }
+            predicate = { it is KtProperty && it.name == property.name.asString() },
-                .orEmpty()
+        )
+    }
+
+    private fun <T : KtCallableDeclaration> provideSourceForCallable(
+        callable: FirCallableDeclaration,
+        project: Project,
+        topLevelSearcher: (KotlinDeclarationProvider) -> Collection<T>,
+        predicate: (KtCallableDeclaration) -> Boolean
+    ): PsiElement? {
+        fun chooseCandidate(candidates: Collection<KtCallableDeclaration>): PsiElement? {
+            return callable.unwrapFakeOverrides().chooseCorrespondingPsi(candidates)
         }
 
-        return candidates.firstOrNull(KtElement::isCompiled)
+        if (callable.isTopLevel) {
+            return getTargetScopes(callable, project)
+                .asSequence()
+                .map { scope -> topLevelSearcher(project.createDeclarationProvider(scope)).filter(KtElement::isCompiled) }
+                .filter { it.isNotEmpty() }
+                .firstNotNullOfOrNull { chooseCandidate(it) }
+        }
+
+        val containingKtClass = getContainingKtClassOrObject(callable, getTargetScopes(callable, project), project)
+        val containingKtFile = containingKtClass?.containingKtFile
+
+        if (containingKtFile == null || !containingKtFile.isCompiled) {
+            return null
+        }
+
+        val candidates = containingKtClass.declarations.filterIsInstanceWithChecker<KtCallableDeclaration>(predicate)
+        return chooseCandidate(candidates)
     }
 
-    private fun provideSourceForClass(klass: FirClass, project: Project): PsiElement? =
+    private fun provideSourceForClass(klass: FirClass, project: Project): PsiElement? {
-        classByClassId(klass.symbol.classId, klass.scope(project), project)
+        return getTargetScopes(klass, project)
-
+            .asSequence()
-    private fun provideSourceForTypeAlias(alias: FirTypeAlias, project: Project): PsiElement? {
+            .mapNotNull { scope -> classByClassId(klass.symbol.classId, scope, project) }
-        val candidates = project.createDeclarationProvider(alias.scope(project)).getAllTypeAliasesByClassId(alias.symbol.classId)
+            .filter { it.isCompiled() }
-        return candidates.firstOrNull(KtElement::isCompiled)
+            .firstOrNull()
     }
 
-    private fun provideSourceForConstructor(
+    private fun provideSourceForTypeAlias(typeAlias: FirTypeAlias, project: Project): PsiElement? {
-        constructor: FirConstructor,
+        return getTargetScopes(typeAlias, project)
-        project: Project
+            .asSequence()
-    ): PsiElement? {
+            .flatMap { scope -> project.createDeclarationProvider(scope).getAllTypeAliasesByClassId(typeAlias.symbol.classId) }
-        val containingKtClass = constructor.containingKtClass(project) ?: return null
+            .firstOrNull { it.isCompiled() }
-        if (constructor.isPrimary) return containingKtClass.primaryConstructor
-
-        return constructor.unwrapFakeOverrides().chooseCorrespondingPsi(containingKtClass.secondaryConstructors)
     }
 
-    private fun FirFunction.chooseCorrespondingPsi(
+    private fun provideSourceForConstructor(constructor: FirConstructor, project: Project): PsiElement? {
-        candidates: Collection<KtFunction>
+        val scopes = getTargetScopes(constructor, project)
-    ): KtFunction? {
+        val containingKtClass = getContainingKtClassOrObject(constructor, scopes, project) ?: return null
+
+        return if (constructor.isPrimary) {
+            containingKtClass.primaryConstructor
+        } else {
+            constructor.unwrapFakeOverrides()
+                .chooseCorrespondingPsi(containingKtClass.secondaryConstructors)
+        }
+    }
+
+    private fun FirCallableDeclaration.chooseCorrespondingPsi(candidates: Collection<KtCallableDeclaration>): KtCallableDeclaration? {
         if (candidates.isEmpty()) return null
         for (candidate in candidates) {
             assert(candidate.isCompiled()) {
@@ -136,32 +167,52 @@ internal object FirDeserializedDeclarationSourceProvider {
         return null
     }
 
-    private fun provideSourceForEnumEntry(
+    private fun provideSourceForEnumEntry(enumEntry: FirEnumEntry, project: Project): PsiElement? {
-        enumEntry: FirEnumEntry,
+        val scopes = getTargetScopes(enumEntry, project)
+        val containingKtClass = getContainingKtClassOrObject(enumEntry, scopes, project) ?: return null
+
+        if (containingKtClass.isCompiled()) {
+            return containingKtClass.body?.enumEntries?.firstOrNull { it.name == enumEntry.name.asString() }
+        }
+
+        return null
+    }
+
+    private fun getTargetScopes(declaration: FirDeclaration, project: Project): List<GlobalSearchScope> {
+        val originalDeclaration = when (declaration) {
+            is FirCallableDeclaration -> declaration.unwrapFakeOverrides()
+            else -> declaration
+        }
+
+        val moduleData = originalDeclaration.llFirModuleData
+        val module = moduleData.ktModule
+
+        return buildList {
+            add(module.contentScope)
+
+            when (module) {
+                is KtBuiltinsModule -> {
+                    add(ProjectScope.getLibrariesScope(project))
+                    add(ProjectScope.getContentScope(project))
+                }
+                is KtLibrarySourceModule, is KtLibraryModule -> {
+                    // Search in other libraries (that might be dependencies of this library)
+                    add(ProjectScope.getLibrariesScope(project))
+                }
+                else -> {}
+            }
+        }
+    }
+
+    private fun getContainingKtClassOrObject(
+        firCallable: FirCallableDeclaration,
+        scopes: List<GlobalSearchScope>,
         project: Project
-    ): PsiElement? {
+    ): KtClassOrObject? {
-        val candidates = enumEntry.containingKtClass(project)?.body?.enumEntries
+        val classId = firCallable.unwrapFakeOverrides().containingClassLookupTag()?.classId ?: return null
-            ?.filter { it.name == enumEntry.name.asString() }
+        return scopes.firstNotNullOfOrNull { scope -> classByClassId(classId, scope, project) }
-            .orEmpty()
-
-        return candidates.firstOrNull(KtElement::isCompiled)
     }
 
-    private fun FirDeclaration.scope(project: Project): GlobalSearchScope {
-        val original = when (this) {
-            is FirCallableDeclaration -> unwrapFakeOverrides()
-            else -> this
-        }
-        val moduleData = original.llFirModuleData
-        return when (val ktModule = moduleData.ktModule) {
-            is KtBuiltinsModule -> GlobalSearchScope.allScope(project) // TODO should be some stdlib
-            else -> ktModule.contentScope
-        }
-    }
-
-    private fun FirCallableDeclaration.containingKtClass(project: Project): KtClassOrObject? =
-        unwrapFakeOverrides().containingClassLookupTag()?.classId?.let { classByClassId(it, scope(project), project) }
-
     private fun classByClassId(classId: ClassId, scope: GlobalSearchScope, project: Project): KtClassOrObject? {
         val correctedClassId = classIdMapping[classId] ?: classId
         return project.createDeclarationProvider(scope)