[FIR-IDE] Changes after review

Move usage of SingleCandidateResolver into dedicated component from KtFirAnalysisSession.
Add validity assertion, use concurrent map for context cache, provide original PSI when building context.
Update names, make other minor changes.

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.completion
 import com.intellij.codeInsight.completion.*
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PsiJavaPatterns
+import com.intellij.psi.PsiElement
 import com.intellij.util.ProcessingContext
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
@@ -79,12 +80,19 @@ private object KotlinAvailableScopesCompletionProvider {
         val explicitReceiver = nameExpression.getQualifiedExpressionForSelector()?.receiverExpression
 
         with(getAnalysisSessionFor(originalFile).createContextDependentCopy()) {
-            val (implicitScopes, implicitReceivers) = originalFile.getScopeContextForPosition(nameExpression)
+            val (implicitScopes, _) = originalFile.getScopeContextForPosition(nameExpression)
 
             when {
                 nameExpression.parent is KtUserType -> collectTypesCompletion(result, implicitScopes)
-                explicitReceiver != null -> collectDotCompletion(result, implicitScopes, explicitReceiver)
-                else -> collectDefaultCompletion(result, implicitScopes, implicitReceivers)
+                explicitReceiver != null -> collectDotCompletion(
+                    result,
+                    implicitScopes,
+                    explicitReceiver,
+                    nameExpression,
+                    originalFile,
+                    parameters.originalPosition
+                )
+                else -> collectDefaultCompletion(result, implicitScopes, nameExpression, originalFile, parameters.originalPosition)
             }
         }
     }
