[LL FIR] KT-57220 Add LLFirSelectingCombinedSymbolProvider

- Classpath order disambiguation will also be needed by the combined
  Java symbol provider, so it makes sense to move it to its own base
  class.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirCombinedKotlinSymbolProvider.kt

@@ -10,7 +10,6 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirSymbolProviderNameCache
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
@@ -25,9 +24,7 @@ import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtClassLikeDeclaration
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.utils.mapToIndex
 
 /**
  * [LLFirCombinedKotlinSymbolProvider] combines multiple [LLFirProvider.SymbolProvider]s with the following advantages:
@@ -46,22 +43,7 @@ internal class LLFirCombinedKotlinSymbolProvider(
     private val project: Project,
     private val providers: List<LLFirProvider.SymbolProvider>,
     private val declarationProvider: KotlinDeclarationProvider,
-) : FirSymbolProvider(session) {
+) : LLFirSelectingCombinedSymbolProvider<LLFirProvider.SymbolProvider>(session, project, providers) {
-    private val providersByKtModule: Map<KtModule, LLFirProvider.SymbolProvider> =
-        providers
-            .groupBy { it.session.llFirModuleData.ktModule }
-            .mapValues { (module, list) -> list.singleOrNull() ?: error("$module must have a unique Kotlin symbol providers.") }
-
-    /**
-     * [KtModule] precedence must be checked in case the index finds multiple elements and classpath order needs to be preserved.
-     */
-    private val modulePrecedenceMap: Map<KtModule, Int> = providers.map { it.session.llFirModuleData.ktModule }.mapToIndex()
-
-    /**
-     * Cache [ProjectStructureProvider] to avoid service access when getting [KtModule]s.
-     */
-    private val projectStructureProvider: ProjectStructureProvider = project.getService(ProjectStructureProvider::class.java)
-
     private val symbolNameCache = object : LLFirSymbolProviderNameCache(session) {
         override fun computeClassifierNames(packageFqName: FqName): Set<String>? = buildSet {
             providers.forEach { addAll(it.knownTopLevelClassifiersInPackage(packageFqName) ?: return null) }
@@ -72,39 +54,17 @@ internal class LLFirCombinedKotlinSymbolProvider(
         }
     }
 
-    @OptIn(FirSymbolProviderInternals::class)
     override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
         if (!symbolNameCache.mayHaveTopLevelClassifier(classId, mayHaveFunctionClass = false)) return null
 
+        return computeClassLikeSymbolByClassId(classId)
+    }
+
+    @OptIn(FirSymbolProviderInternals::class)
+    private fun computeClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
         val candidates = declarationProvider.getAllClassesByClassId(classId) + declarationProvider.getAllTypeAliasesByClassId(classId)
