Analysis API: Move AllCandidatesResolver closer to

SingleCandidateResolver and share common code.

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

@@ -12,10 +12,7 @@ import org.jetbrains.kotlin.analysis.api.fir.evaluate.FirAnnotationValueConverte
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
-import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
-import org.jetbrains.kotlin.fir.expressions.FirImplicitInvokeCall
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirNamedReference
@@ -23,9 +20,13 @@ import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtCallElement
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 /**
  * Checks if the function call is an implicit invoke call with a simple qualified receiver, or looks like it.

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

@@ -15,24 +15,22 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayTypeToArrayOfCall
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.impl.barebone.parentOfType
-import org.jetbrains.kotlin.analysis.api.impl.barebone.parentsOfType
 import org.jetbrains.kotlin.analysis.api.impl.base.components.AbstractKtCallResolver
 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.getOrBuildFir
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.AllCandidatesResolver
 import org.jetbrains.kotlin.diagnostics.KtDiagnostic
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
-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.expressions.builder.buildFunctionCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
+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
@@ -40,15 +38,12 @@ 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.dfa.unwrapSmartcastExpression
 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
@@ -732,10 +727,7 @@ internal class KtFirCallResolver(
             ) ?: emptyList()
         }
         return when (this) {
-            is FirFunctionCall -> {
-                val firFile = psi.containingKtFile.getOrBuildFirFile(firResolveState)
-                AllCandidatesResolver(analysisSession.rootModuleSession, firFile).getAllCandidates(this, psi, resolveFragmentOfCall)
-            }
+            is FirFunctionCall -> collectCallCandidates(psi, resolveFragmentOfCall)
             is FirSafeCallExpression -> selector.collectCallCandidates(
                 psi,
                 resolveCalleeExpressionOfFunctionCall,
@@ -757,98 +749,68 @@ internal class KtFirCallResolver(
         }
     }
 
-    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)
-
-        fun getAllCandidates(functionCall: FirFunctionCall, element: KtElement, resolveFragmentOfCall: Boolean): List<KtCallInfo> {
-            initializeBodyResolveContext(element)
-
-            // If a function call is resolved to an implicit invoke call, the FirImplicitInvokeCall will have the `invoke()` function as the
-            // callee and the variable as the explicit receiver. To correctly get all candidates, we need to get the original function
-            // call's explicit receiver (if there is any) and callee (i.e., the variable).
-            val unwrappedExplicitReceiver = functionCall.explicitReceiver?.unwrapSmartcastExpression()
-            val originalFunctionCall =
-                if (functionCall is FirImplicitInvokeCall && unwrappedExplicitReceiver is FirPropertyAccessExpression) {
-                    val originalCallee = unwrappedExplicitReceiver.calleeReference.safeAs<FirNamedReference>() ?: return emptyList()
-                    buildFunctionCall {
-                        source = functionCall.source
-                        annotations.addAll(functionCall.annotations)
-                        typeArguments.addAll(functionCall.typeArguments)
-                        explicitReceiver = unwrappedExplicitReceiver.explicitReceiver
-                        argumentList = functionCall.argumentList
-                        calleeReference = originalCallee
-                    }
-                } else {
-                    functionCall
+    private fun FirFunctionCall.collectCallCandidates(
+        psi: KtElement,
+        resolveFragmentOfCall: Boolean
+    ): List<KtCallInfo> {
+        // If a function call is resolved to an implicit invoke call, the FirImplicitInvokeCall will have the `invoke()` function as the
+        // callee and the variable as the explicit receiver. To correctly get all candidates, we need to get the original function
+        // call's explicit receiver (if there is any) and callee (i.e., the variable).
+        val unwrappedExplicitReceiver = explicitReceiver?.unwrapSmartcastExpression()
+        val originalFunctionCall =
+            if (this is FirImplicitInvokeCall && unwrappedExplicitReceiver is FirPropertyAccessExpression) {
+                val originalCallee = unwrappedExplicitReceiver.calleeReference.safeAs<FirNamedReference>() ?: return emptyList()
+                buildFunctionCall {
+                    // NOTE: We only need to copy the explicit receiver and not the dispatch and extension receivers as only the explicit
+                    // receiver is needed by the resolver. The dispatch and extension receivers are only assigned after resolution when a
+                    // candidate is selected.
+                    source = this@collectCallCandidates.source
+                    annotations.addAll(this@collectCallCandidates.annotations)
+                    typeArguments.addAll(this@collectCallCandidates.typeArguments)
+                    explicitReceiver = unwrappedExplicitReceiver.explicitReceiver
+                    argumentList = this@collectCallCandidates.argumentList
+                    calleeReference = originalCallee
                 }
-
-            val calleeName = originalFunctionCall.calleeOrCandidateName ?: return emptyList()
-            val candidates = bodyResolveComponents.context.withFile(firFile, bodyResolveComponents) {
-                bodyResolveComponents.callResolver.collectAllCandidates(
-                    originalFunctionCall,
-                    calleeName,
-                    bodyResolveComponents.context.containers,
-                    resolutionContext
-                )
-            }
-            return candidates.mapNotNull {
-                convertToKtCallInfo(
-                    originalFunctionCall,
-                    element,
-                    it.candidate,
-                    it.isInBestCandidates,
-                    resolveFragmentOfCall
-                )
-            }
-        }
-
-        @OptIn(PrivateForInline::class, SymbolInternals::class)
-        private fun initializeBodyResolveContext(element: KtElement) {
-            // Set up needed context to get all candidates.
-            val towerContext = firResolveState.getTowerContextProvider(element.containingKtFile).getClosestAvailableParentContext(element)
-            towerContext?.let { bodyResolveComponents.context.replaceTowerDataContext(it) }
-            val containingDeclarations =
-                element.parentsOfType<KtDeclaration>().map { it.resolveToFirSymbol(firResolveState).fir }.toList().asReversed()
-            bodyResolveComponents.context.containers.addAll(containingDeclarations)
-        }
-
-        private fun convertToKtCallInfo(
-            functionCall: FirFunctionCall,
-            element: KtElement,
-            candidate: Candidate,
-            isInBestCandidates: Boolean,
-            resolveFragmentOfCall: Boolean
-        ): KtCallInfo? {
-            val call = createKtCall(element, functionCall, candidate, 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)
+                this
             }
+
+        val calleeName = originalFunctionCall.calleeOrCandidateName ?: return emptyList()
+        val candidates = AllCandidatesResolver(analysisSession.rootModuleSession).getAllCandidates(
+            analysisSession.firResolveState,
+            originalFunctionCall,
+            calleeName,
+            psi
+        )
+        return candidates.mapNotNull {
+            convertToKtCallInfo(
+                originalFunctionCall,
+                psi,
+                it.candidate,
+                it.isInBestCandidates,
+                resolveFragmentOfCall
+            )
+        }
+    }
+
+    private fun convertToKtCallInfo(
+        functionCall: FirFunctionCall,
+        element: KtElement,
+        candidate: Candidate,
+        isInBestCandidates: Boolean,
+        resolveFragmentOfCall: Boolean
+    ): KtCallInfo? {
+        val call = createKtCall(element, functionCall, candidate, 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)
         }
     }
 
@@ -873,7 +835,7 @@ internal class KtFirCallResolver(
                 is KtNameReferenceExpression -> psi.getReferencedNameAsName()
                 else -> {
                     // This could be KtArrayAccessExpression or KtOperationReferenceExpression.
-                    // Note: All candidate symbols should have the same name. We go by the symbol because `originalCallee.name` will include
+                    // Note: All candidate symbols should have the same name. We go by the symbol because `calleeReference.name` will include
                     // the applicability if not successful.
                     calleeReference.getCandidateSymbols().firstOrNull()?.safeAs<FirCallableSymbol<*>>()?.name
                 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/AllCandidatesResolver.kt

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2022 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.resolver
+
+import org.jetbrains.kotlin.analysis.api.impl.barebone.parentsOfType
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirModuleResolveState
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.OverloadCandidate
+import org.jetbrains.kotlin.fir.PrivateForInline
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtElement
+
+class AllCandidatesResolver(firSession: FirSession) {
+    private val bodyResolveComponents = createStubBodyResolveComponents(firSession)
+    private val resolutionContext = ResolutionContext(firSession, bodyResolveComponents, bodyResolveComponents.transformer.context)
+
+    fun getAllCandidates(
+        resolveState: LLFirModuleResolveState,
+        functionCall: FirFunctionCall,
+        calleeName: Name,
+        element: KtElement
+    ): List<OverloadCandidate> {
+        initializeBodyResolveContext(resolveState, element)
+
+        val firFile = element.containingKtFile.getOrBuildFirFile(resolveState)
+        return bodyResolveComponents.context.withFile(firFile, bodyResolveComponents) {
+            bodyResolveComponents.callResolver.collectAllCandidates(
+                functionCall,
+                calleeName,
+                bodyResolveComponents.context.containers,
+                resolutionContext
+            )
+        }
+    }
+
+    @OptIn(PrivateForInline::class, SymbolInternals::class)
+    private fun initializeBodyResolveContext(resolveState: LLFirModuleResolveState, element: KtElement) {
+        // Set up needed context to get all candidates.
+        val towerContext = resolveState.getTowerContextProvider(element.containingKtFile).getClosestAvailableParentContext(element)
+        towerContext?.let { bodyResolveComponents.context.replaceTowerDataContext(it) }
+        val containingDeclarations =
+            element.parentsOfType<KtDeclaration>().map { it.resolveToFirSymbol(resolveState).fir }.toList().asReversed()
+        bodyResolveComponents.context.containers.addAll(containingDeclarations)
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/SingleCandidateResolver.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.resolver
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirEmptyArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
@@ -16,8 +15,6 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
-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.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -29,22 +26,9 @@ class SingleCandidateResolver(
 ) {
     private val scopeSession = ScopeSession()
 
-    // TODO This transformer is not intended for actual transformations and created here only to simplify access to body resolve components
-    private val stubBodyResolveTransformer = object : FirBodyResolveTransformer(
-        session = firSession,
-        phase = FirResolvePhase.BODY_RESOLVE,
-        implicitTypeOnly = false,
-        scopeSession = scopeSession,
-    ) {}
-    private val bodyResolveComponents =
-        FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents(
-            firSession,
-            scopeSession,
-            stubBodyResolveTransformer,
-            stubBodyResolveTransformer.context,
-        )
+    private val bodyResolveComponents = createStubBodyResolveComponents(firSession)
     private val firCallCompleter = FirCallCompleter(
-        stubBodyResolveTransformer,
+        bodyResolveComponents.transformer,
         bodyResolveComponents,
     )
     private val resolutionStageRunner = ResolutionStageRunner()
@@ -62,7 +46,7 @@ class SingleCandidateResolver(
         val dispatchReceiverValue = infoProvider.dispatchReceiverValue()
         val implicitExtensionReceiverValue = infoProvider.implicitExtensionReceiverValue()
 
-        val resolutionContext = stubBodyResolveTransformer.resolutionContext
+        val resolutionContext = bodyResolveComponents.transformer.resolutionContext
 
         val candidate = CandidateFactory(resolutionContext, callInfo).createCandidate(
             callInfo,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/resolverUtils.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2022 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.resolver
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+
+internal fun createStubBodyResolveComponents(firSession: FirSession): FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents {
+    val scopeSession = ScopeSession()
+
+    // This transformer is not intended for actual transformations and created here only to simplify access to resolve components
+    val stubBodyResolveTransformer = FirBodyResolveTransformer(
+        session = firSession,
+        phase = FirResolvePhase.BODY_RESOLVE,
+        implicitTypeOnly = false,
+        scopeSession = scopeSession,
+    )
+
+    return FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents(
+        firSession,
+        scopeSession,
+        stubBodyResolveTransformer,
+        stubBodyResolveTransformer.context,
+    )
+}