[kotlin] Reuse applicability checker between extension candidates

Before, auto-completion was based on on-air resolve, and
'FirTowerDataContext' was computed once, during creation of the
dependent analysis session (with using 'FirTowerDataContextCollector').
This allowed to return the context almost instantly, no matter how many
completion candidates there were.

Dangling file resolution doesn't use 'FirTowerDataContextCollector', as
now there is more precise 'ContextCollector'. However,
'ContextCollector' doesn't cache results by itself, and, if asked, it
will compute the same context over and over again. This exact happened
with extension applicability checker.

The optimized implementation maintains a checker object which caches the
computed context. The old end-point is left to give the IDE plugin time
to adapt to the change.

^KTIJ-28445

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

@@ -7,11 +7,11 @@ package org.jetbrains.kotlin.analysis.api.descriptors.components
 
 import org.jetbrains.kotlin.analysis.api.components.KtCompletionCandidateChecker
 import org.jetbrains.kotlin.analysis.api.components.KtExtensionApplicabilityResult
+import org.jetbrains.kotlin.analysis.api.components.KtCompletionExtensionCandidateChecker
 import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
 import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnalysisSessionComponent
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -22,6 +22,15 @@ internal class KtFe10CompletionCandidateChecker(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
+    override fun createExtensionCandidateChecker(
+        originalFile: KtFile,
+        nameExpression: KtSimpleNameExpression,
+        explicitReceiver: KtExpression?
+    ): KtCompletionExtensionCandidateChecker {
+        throw NotImplementedError("Method is not implemented for FE 1.0")
+    }
+
+    @Suppress("OVERRIDE_DEPRECATION")
     override fun checkExtensionFitsCandidate(
         firSymbolForCandidate: KtCallableSymbol,
         originalFile: KtFile,

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

@@ -5,11 +5,15 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.components
 
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.components.KtCompletionCandidateChecker
 import org.jetbrains.kotlin.analysis.api.components.KtExtensionApplicabilityResult
+import org.jetbrains.kotlin.analysis.api.components.KtCompletionExtensionCandidateChecker
+import org.jetbrains.kotlin.analysis.api.components.KtExtensionApplicabilityResult.*
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
@@ -18,14 +22,16 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.SingleCandidateR
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.SingleCandidateResolver
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+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.declarations.FirVariable
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.fir.resolve.calls.FirErrorReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.receiverType
 import org.jetbrains.kotlin.psi.KtExpression
@@ -41,97 +47,147 @@ internal class KtFirCompletionCandidateChecker(
     override val analysisSession: KtFirAnalysisSession,
     override val token: KtLifetimeToken,
 ) : KtCompletionCandidateChecker(), KtFirAnalysisSessionComponent {
+    override fun createExtensionCandidateChecker(
+        originalFile: KtFile,
+        nameExpression: KtSimpleNameExpression,
+        explicitReceiver: KtExpression?
+    ): KtCompletionExtensionCandidateChecker = analysisSession.withValidityAssertion {
+        return LazyKtCompletionExtensionCandidateChecker {
+            // Double validity check is needed, as the checker may be requested some time later
+            analysisSession.withValidityAssertion {
+                KtFirCompletionExtensionCandidateChecker(analysisSession, nameExpression, explicitReceiver, originalFile)
+            }
+        }
+    }
+
+    @Suppress("OVERRIDE_DEPRECATION")
     override fun checkExtensionFitsCandidate(
         firSymbolForCandidate: KtCallableSymbol,
         originalFile: KtFile,
         nameExpression: KtSimpleNameExpression,
         possibleExplicitReceiver: KtExpression?,
-    ): KtExtensionApplicabilityResult {
+    ): KtExtensionApplicabilityResult = analysisSession.withValidityAssertion {
-        require(firSymbolForCandidate is KtFirSymbol<*>)
+        val checker = KtFirCompletionExtensionCandidateChecker(analysisSession, nameExpression, possibleExplicitReceiver, originalFile)
-        firSymbolForCandidate.firSymbol.lazyResolveToPhase(FirResolvePhase.STATUS)
+        return with(analysisSession) {
-        val declaration = firSymbolForCandidate.firSymbol.fir as FirCallableDeclaration
+            checker.computeApplicability(firSymbolForCandidate)
-        return checkExtension(declaration, originalFile, nameExpression, possibleExplicitReceiver)
-    }
-
-    private fun checkExtension(
-        candidateSymbol: FirCallableDeclaration,
-        originalFile: KtFile,
-        nameExpression: KtSimpleNameExpression,
-        possibleExplicitReceiver: KtExpression?,
-    ): KtExtensionApplicabilityResult {
-        val file = originalFile.getOrBuildFirFile(firResolveSession)
-        val explicitReceiverExpression = possibleExplicitReceiver?.getMatchingFirExpressionForCallReceiver()
-        val resolver = SingleCandidateResolver(firResolveSession.useSiteFirSession, file)
-        val implicitReceivers = getImplicitReceivers(nameExpression)
-        for (implicitReceiverValue in implicitReceivers) {
-            val resolutionParameters = ResolutionParameters(
-                singleCandidateResolutionMode = SingleCandidateResolutionMode.CHECK_EXTENSION_FOR_COMPLETION,
-                callableSymbol = candidateSymbol.symbol,
-                implicitReceiver = implicitReceiverValue,
-                explicitReceiver = explicitReceiverExpression,
-                allowUnsafeCall = true,
-                allowUnstableSmartCast = true,
-            )
-            resolver.resolveSingleCandidate(resolutionParameters)?.let { call ->
-                val substitutor = call.createSubstitutorFromTypeArguments() ?: return@let null
-                val receiverCastRequired = call.calleeReference is FirErrorReferenceWithCandidate
-
-                return when {
-                    candidateSymbol is FirVariable && candidateSymbol.symbol.resolvedReturnType.receiverType(rootModuleSession) != null -> {
-                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall(substitutor, receiverCastRequired, token)
-                    }
-                    else -> {
-                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable(substitutor, receiverCastRequired, token)
-                    }
-                }
-            }
         }
-        return KtExtensionApplicabilityResult.NonApplicable(token)
+    }
+}
+
+private class KtFirCompletionExtensionCandidateChecker(
+    override val analysisSession: KtFirAnalysisSession,
+    private val nameExpression: KtSimpleNameExpression,
+    explicitReceiver: KtExpression?,
+    originalFile: KtFile,
+) : KtCompletionExtensionCandidateChecker, KtFirAnalysisSessionComponent {
+    private val implicitReceivers: List<ImplicitReceiverValue<*>>
+    private val firCallSiteSession: FirSession
+    private val firOriginalFile: FirFile
+    private val firExplicitReceiver: FirExpression?
+
+    init {
+        val fakeFile = nameExpression.containingKtFile
+        val firFakeFile = fakeFile.getOrBuildFirFile(firResolveSession)
+
+        implicitReceivers = computeImplicitReceivers(firFakeFile)
+        firCallSiteSession = firFakeFile.llFirSession
+        firOriginalFile = originalFile.getOrBuildFirFile(firResolveSession)
+        firExplicitReceiver = explicitReceiver?.let(::findReceiverFirExpression)
     }
 
-    private fun getImplicitReceivers(fakeNameExpression: KtSimpleNameExpression): Sequence<ImplicitReceiverValue<*>?> {
+    context(KtAnalysisSession)
-        val fakeFile = fakeNameExpression.containingKtFile
+    override fun computeApplicability(candidate: KtCallableSymbol): KtExtensionApplicabilityResult {
-        val fakeFirFile = fakeFile.getOrBuildFirFile(firResolveSession)
+        require(candidate is KtFirSymbol<*>)
 
+        analysisSession.withValidityAssertion {
+            val firSymbol = candidate.firSymbol as FirCallableSymbol<*>
+            firSymbol.lazyResolveToPhase(FirResolvePhase.STATUS)
+
+            val resolver = SingleCandidateResolver(firCallSiteSession, firOriginalFile)
+            val token = analysisSession.token
+
+            fun processReceiver(implicitReceiverValue: ImplicitReceiverValue<*>?): KtExtensionApplicabilityResult? {
+                val resolutionParameters = ResolutionParameters(
+                    singleCandidateResolutionMode = SingleCandidateResolutionMode.CHECK_EXTENSION_FOR_COMPLETION,
+                    callableSymbol = firSymbol,
+                    implicitReceiver = implicitReceiverValue,
+                    explicitReceiver = firExplicitReceiver,
+                    allowUnsafeCall = true,
+                    allowUnstableSmartCast = true,
+                )
+
+                val firResolvedCall = resolver.resolveSingleCandidate(resolutionParameters) ?: return null
+                val substitutor = firResolvedCall.createSubstitutorFromTypeArguments() ?: return null
+
+                val receiverCastRequired = firResolvedCall.calleeReference is FirErrorReferenceWithCandidate
+
+                if (firSymbol is FirVariableSymbol<*> && firSymbol.resolvedReturnType.receiverType(firCallSiteSession) != null) {
+                    return ApplicableAsFunctionalVariableCall(substitutor, receiverCastRequired, token)
+                }
+
+                return ApplicableAsExtensionCallable(substitutor, receiverCastRequired, token)
+            }
+
+            return implicitReceivers.firstNotNullOfOrNull(::processReceiver)
+                ?: processReceiver(null)
+                ?: NonApplicable(token)
+        }
+    }
+
+    private fun computeImplicitReceivers(firFakeFile: FirFile): List<ImplicitReceiverValue<*>> {
         val sessionHolder = run {
-            val firSession = fakeFirFile.llFirSession
+            val firSession = firFakeFile.llFirSession
             val scopeSession = firResolveSession.getScopeSessionFor(firSession)
             SessionHolderImpl(firSession, scopeSession)
         }
 
-        val elementContext = ContextCollector.process(fakeFirFile, sessionHolder, fakeNameExpression, bodyElement = null)
+        val elementContext = ContextCollector.process(firFakeFile, sessionHolder, nameExpression, bodyElement = null)
 
         val towerDataContext = elementContext?.towerDataContext
-            ?: errorWithAttachment("Cannot find enclosing declaration for ${fakeNameExpression::class}") {
+            ?: errorWithAttachment("Cannot find enclosing declaration for ${nameExpression::class}") {
-                withPsiEntry("fakeNameExpression", fakeNameExpression)
+                withPsiEntry("fakeNameExpression", nameExpression)
             }
 
-        return sequence {
+        return buildList {
-            yield(null) // otherwise explicit receiver won't be checked when there are no implicit receivers in completion position
+            addAll(towerDataContext.implicitReceiverStack)
-            yieldAll(towerDataContext.implicitReceiverStack)
             for (towerDataElement in towerDataContext.towerDataElements) {
-                yieldAll(towerDataElement.contextReceiverGroup.orEmpty())
+                addAll(towerDataElement.contextReceiverGroup.orEmpty())
             }
         }
     }
 
     /**
-     * It is not enough to just call the `getOrBuildFirOfType` on [this] receiver expression, because for calls
+     * Returns a [FirExpression] matching the given PSI [receiverExpression].
-     * like `foo?.bar()` the receiver is additionally wrapped into `FirCheckedSafeCallSubject`, which is important
-     * for type-checks during resolve.
      *
-     * @receiver PSI receiver expression in some qualified expression (e.g. `foo` in `foo?.bar()`, `a` in `a.b`)
+     * @param receiverExpression a qualified expression receiver (e.g., `foo` in `foo?.bar()`, or in `foo.bar`).
-     * @return A FIR expression which most precisely represents the receiver for the corresponding FIR call.
+     *
+     * The function unwraps certain receiver expressions. For instance, for safe calls direct counterpart to a [KtSafeQualifiedExpression]
+     * is (FirCheckedSafeCallSubject)[org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject] which requires additional unwrapping
+     * to be used for call resolution.
      */
-    private fun KtExpression.getMatchingFirExpressionForCallReceiver(): FirExpression? {
+    private fun findReceiverFirExpression(receiverExpression: KtExpression): FirExpression? {
-        // FIR for KtStatementExpression is not FirExpression
+        if (receiverExpression is KtStatementExpression) {
-        if (this is KtStatementExpression) {
+            // FIR for 'KtStatementExpression' is not a 'FirExpression'
             return null
         }
-        val psiWholeCall = this.getQualifiedExpressionForReceiver()
-        if (psiWholeCall !is KtSafeQualifiedExpression) return this.getOrBuildFirOfType<FirExpression>(firResolveSession)
 
-        val firSafeCall = psiWholeCall.getOrBuildFirOfType<FirSafeCallExpression>(firResolveSession)
+        val parentCall = receiverExpression.getQualifiedExpressionForReceiver()
+        if (parentCall !is KtSafeQualifiedExpression) {
+            return receiverExpression.getOrBuildFirOfType<FirExpression>(firResolveSession)
+        }
+
+        val firSafeCall = parentCall.getOrBuildFirOfType<FirSafeCallExpression>(firResolveSession)
         return firSafeCall.checkedSubjectRef.value
     }
 }
+
+private class LazyKtCompletionExtensionCandidateChecker(
+    delegateFactory: () -> KtCompletionExtensionCandidateChecker
+) : KtCompletionExtensionCandidateChecker {
+    private val delegate: KtCompletionExtensionCandidateChecker by lazy(delegateFactory)
+
+    context(KtAnalysisSession)
+    override fun computeApplicability(candidate: KtCallableSymbol): KtExtensionApplicabilityResult {
+        return delegate.computeApplicability(candidate)
+    }
+}

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
@@ -15,6 +16,13 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 
 public abstract class KtCompletionCandidateChecker : KtAnalysisSessionComponent() {
+    public abstract fun createExtensionCandidateChecker(
+        originalFile: KtFile,
+        nameExpression: KtSimpleNameExpression,
+        explicitReceiver: KtExpression?
+    ): KtCompletionExtensionCandidateChecker
+
+    @Deprecated("Use createExtensionCandidateChecker() instead.")
     public abstract fun checkExtensionFitsCandidate(
         firSymbolForCandidate: KtCallableSymbol,
         originalFile: KtFile,
@@ -23,6 +31,10 @@ public abstract class KtCompletionCandidateChecker : KtAnalysisSessionComponent(
     ): KtExtensionApplicabilityResult
 }
 
+public interface KtCompletionExtensionCandidateChecker {
+    context(KtAnalysisSession)
+    public fun computeApplicability(candidate: KtCallableSymbol): KtExtensionApplicabilityResult
+}
 
 public sealed class KtExtensionApplicabilityResult : KtLifetimeOwner {
     public sealed class Applicable : KtExtensionApplicabilityResult() {
@@ -54,11 +66,36 @@ public sealed class KtExtensionApplicabilityResult : KtLifetimeOwner {
 }
 
 public interface KtCompletionCandidateCheckerMixIn : KtAnalysisSessionMixIn {
+    /**
+     * Returns an extension applicability checker for the given context [nameExpression].
+     * The function is meant to only be used for providing auto-completion for Kotlin in IntelliJ IDEA.
+     *
+     * The returned checker does not cache the results for individual callable candidates.
+     *
+     * @param originalFile The file being edited.
+     * @param nameExpression The expression under the caret in an in-memory copy of [originalFile]
+     *     with a dummy identifier inserted. Also see `CompletionUtilCore.DUMMY_IDENTIFIER` in IntelliJ IDEA.
+     * @param explicitReceiver A receiver expression, if available (also from the in-memory copy of [originalFile]).
+     */
+    public fun createExtensionCandidateChecker(
+        originalFile: KtFile,
+        nameExpression: KtSimpleNameExpression,
+        explicitReceiver: KtExpression?
+    ): KtCompletionExtensionCandidateChecker {
+        return analysisSession.completionCandidateChecker.createExtensionCandidateChecker(
+            originalFile,
+            nameExpression,
+            explicitReceiver
+        )
+    }
+
+    @Deprecated("Use createExtensionCandidateChecker() instead.")
     public fun KtCallableSymbol.checkExtensionIsSuitable(
         originalPsiFile: KtFile,
         psiFakeCompletionExpression: KtSimpleNameExpression,
         psiReceiverExpression: KtExpression?,
     ): KtExtensionApplicabilityResult = withValidityAssertion {
+        @Suppress("DEPRECATION")
         analysisSession.completionCandidateChecker.checkExtensionFitsCandidate(
             this,
             originalPsiFile,