-        if (candidates.isEmpty()) return null
+        val (ktClass, provider) = selectFirstElementInClasspathOrder(candidates) ?: return null
-
+        return provider.getClassLikeSymbolByClassId(classId, ktClass)
-        // Find the `KtClassLikeDeclaration` with the highest module precedence. (We're using a custom implementation instead of `minBy` so
-        // that `ktModule` doesn't need to be fetched twice.)
-        // TODO (marco): This algorithm can be applied to other combined symbol providers as well, such as Java symbol providers.
-        var ktClassCandidate: KtClassLikeDeclaration? = null
-        var currentPrecedence: Int = Int.MAX_VALUE
-        var ktModule: KtModule? = null
-
-        candidates.forEach { candidate ->
-            val candidateKtModule = projectStructureProvider.getKtModuleForKtElement(candidate)
-
-            // If `candidateKtModule` cannot be found in the map, `candidate` cannot be processed by any of the available providers, because
-            // none of them belong to the correct module. We can skip in that case, because iterating through all providers wouldn't lead to
-            // any results for `candidate`.
-            val precedence = modulePrecedenceMap[candidateKtModule] ?: return@forEach
-            if (precedence < currentPrecedence) {
-                ktClassCandidate = candidate
-                currentPrecedence = precedence
-                ktModule = candidateKtModule
-            }
-        }
-
-        val ktClass = ktClassCandidate ?: return null
-
-        // The provider will always be found at this point, because `modulePrecedenceMap` contains the same keys as `providersByKtModule`
-        // and a precedence for `ktModule` must have been found in the previous step.
-        return providersByKtModule[ktModule]!!.getClassLikeSymbolByClassId(classId, ktClass)
     }
 
     @FirSymbolProviderInternals

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirSelectingCombinedSymbolProvider.kt

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2010-2023 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.providers
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.utils.mapToIndex
+
+/**
+ * A combined symbol provider which *selects* a subset of its [providers] to delegate to by an element's [KtModule]. For example, a
+ * selecting symbol provider might perform an index access, get a number of candidate PSI elements, and delegate to the appropriate provider
+ * for the element directly.
+ *
+ * Classpath order must be preserved with [selectFirstElementInClasspathOrder] in case a single result is required.
+ */
+abstract class LLFirSelectingCombinedSymbolProvider<PROVIDER : FirSymbolProvider>(
+    session: FirSession,
+    project: Project,
+    private val providers: List<PROVIDER>,
+) : FirSymbolProvider(session) {
+    protected val providersByKtModule: Map<KtModule, PROVIDER> =
+        providers
+            .groupingBy { it.session.llFirModuleData.ktModule }
+            // `reduce` invokes the `error` operation if it encounters a second element.
+            .reduce { module, _, _ -> error("$module must have a unique symbol provider.") }
+
+    /**
+     * [KtModule] precedence must be checked in case of multiple candidates to preserve classpath order.
+     */
+    private val modulePrecedenceMap: Map<KtModule, Int> = providers.map { it.session.llFirModuleData.ktModule }.mapToIndex()
+
+    /**
+     * Cache [ProjectStructureProvider] to avoid service access when getting [KtModule]s.
+     */
+    private val projectStructureProvider: ProjectStructureProvider = project.getService(ProjectStructureProvider::class.java)
+
+    /**
+     * Selects the element with the highest module precedence in [candidates], returning the element and the provider to which resolution
+     * should be delegated. This is a post-processing step that preserves classpath order when, for example, an index access with a combined
+     * scope isn't guaranteed to return the first element in classpath order.
+     */
+    protected fun <ELEMENT : PsiElement> selectFirstElementInClasspathOrder(candidates: Collection<ELEMENT>): Pair<ELEMENT, PROVIDER>? {
+        if (candidates.isEmpty()) return null
+
+        // We're using a custom implementation instead of `minBy` so that `ktModule` doesn't need to be fetched twice.
+        var currentElement: ELEMENT? = null
+        var currentPrecedence: Int = Int.MAX_VALUE
+        var currentKtModule: KtModule? = null
+
+        for (candidate in candidates) {
+            val ktModule = projectStructureProvider.getKtModuleForKtElement(candidate)
+
+            // If `ktModule` cannot be found in the map, `candidate` cannot be processed by any of the available providers, because none of
+            // them belong to the correct module. We can skip in that case, because iterating through all providers wouldn't lead to any
+            // results for `candidate`.
+            val precedence = modulePrecedenceMap[ktModule] ?: continue
+            if (precedence < currentPrecedence) {
+                currentElement = candidate
+                currentPrecedence = precedence
+                currentKtModule = ktModule
+            }
+        }
+
+        val element = currentElement ?: return null
+        val ktModule = currentKtModule ?: error("`currentKtModule` must not be `null` when `currentElement` has been found.")
+
+        // The provider will always be found at this point, because `modulePrecedenceMap` contains the same keys as `providersByKtModule`
+        // and a precedence for `currentKtModule` must have been found in the previous step.
+        val provider = providersByKtModule.getValue(ktModule)
+
+        return Pair(element, provider)
+    }
+}