[Analysis API] consider applicable extensions that require receiver cast

^KTIJ-23715

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

@@ -11,7 +11,6 @@ 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.symbols.KtCallableSymbol
-import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.ResolutionParameters
@@ -23,6 +22,7 @@ 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.calls.FirErrorReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.coneType
@@ -64,21 +64,25 @@ internal class KtFirCompletionCandidateChecker(
                 singleCandidateResolutionMode = SingleCandidateResolutionMode.CHECK_EXTENSION_FOR_COMPLETION,
                 callableSymbol = candidateSymbol.symbol,
                 implicitReceiver = implicitReceiverValue,
-                explicitReceiver = explicitReceiverExpression
+                explicitReceiver = explicitReceiverExpression,
+                allowUnsafeCall = true,
+                allowUnstableSmartCast = true,
             )
-            resolver.resolveSingleCandidate(resolutionParameters)?.let {
-                val substitutor = it.createSubstitutorFromTypeArguments() ?: return@let null
+            resolver.resolveSingleCandidate(resolutionParameters)?.let { call ->
+                val substitutor = call.createSubstitutorFromTypeArguments() ?: return@let null
+                val receiverCastRequired = call.calleeReference is FirErrorReferenceWithCandidate
+
                 return when {
                     candidateSymbol is FirVariable && candidateSymbol.returnTypeRef.coneType.receiverType(rootModuleSession) != null -> {
-                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall(substitutor, token)
+                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall(substitutor, receiverCastRequired, token)
                     }
                     else -> {
-                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable(substitutor, token)
+                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable(substitutor, receiverCastRequired, token)
                     }
                 }
             }
         }
-        return KtExtensionApplicabilityResult.NonApplicable(KtSubstitutor.Empty(token), token)
+        return KtExtensionApplicabilityResult.NonApplicable(token)
     }
 
     private fun getImplicitReceivers(

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

@@ -25,33 +25,32 @@ public abstract class KtCompletionCandidateChecker : KtAnalysisSessionComponent(
 
 
 public sealed class KtExtensionApplicabilityResult : KtLifetimeOwner {
-    public abstract val isApplicable: Boolean
-    public abstract val substitutor: KtSubstitutor
+    public sealed class Applicable : KtExtensionApplicabilityResult() {
+        public abstract val receiverCastRequired: Boolean
+        public abstract val substitutor: KtSubstitutor
+    }
 
     public class ApplicableAsExtensionCallable(
         private val _substitutor: KtSubstitutor,
+        private val _receiverCastRequired: Boolean,
         override val token: KtLifetimeToken
-    ) : KtExtensionApplicabilityResult() {
+    ) : Applicable() {
         override val substitutor: KtSubstitutor = withValidityAssertion { _substitutor }
-        override val isApplicable: Boolean get() = withValidityAssertion { true }
+        override val receiverCastRequired: Boolean get() = withValidityAssertion { _receiverCastRequired }
     }
 
     public class ApplicableAsFunctionalVariableCall(
         private val _substitutor: KtSubstitutor,
+        private val _receiverCastRequired: Boolean,
         override val token: KtLifetimeToken
-    ) : KtExtensionApplicabilityResult() {
+    ) : Applicable() {
         override val substitutor: KtSubstitutor get() = withValidityAssertion { _substitutor }
-        override val isApplicable: Boolean get() = withValidityAssertion { true }
+        override val receiverCastRequired: Boolean get() = withValidityAssertion { _receiverCastRequired }
     }
 
-
     public class NonApplicable(
-        private val _substitutor: KtSubstitutor,
         override val token: KtLifetimeToken
-    ) : KtExtensionApplicabilityResult() {
-        override val substitutor: KtSubstitutor = withValidityAssertion { _substitutor }
-        override val isApplicable: Boolean get() = withValidityAssertion { false }
-    }
+    ) : KtExtensionApplicabilityResult()
 }
 
 public interface KtCompletionCandidateCheckerMixIn : KtAnalysisSessionMixIn {

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

@@ -12,12 +12,13 @@ import org.jetbrains.kotlin.fir.expressions.FirEmptyArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 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.createConeDiagnosticForCandidateWithError
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 
 class SingleCandidateResolver(
@@ -61,10 +62,22 @@ class SingleCandidateResolver(
         )
 
         val applicability = resolutionStageRunner.processCandidate(candidate, resolutionContext, stopOnFirstError = true)
-        if (applicability.isSuccess) {
-            return completeResolvedCandidate(candidate, resolutionParameters)
+
+        val fakeCall = if (applicability.isSuccess) {
+            buildCallForResolvedCandidate(candidate, resolutionParameters)
+        } else if (
+            resolutionParameters.allowUnsafeCall && applicability == CandidateApplicability.UNSAFE_CALL ||
+            resolutionParameters.allowUnstableSmartCast && applicability == CandidateApplicability.UNSTABLE_SMARTCAST
+        ) {
+            buildCallForCandidateWithError(candidate, applicability, resolutionParameters)
+        } else {
+            return null
         }
-        return null
+
+        val expectedType = resolutionParameters.expectedType ?: bodyResolveComponents.noExpectedType
+        val completionResult = firCallCompleter.completeCall(fakeCall, expectedType)
+
+        return completionResult.takeIf { it.callCompleted }?.result
     }
 
     private fun createCandidateInfoProvider(resolutionParameters: ResolutionParameters): CandidateInfoProvider {
@@ -77,22 +90,32 @@ class SingleCandidateResolver(
         }
     }
 
-    private fun completeResolvedCandidate(candidate: Candidate, resolutionParameters: ResolutionParameters): FirFunctionCall? {
-        val fakeCall = buildFunctionCall {
+    private fun buildCallForResolvedCandidate(candidate: Candidate, resolutionParameters: ResolutionParameters): FirFunctionCall =
+        buildFunctionCall {
             calleeReference = FirNamedReferenceWithCandidate(
                 source = null,
                 name = resolutionParameters.callableSymbol.callableId.callableName,
                 candidate = candidate
             )
         }
-        val expectedType = resolutionParameters.expectedType ?: bodyResolveComponents.noExpectedType
-        val completionResult = firCallCompleter.completeCall(fakeCall, expectedType)
-        return if (completionResult.callCompleted) {
-            completionResult.result
-        } else null
+
+    private fun buildCallForCandidateWithError(
+        candidate: Candidate,
+        applicability: CandidateApplicability,
+        resolutionParameters: ResolutionParameters
+    ): FirFunctionCall {
+        val diagnostic = createConeDiagnosticForCandidateWithError(applicability, candidate)
+        val name = resolutionParameters.callableSymbol.callableId.callableName
+        return buildFunctionCall {
+            calleeReference = FirErrorReferenceWithCandidate(source = null, name, candidate, diagnostic)
+        }
     }
 }
 
+/**
+ * @param allowUnsafeCall if true, then candidate is resolved even if receiver's nullability doesn't match
+ * @param allowUnstableSmartCast if true, then candidate is resolved even if it requires unstable smart cast
+ */
 class ResolutionParameters(
     val singleCandidateResolutionMode: SingleCandidateResolutionMode,
     val callableSymbol: FirCallableSymbol<*>,
@@ -101,6 +124,8 @@ class ResolutionParameters(
     val explicitReceiver: FirExpression? = null,
     val argumentList: FirArgumentList = FirEmptyArgumentList,
     val typeArgumentList: List<FirTypeProjection> = emptyList(),
+    val allowUnsafeCall: Boolean = false,
+    val allowUnstableSmartCast: Boolean = false,
 )
 
 enum class SingleCandidateResolutionMode {