@@ -97,7 +105,10 @@ private object KotlinAvailableScopesCompletionProvider {
     private fun KtAnalysisSession.collectDotCompletion(
         result: CompletionResultSet,
         implicitScopes: KtCompositeScope,
-        explicitReceiver: KtExpression
+        explicitReceiver: KtExpression,
+        nameExpression: KtSimpleNameExpression,
+        originalFile: KtFile,
+        originalPosition: PsiElement?,
     ) {
         val typeOfPossibleReceiver = explicitReceiver.getKtType()
         val possibleReceiverScope = typeOfPossibleReceiver.getTypeScope() ?: return
@@ -108,16 +119,25 @@ private object KotlinAvailableScopesCompletionProvider {
 
         val extensionNonMembers = implicitScopes
             .getCallableSymbols()
-            .filter { it.isExtension && it.canBeCalledWith(listOf(typeOfPossibleReceiver)) }
+            .filter {
+                it.isExtension && it.checkExtensionIsSuitable(
+                    originalFile,
+                    originalPosition,
+                    nameExpression,
+                    explicitReceiver
+                )
+            }
 
         nonExtensionMembers.forEach { result.addSymbolToCompletion(it) }
         extensionNonMembers.forEach { result.addSymbolToCompletion(it) }
     }
 
-    private fun collectDefaultCompletion(
+    private fun KtAnalysisSession.collectDefaultCompletion(
         result: CompletionResultSet,
         implicitScopes: KtCompositeScope,
-        implicitReceivers: List<KtType>
+        nameExpression: KtSimpleNameExpression,
+        originalFile: KtFile,
+        originalPosition: PsiElement?,
     ) {
         val availableNonExtensions = implicitScopes
             .getCallableSymbols()
@@ -125,7 +145,14 @@ private object KotlinAvailableScopesCompletionProvider {
 
         val extensionsWhichCanBeCalled = implicitScopes
             .getCallableSymbols()
-            .filter { it.isExtension && it.canBeCalledWith(implicitReceivers) }
+            .filter {
+                it.isExtension && it.checkExtensionIsSuitable(
+                    originalFile,
+                    originalPosition,
+                    nameExpression,
+                    null
+                )
+            }
 
         availableNonExtensions.forEach { result.addSymbolToCompletion(it) }
         extensionsWhichCanBeCalled.forEach { result.addSymbolToCompletion(it) }
@@ -133,10 +160,3 @@ private object KotlinAvailableScopesCompletionProvider {
         collectTypesCompletion(result, implicitScopes)
     }
 }
-
-private fun KtCallableSymbol.canBeCalledWith(implicitReceivers: List<KtType>): Boolean {
-    val requiredReceiverType = (this as? KtPossibleExtensionSymbol)?.receiverType
-        ?: error("Extension receiver type should be present on $this")
-
-    return implicitReceivers.any { it.isSubTypeOf(requiredReceiverType) }
-}

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -40,6 +40,7 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
     protected abstract val containingDeclarationProvider: KtSymbolContainingDeclarationProvider
     protected abstract val symbolProvider: KtSymbolProvider
     protected abstract val callResolver: KtCallResolver
+    protected abstract val completionCandidateChecker: KtCompletionCandidateChecker
 
 
     /// TODO: get rid of
@@ -119,10 +120,16 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
         return resolveToSymbols().singleOrNull()
     }
 
-    abstract fun resolveAndCheckReceivers(
-        firSymbolForCandidate: KtCallableSymbol,
-        originalFile: KtFile,
-        nameExpression: KtSimpleNameExpression,
-        possibleReceiver: KtExpression?,
-    ): Boolean
+    fun KtCallableSymbol.checkExtensionIsSuitable(
+        originalPsiFile: KtFile,
+        originalPosition: PsiElement?,
+        psiFakeCompletionExpression: KtSimpleNameExpression,
+        psiReceiverExpression: KtExpression?,
+    ): Boolean = completionCandidateChecker.checkExtensionFitsCandidate(
+        this,
+        originalPsiFile,
+        originalPosition,
+        psiFakeCompletionExpression,
+        psiReceiverExpression
+    )
 }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtCompletionCandidateChecker.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2010-2020 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.idea.frontend.api.components
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+
+abstract class KtCompletionCandidateChecker : KtAnalysisSessionComponent() {
+    abstract fun checkExtensionFitsCandidate(
+        firSymbolForCandidate: KtCallableSymbol,
+        originalFile: KtFile,
+        originalPosition: PsiElement?,
+        nameExpression: KtSimpleNameExpression,
+        possibleExplicitReceiver: KtExpression?,
+    ): Boolean
+}

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

@@ -49,14 +49,14 @@ class SingleCandidateResolver(
         resolutionParameters: ResolutionParameters
     ): FirFunctionCall? {
 
-        val partProvider = createCandidatePartsProvider(resolutionParameters)
-        if (partProvider.shouldFailBeforeResolve())
+        val infoProvider = createCandidateInfoProvider(resolutionParameters)
+        if (infoProvider.shouldFailBeforeResolve())
             return null
 
-        val callInfo = partProvider.callInfo()
-        val explicitReceiverKind = partProvider.explicitReceiverKind()
-        val dispatchReceiverValue = partProvider.dispatchReceiverValue()
-        val implicitExtensionReceiverValue = partProvider.implicitExtensionReceiverValue()
+        val callInfo = infoProvider.callInfo()
+        val explicitReceiverKind = infoProvider.explicitReceiverKind()
+        val dispatchReceiverValue = infoProvider.dispatchReceiverValue()
+        val implicitExtensionReceiverValue = infoProvider.implicitExtensionReceiverValue()
 
         val candidate = CandidateFactory(bodyResolveComponents, callInfo).createCandidate(
             resolutionParameters.callableSymbol,
@@ -72,7 +72,7 @@ class SingleCandidateResolver(
         return null
     }
 
-    private fun createCandidatePartsProvider(resolutionParameters: ResolutionParameters): CandidateInfoProvider {
+    private fun createCandidateInfoProvider(resolutionParameters: ResolutionParameters): CandidateInfoProvider {
         return when (resolutionParameters.singleCandidateResolutionMode) {
             SingleCandidateResolutionMode.CHECK_EXTENSION_FOR_COMPLETION -> CheckExtensionForCompletionCandidateInfoProvider(
                 resolutionParameters,

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

@@ -35,7 +35,7 @@ abstract class AbstractCandidateInfoProvider(
     protected val firSession: FirSession,
 ) : CandidateInfoProvider {
     override fun callInfo(): CallInfo = with(resolutionParameters) {
-        return CallInfo(
+        CallInfo(
             callKind = callKind(),
             name = callableSymbol.callableId.callableName,
             explicitReceiver = explicitReceiver,
@@ -81,7 +81,7 @@ class CheckExtensionForCompletionCandidateInfoProvider(
     // Passing them through can lead to false positives.
     override fun shouldFailBeforeResolve(): Boolean = with(resolutionParameters) {
         val callHasExtensionReceiver = explicitReceiverKind() == ExplicitReceiverKind.EXTENSION_RECEIVER
-                    || implicitExtensionReceiverValue() != null
+                || implicitExtensionReceiverValue() != null
         val candidateHasExtensionReceiver = callableSymbol.fir.receiverTypeRef != null
         callHasExtensionReceiver != candidateHasExtensionReceiver
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSession.kt

@@ -5,9 +5,6 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir
 
-import com.jetbrains.rd.util.getOrCreate
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
@@ -20,20 +17,10 @@ import org.jetbrains.kotlin.idea.frontend.api.components.KtDiagnosticProvider
 import org.jetbrains.kotlin.idea.frontend.api.components.KtSmartCastProvider
 import org.jetbrains.kotlin.idea.frontend.api.components.KtTypeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.components.*
-import org.jetbrains.kotlin.idea.fir.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolver
-import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolutionMode
-import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.ResolutionParameters
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFunctionSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbolProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.threadLocal
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolProvider
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.idea.frontend.api.symbols.*
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class KtFirAnalysisSession
 private constructor(
@@ -57,6 +44,7 @@ private constructor(
     override val scopeProvider by threadLocal { KtFirScopeProvider(this, firSymbolBuilder, project, firResolveState, token) }
     override val symbolProvider: KtSymbolProvider =
         KtFirSymbolProvider(this, firResolveState.firIdeLibrariesSession.firSymbolProvider, firResolveState, firSymbolBuilder, token)
+    override val completionCandidateChecker: KtCompletionCandidateChecker = KtFirCompletionCandidateChecker(this, token)
 
 
     override fun createContextDependentCopy(): KtAnalysisSession {
@@ -71,53 +59,6 @@ private constructor(
         )
     }
 
-    private val completionContextCache = HashMap<Pair<FirFile, KtNamedFunction>, LowLevelFirApiFacade.FirCompletionContext>()
-
-    // Temporary kludge to check single candidate resolver in tests
-    override fun resolveAndCheckReceivers(
-        firSymbolForCandidate: KtCallableSymbol,
-        originalFile: KtFile,
-        nameExpression: KtSimpleNameExpression,
-        possibleReceiver: KtExpression?,
-    ): Boolean {
-        firSymbolForCandidate.safeAs<KtFirFunctionSymbol>()?.firRef?.withFir { firFunction ->
-            val file = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
-            val explicitReceiverExpression = possibleReceiver?.getOrBuildFirOfType<FirExpression>(firResolveState)
-            val resolver = SingleCandidateResolver(firResolveState.firIdeSourcesSession, file)
-            val enclosingFunction = nameExpression.getNonStrictParentOfType<KtNamedFunction>()
-                ?: error("Cannot find enclosing function for ${nameExpression.getElementTextInContext()}")
-
-            val completionContext = completionContextCache.getOrCreate(file to enclosingFunction) {
-                LowLevelFirApiFacade.buildCompletionContextForFunction(
-                    file,
-                    enclosingFunction,
-                    enclosingFunction,
-                    state = firResolveState
-                )
-            }
-
-            val towerDataContext = completionContext.getTowerDataContext(nameExpression)
-
-            val implicitReceivers = sequence {
-                yield(null)
-                yieldAll(towerDataContext.implicitReceiverStack)
-            }
-            for (implicitReceiverValue in implicitReceivers) {
-                val resolutionParameters = ResolutionParameters(
-                    singleCandidateResolutionMode = SingleCandidateResolutionMode.CHECK_EXTENSION_FOR_COMPLETION,
-                    callableSymbol = firFunction.symbol,
-                    implicitReceiver = implicitReceiverValue,
-                    explicitReceiver = explicitReceiverExpression
-                )
-                resolver.resolveSingleCandidate(resolutionParameters)?.let {
-                    // not null if resolved and completed successfully
-                    return true
-                }
-            }
-        }
-        return false
-    }
-
     companion object {
         @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
         internal fun createForElement(element: KtElement): KtFirAnalysisSession {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2010-2020 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.idea.frontend.api.fir.components
+
+import com.intellij.psi.PsiElement
+import com.jetbrains.rd.util.getOrCreate
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
+import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.ResolutionParameters
+import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolutionMode
+import org.jetbrains.kotlin.idea.fir.low.level.api.resolver.SingleCandidateResolver
+import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.components.KtCompletionCandidateChecker
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFunctionSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPropertySymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import java.util.concurrent.ConcurrentHashMap
+
+internal class KtFirCompletionCandidateChecker(
+    override val analysisSession: KtFirAnalysisSession,
+    override val token: ValidityToken,
+) : KtCompletionCandidateChecker(), KtFirAnalysisSessionComponent {
+    private val completionContextCache = ConcurrentHashMap<Pair<FirFile, KtNamedFunction>, LowLevelFirApiFacade.FirCompletionContext>()
+
+    override fun checkExtensionFitsCandidate(
+        firSymbolForCandidate: KtCallableSymbol,
+        originalFile: KtFile,
+        originalPosition: PsiElement?,
+        nameExpression: KtSimpleNameExpression,
+        possibleExplicitReceiver: KtExpression?,
+    ): Boolean = withValidityAssertion {
+        val originalEnclosingFunction = originalPosition?.getNonStrictParentOfType<KtNamedFunction>()
+            ?: error("Cannot find enclosing function for completion in provided position (or position is absent)")
+
+        val functionFits = firSymbolForCandidate.withResolvedFirOfType<KtFirFunctionSymbol, FirSimpleFunction, Boolean> { firFunction ->
+            checkExtension(firFunction, originalFile, originalEnclosingFunction, nameExpression, possibleExplicitReceiver)
+        }
+        val propertyFits = firSymbolForCandidate.withResolvedFirOfType<KtFirPropertySymbol, FirProperty, Boolean> { firProperty ->
+            checkExtension(firProperty, originalFile, originalEnclosingFunction, nameExpression, possibleExplicitReceiver)
+        }
+
+        functionFits ?: propertyFits ?: false
+    }
+
+    private inline fun <reified T : KtFirSymbol<F>, F : FirDeclaration, R> KtCallableSymbol.withResolvedFirOfType(
+        noinline action: (F) -> R,
+    ): R? = this.safeAs<T>()?.firRef?.withFir(FirResolvePhase.BODY_RESOLVE, action)
+
+    private fun checkExtension(
+        candidateSymbol: FirCallableDeclaration<*>,
+        originalFile: KtFile,
+        originalEnclosingFunction: KtNamedFunction,
+        nameExpression: KtSimpleNameExpression,
+        possibleExplicitReceiver: KtExpression?,
+    ): Boolean {
+        val file = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
+        val explicitReceiverExpression = possibleExplicitReceiver?.getOrBuildFirOfType<FirExpression>(firResolveState)
+        val resolver = SingleCandidateResolver(firResolveState.firIdeSourcesSession, file)
+        val implicitReceivers = getImplicitReceivers(file, nameExpression, originalEnclosingFunction)
+
+        for (implicitReceiverValue in implicitReceivers) {
+            val resolutionParameters = ResolutionParameters(
+                singleCandidateResolutionMode = SingleCandidateResolutionMode.CHECK_EXTENSION_FOR_COMPLETION,
+                callableSymbol = candidateSymbol.symbol,
+                implicitReceiver = implicitReceiverValue,
+                explicitReceiver = explicitReceiverExpression
+            )
+            resolver.resolveSingleCandidate(resolutionParameters)?.let {
+                // not null if resolved and completed successfully
+                return true
+            }
+        }
+        return false
+    }
+
+    private fun getImplicitReceivers(
+        file: FirFile,
+        fakeNameExpression: KtElement,
+        originalEnclosingFunction: KtNamedFunction,
+    ): Sequence<ImplicitReceiverValue<*>?> {
+        val fakeEnclosingFunction = fakeNameExpression.getNonStrictParentOfType<KtNamedFunction>()
+            ?: error("Cannot find enclosing function for ${fakeNameExpression.getElementTextInContext()}")
+
+        val completionContext = completionContextCache.getOrCreate(file to fakeEnclosingFunction) {
+            LowLevelFirApiFacade.buildCompletionContextForFunction(
+                file,
+                fakeEnclosingFunction,
+                originalEnclosingFunction,
+                state = firResolveState
+            )
+        }
+
+        val towerDataContext = completionContext.getTowerDataContext(fakeNameExpression)
+
+        return sequence {
+            yield(null) // otherwise explicit receiver won't be checked when there are no implicit receivers in completion position
+            yieldAll(towerDataContext.implicitReceiverStack)
+        }
+    }
+}