[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
This commit is contained in:
Yan Zhulanow
2024-01-23 17:32:15 +09:00
committed by Space Team
parent 4f521aa05c
commit 57dcd645b1
3 changed files with 167 additions and 65 deletions
@@ -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,
@@ -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 {
require(firSymbolForCandidate is KtFirSymbol<*>)
firSymbolForCandidate.firSymbol.lazyResolveToPhase(FirResolvePhase.STATUS)
val declaration = firSymbolForCandidate.firSymbol.fir as FirCallableDeclaration
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)
}
}
}
): KtExtensionApplicabilityResult = analysisSession.withValidityAssertion {
val checker = KtFirCompletionExtensionCandidateChecker(analysisSession, nameExpression, possibleExplicitReceiver, originalFile)
return with(analysisSession) {
checker.computeApplicability(firSymbolForCandidate)
}
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<*>?> {
val fakeFile = fakeNameExpression.containingKtFile
val fakeFirFile = fakeFile.getOrBuildFirFile(firResolveSession)
context(KtAnalysisSession)
override fun computeApplicability(candidate: KtCallableSymbol): KtExtensionApplicabilityResult {
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}") {
withPsiEntry("fakeNameExpression", fakeNameExpression)
?: errorWithAttachment("Cannot find enclosing declaration for ${nameExpression::class}") {
withPsiEntry("fakeNameExpression", nameExpression)
}
return sequence {
yield(null) // otherwise explicit receiver won't be checked when there are no implicit receivers in completion position
yieldAll(towerDataContext.implicitReceiverStack)
return buildList {
addAll(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
* like `foo?.bar()` the receiver is additionally wrapped into `FirCheckedSafeCallSubject`, which is important
* for type-checks during resolve.
* Returns a [FirExpression] matching the given PSI [receiverExpression].
*
* @receiver PSI receiver expression in some qualified expression (e.g. `foo` in `foo?.bar()`, `a` in `a.b`)
* @return A FIR expression which most precisely represents the receiver for the corresponding FIR call.
* @param receiverExpression a qualified expression receiver (e.g., `foo` in `foo?.bar()`, or in `foo.bar`).
*
* 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? {
// FIR for KtStatementExpression is not FirExpression
if (this is KtStatementExpression) {
private fun findReceiverFirExpression(receiverExpression: KtExpression): FirExpression? {
if (receiverExpression 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)
}
}
@@ -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,