Analysis API: Add API functions to get all the candidates (with argument

mapping and substitutor) for a function call.

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10CallResolver.kt

@@ -146,6 +146,9 @@ internal class KtFe10CallResolver(
         } ?: handleResolveErrors(this, psi)
     }
 
+    // TODO: See Call.resolveCandidates() in plugins/kotlin/core/src/org/jetbrains/kotlin/idea/core/Utils.kt
+    override fun resolveCandidates(psi: KtElement): List<KtCallInfo> = TODO()
+
     private fun handleAsCompoundAssignment(context: BindingContext, binaryExpression: KtBinaryExpression): KtCallInfo? {
         val left = binaryExpression.left ?: return null
         val right = binaryExpression.right

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.api.fir.components
 import org.jetbrains.kotlin.analysis.api.calls.*
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtNonBoundToPsiErrorDiagnostic
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOf
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOfSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayTypeToArrayOfCall
@@ -18,24 +19,32 @@ import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir.getTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.diagnostics.KtDiagnostic
-import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
+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.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.calls.AbstractCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
+import org.jetbrains.kotlin.fir.resolve.createConeDiagnosticForCandidateWithError
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeDiagnosticWithCandidates
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeHiddenCandidateError
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
 import org.jetbrains.kotlin.fir.scopes.processOverriddenFunctions
@@ -628,6 +637,104 @@ internal class KtFirCallResolver(
     private fun FirValueParameterSymbol.toKtSymbol(): KtValueParameterSymbol =
         firSymbolBuilder.variableLikeBuilder.buildValueParameterSymbol(this)
 
+    override fun resolveCandidates(psi: KtElement): List<KtCallInfo> = withValidityAssertion {
+        val firCall = when (val fir = psi.getOrBuildFir(firResolveState)) {
+            is FirFunctionCall -> fir
+            is FirSafeCallExpression -> fir.regularQualifiedAccess.safeAs<FirFunctionCall>()
+            // TODO: FirDelegatedConstructorColl, FirAnnotationCall, FirArrayOfCall, FirConstructor
+            else -> null
+        } ?: return@withValidityAssertion emptyList()
+        val firFile = psi.containingKtFile.getOrBuildFirFile(firResolveState)
+        AllCandidatesResolver(analysisSession.rootModuleSession, firFile).getAllCandidates(firCall, psi)
+    }
+
+    private inner class AllCandidatesResolver(private val firSession: FirSession, private val firFile: FirFile) {
+        private val scopeSession = ScopeSession()
+
+        // TODO: This transformer is not intended for actual transformations and created here only to simplify access to call resolver
+        private val stubBodyResolveTransformer = FirBodyResolveTransformer(
+            session = firSession,
+            phase = FirResolvePhase.BODY_RESOLVE,
+            implicitTypeOnly = false,
+            scopeSession = scopeSession,
+        )
+
+        private val bodyResolveComponents =
+            FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents(
+                firSession,
+                scopeSession,
+                stubBodyResolveTransformer,
+                stubBodyResolveTransformer.context,
+            )
+
+        private val resolutionContext = ResolutionContext(firSession, bodyResolveComponents, stubBodyResolveTransformer.context)
+
+        @OptIn(PrivateForInline::class)
+        fun getAllCandidates(functionCall: FirFunctionCall, element: KtElement): List<KtCallInfo> {
+            val towerContext = firResolveState.getTowerContextProvider().getClosestAvailableParentContext(element)
+            towerContext?.let { bodyResolveComponents.context.replaceTowerDataContext(it) }
+            // Note: All candidate symbols should have the same name
+            val name =
+                functionCall.calleeReference.getCandidateSymbols().firstOrNull()?.safeAs<FirFunctionSymbol<*>>()?.name ?: return emptyList()
+            return bodyResolveComponents.context.withFile(firFile, bodyResolveComponents) {
+                val candidates = bodyResolveComponents.callResolver.collectAllCandidates(
+                    functionCall,
+                    name,
+                    element.getContainingDeclarations(),
+                    resolutionContext
+                )
+                candidates.mapNotNull { convertToKtCallInfo(functionCall, element, it.candidate, it.isInBestCandidates) }
+            }
+        }
+
+        private fun KtElement.getContainingDeclarations(): List<FirDeclaration> {
+            fun KtElement.getContainingKtDeclaration(): KtDeclaration? =
+                when (val container = this.parentOfType<KtDeclaration>()) {
+                    is KtDestructuringDeclaration -> container.parentOfType()
+                    else -> container
+                }
+
+            val containingDeclarations = mutableListOf<FirDeclaration>()
+            var current = getContainingKtDeclaration()
+            while (current != null) {
+                val firElement = current.getOrBuildFir(firResolveState)
+                val firDeclaration = when (firElement) {
+                    is FirAnonymousObjectExpression -> firElement.anonymousObject
+                    is FirAnonymousFunctionExpression -> firElement.anonymousFunction
+                    is FirDeclaration -> firElement
+                    else -> error(
+                        "Expected a FirDeclaration for KtDeclaration (type: ${current::class.simpleName}) " +
+                                "but was ${firElement?.let { it::class.simpleName }}. KtDeclaration text:\n${current.text}"
+                    )
+                }
+                containingDeclarations += firDeclaration
+                current = current.getContainingKtDeclaration()
+            }
+            return containingDeclarations.asReversed()
+        }
+
+        @OptIn(SymbolInternals::class)
+        private fun convertToKtCallInfo(
+            functionCall: FirFunctionCall,
+            element: KtElement,
+            candidate: Candidate,
+            isInBestCandidates: Boolean
+        ): KtCallInfo? {
+            val call = createKtCall(element, functionCall, candidate, resolveFragmentOfCall = false) // TODO: resolveFragmentOfCall
+                ?: error("expect `createKtCall` to succeed for candidate")
+            return if (candidate.isSuccessful && isInBestCandidates) {
+                KtSuccessCallInfo(call)
+            } else {
+                val diagnostic = createConeDiagnosticForCandidateWithError(candidate.currentApplicability, candidate)
+                if (diagnostic is ConeHiddenCandidateError) return null
+                val ktDiagnostic =
+                    functionCall.source?.let { diagnostic.asKtDiagnostic(it, element.toKtPsiSourceElement(), diagnosticCache) }
+                        ?: KtNonBoundToPsiErrorDiagnostic(factoryName = null, diagnostic.reason, token)
+                KtErrorCallInfo(listOf(call), ktDiagnostic, token)
+            }
+        }
+    }
+
     private fun FirArrayOfCall.toKtCallInfo(): KtCallInfo? {
         val arrayOfSymbol = with(analysisSession) {
             val type = typeRef.coneTypeSafe<ConeClassLikeType>()

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtCallResolver.kt

@@ -6,10 +6,14 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.calls.KtCallInfo
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtArrayAccessExpression
+import org.jetbrains.kotlin.psi.KtCallElement
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtUnaryExpression
 
 public abstract class KtCallResolver : KtAnalysisSessionComponent() {
     public abstract fun resolveCall(psi: KtElement): KtCallInfo?
+    public abstract fun resolveCandidates(psi: KtElement): List<KtCallInfo>
 }
 
 public interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
@@ -26,4 +30,10 @@ public interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
 
     public fun KtArrayAccessExpression.resolveCall(): KtCallInfo =
         analysisSession.callResolver.resolveCall(this) ?: error("KtArrayAccessExpression should always resolve to a KtCallInfo")
+
+    public fun KtCallElement.resolveCandidates(): List<KtCallInfo> =
+        analysisSession.callResolver.resolveCandidates(this)
+
+    public fun KtArrayAccessExpression.resolveCandidates(): List<KtCallInfo> =
+        analysisSession.callResolver.resolveCandidates(this)
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.references.impl.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.calls.tower.FirTowerResolver
+import org.jetbrains.kotlin.fir.resolve.calls.tower.TowerGroup
 import org.jetbrains.kotlin.fir.resolve.calls.tower.TowerResolveManager
 import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.resolve.inference.ResolvedCallableReferenceAtom
@@ -159,11 +160,43 @@ class FirCallResolver(
         val info: CallInfo, val applicability: CandidateApplicability, val candidates: Collection<Candidate>,
     )
 
+    /** WARNING: This function is public for the analysis API and should only be used there. */
+    fun <T : FirQualifiedAccess> collectAllCandidates(
+        qualifiedAccess: T,
+        name: Name,
+        containingDeclarations: List<FirDeclaration> = transformer.components.containingDeclarations,
+        resolutionContext: ResolutionContext = transformer.resolutionContext
+    ): List<OverloadCandidate> {
+        class AllCandidatesCollector(
+            components: BodyResolveComponents,
+            resolutionStageRunner: ResolutionStageRunner
+        ) : CandidateCollector(components, resolutionStageRunner) {
+            private val allCandidatesSet = mutableSetOf<Candidate>()
+
+            override fun consumeCandidate(group: TowerGroup, candidate: Candidate, context: ResolutionContext): CandidateApplicability {
+                allCandidatesSet += candidate
+                return super.consumeCandidate(group, candidate, context)
+            }
+
+            val allCandidates: List<Candidate>
+                get() = allCandidatesSet.toList()
+        }
+
+        val collector = AllCandidatesCollector(components, components.resolutionStageRunner)
+        val origin = (qualifiedAccess as? FirFunctionCall)?.origin ?: FirFunctionCallOrigin.Regular
+        val result =
+            collectCandidates(qualifiedAccess, name, forceCallKind = null, origin, containingDeclarations, resolutionContext, collector)
+        return collector.allCandidates.map { OverloadCandidate(it, isInBestCandidates = it in result.candidates) }
+    }
+
     private fun <T : FirQualifiedAccess> collectCandidates(
         qualifiedAccess: T,
         name: Name,
         forceCallKind: CallKind? = null,
-        origin: FirFunctionCallOrigin = FirFunctionCallOrigin.Regular
+        origin: FirFunctionCallOrigin = FirFunctionCallOrigin.Regular,
+        containingDeclarations: List<FirDeclaration> = transformer.components.containingDeclarations,
+        resolutionContext: ResolutionContext = transformer.resolutionContext,
+        collector: CandidateCollector? = null
     ): ResolutionResult {
         val explicitReceiver = qualifiedAccess.explicitReceiver
         val argumentList = (qualifiedAccess as? FirFunctionCall)?.argumentList ?: FirEmptyArgumentList
@@ -179,11 +212,15 @@ class FirCallResolver(
             typeArguments,
             session,
             components.file,
-            transformer.components.containingDeclarations,
+            containingDeclarations,
             origin = origin
         )
         towerResolver.reset()
-        val result = towerResolver.runResolver(info, transformer.resolutionContext)
+        val result = if (collector != null) {
+            towerResolver.runResolver(info, resolutionContext, collector)
+        } else {
+            towerResolver.runResolver(info, resolutionContext)
+        }
         val bestCandidates = result.bestCandidates()
 
         fun chooseMostSpecific(): Set<Candidate> {
@@ -773,19 +810,6 @@ class FirCallResolver(
         }
     }
 
-    private fun createConeDiagnosticForCandidateWithError(
-        applicability: CandidateApplicability,
-        candidate: Candidate
-    ): ConeDiagnostic {
-        return when (applicability) {
-            CandidateApplicability.HIDDEN -> ConeHiddenCandidateError(candidate)
-            CandidateApplicability.VISIBILITY_ERROR -> ConeVisibilityError(candidate.symbol)
-            CandidateApplicability.INAPPLICABLE_WRONG_RECEIVER -> ConeInapplicableWrongReceiver(listOf(candidate))
-            CandidateApplicability.NO_COMPANION_OBJECT -> ConeNoCompanionObject(candidate)
-            else -> ConeInapplicableCandidateError(applicability, candidate)
-        }
-    }
-
     private fun buildErrorReference(
         callInfo: CallInfo,
         diagnostic: ConeDiagnostic,
@@ -800,3 +824,6 @@ class FirCallResolver(
         )
     }
 }
+
+/** A candidate in the overload candidate set. */
+data class OverloadCandidate(val candidate: Candidate, val isInBestCandidates: Boolean)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -31,7 +31,7 @@ import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirPropertyWithExplicitBackingFieldResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.dfa.PropertyStability
-import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
+import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
 import org.jetbrains.kotlin.types.ConstantValueKind
+import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.types.SmartcastStability
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import kotlin.contracts.ExperimentalContracts
@@ -509,3 +510,16 @@ fun FirExpression?.isIntegerLiteralOrOperatorCall(): Boolean {
         else -> false
     }
 }
+
+fun createConeDiagnosticForCandidateWithError(
+    applicability: CandidateApplicability,
+    candidate: Candidate
+): ConeDiagnostic {
+    return when (applicability) {
+        CandidateApplicability.HIDDEN -> ConeHiddenCandidateError(candidate)
+        CandidateApplicability.VISIBILITY_ERROR -> ConeVisibilityError(candidate.symbol)
+        CandidateApplicability.INAPPLICABLE_WRONG_RECEIVER -> ConeInapplicableWrongReceiver(listOf(candidate))
+        CandidateApplicability.NO_COMPANION_OBJECT -> ConeNoCompanionObject(candidate)
+        else -> ConeInapplicableCandidateError(applicability, candidate)
+    }
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirArgumentsToParametersMapper.kt

@@ -133,20 +133,19 @@ private class FirCallArgumentsProcessor(
     }
 
     fun processNonLambdaArguments(arguments: List<FirExpression>) {
-        for (argument in arguments) {
-            // process position argument
-            if (argument !is FirNamedArgumentExpression) {
-                if (processPositionArgument(argument, isLastArgument = argument === arguments.last())) {
-                    state = State.VARARG_POSITION
+        for ((argumentIndex, argument) in arguments.withIndex()) {
+            if (argument is FirVarargArgumentsExpression) {
+                // If the argument list was already resolved, any arguments for a vararg parameter will be in a FirVarargArgumentsExpression.
+                // This can happen when getting all the candidates for an already resolved function call.
+                val varargArguments = argument.arguments
+                for ((varargArgumentIndex, varargArgument) in varargArguments.withIndex()) {
+                    processNonLambdaArgument(
+                        varargArgument,
+                        isLastArgument = argumentIndex == arguments.lastIndex && varargArgumentIndex == varargArguments.lastIndex
+                    )
                 }
-            }
-            // process named argument
-            else {
-                if (state == State.VARARG_POSITION) {
-                    completeVarargPositionArguments()
-                }
-
-                processNamedArgument(argument)
+            } else {
+                processNonLambdaArgument(argument, isLastArgument = argumentIndex == arguments.lastIndex)
             }
         }
         if (state == State.VARARG_POSITION) {
@@ -154,6 +153,23 @@ private class FirCallArgumentsProcessor(
         }
     }
 
+    private fun processNonLambdaArgument(argument: FirExpression, isLastArgument: Boolean) {
+        // process position argument
+        if (argument !is FirNamedArgumentExpression) {
+            if (processPositionArgument(argument, isLastArgument)) {
+                state = State.VARARG_POSITION
+            }
+        }
+        // process named argument
+        else {
+            if (state == State.VARARG_POSITION) {
+                completeVarargPositionArguments()
+            }
+
+            processNamedArgument(argument)
+        }
+    }
+
     // return true, if it was mapped to vararg parameter
     private fun processPositionArgument(argument: FirExpression, isLastArgument: Boolean): Boolean {
         if (state == State.NAMED_ONLY_ARGUMENTS) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirTowerResolver.kt

@@ -27,8 +27,16 @@ class FirTowerResolver(
     fun runResolver(
         info: CallInfo,
         context: ResolutionContext,
-        collector: CandidateCollector = this.collector,
-        manager: TowerResolveManager = this.manager
+        externalCollector: CandidateCollector? = null
+    ): CandidateCollector {
+        return runResolver(info, context, externalCollector ?: collector, manager)
+    }
+
+    fun runResolver(
+        info: CallInfo,
+        context: ResolutionContext,
+        collector: CandidateCollector,
+        manager: TowerResolveManager
     ): CandidateCollector {
         val candidateFactoriesAndCollectors = buildCandidateFactoriesAndCollectors(info, collector, context)