[FIR] Prefer the non-smartcasted receiver in resolution if it's visible

Previously, we would ignore the candidate with the non-smartcasted
receiver if they have the same symbol.
Now we prefer them when they're visible or when the smart cast changes
the nullability.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerLevelHandler.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.resolve.calls.CandidateFactory
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 
 internal class CandidateFactoriesAndCollectors(
     // Common calls
@@ -85,8 +86,8 @@ private class TowerScopeLevelProcessor(
         scope: FirScope,
         objectsByName: Boolean,
         isFromOriginalTypeInPresenceOfSmartCast: Boolean,
-    ) {
-        resultCollector.consumeCandidate(
+    ): CandidateApplicability {
+        return resultCollector.consumeCandidate(
             group, candidateFactory.createCandidate(
                 callInfo,
                 symbol,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerLevels.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitAnyTypeRef
 import org.jetbrains.kotlin.fir.utils.exceptions.withConeTypeEntry
 import org.jetbrains.kotlin.name.StandardClassIds.Annotations.HidesMembers
+import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
@@ -59,7 +60,7 @@ abstract class TowerScopeLevel {
             scope: FirScope,
             objectsByName: Boolean = false,
             isFromOriginalTypeInPresenceOfSmartCast: Boolean = false,
-        )
+        ): CandidateApplicability
     }
 }
 
@@ -86,14 +87,13 @@ class MemberScopeTowerLevel(
         val scope = dispatchReceiverValue.scope(session, scopeSession) ?: return ProcessResult.SCOPE_EMPTY
         var (empty, candidates) = scope.collectCandidates(processScopeMembers)
 
-        val scopeWithoutSmartcast = getOriginalReceiverExpressionIfStableSmartCast()
-            ?.resolvedType
-            ?.scope(
-                session,
-                scopeSession,
-                bodyResolveComponents.returnTypeCalculator.callableCopyTypeCalculator,
-                requiredMembersPhase = FirResolvePhase.STATUS,
-            )
+        val receiverTypeWithoutSmartCast = getOriginalReceiverExpressionIfStableSmartCast()?.resolvedType
+        val scopeWithoutSmartcast = receiverTypeWithoutSmartCast?.scope(
+            session,
+            scopeSession,
+            bodyResolveComponents.returnTypeCalculator.callableCopyTypeCalculator,
+            requiredMembersPhase = FirResolvePhase.STATUS,
+        )
 
         if (scopeWithoutSmartcast == null) {
             consumeCandidates(
@@ -107,12 +107,22 @@ class MemberScopeTowerLevel(
             scopeWithoutSmartcast.collectCandidates(processScopeMembers).let { (isEmpty, originalCandidates) ->
                 empty = empty && isEmpty
                 for (originalCandidate in originalCandidates) {
-                    map[originalCandidate.member] = MemberFromSmartcastScope(originalCandidate, cameFromSmartcast = false)
+                    map[originalCandidate.member] = MemberFromSmartcastScope(originalCandidate, DispatchReceiverToUse.UnwrapSmartcast)
                 }
             }
 
             for (candidateFromSmartCast in candidates) {
-                map[candidateFromSmartCast.member] = MemberFromSmartcastScope(candidateFromSmartCast, cameFromSmartcast = true)
+                val existing = map[candidateFromSmartCast.member]
+
+                // If both scopes return the same symbol, we want to prefer the candidate from the original scope without smartcast
+                // with two exceptions:
+                // - When the smart-casted type is always null, we want to return it and report UNSAFE_CALL.
+                // - When the original type can be null, in this case the smart-case either makes it not-null or the call is red anyway.
+                if (existing == null || dispatchReceiverValue.type.isNullableNothing || receiverTypeWithoutSmartCast.canBeNull) {
+                    map[candidateFromSmartCast.member] = MemberFromSmartcastScope(candidateFromSmartCast, DispatchReceiverToUse.SmartcastWithoutUnwrapping)
+                } else {
+                    existing.dispatchReceiverToUse = DispatchReceiverToUse.SmartcastIfUnwrappedInvisible
+                }
             }
 
             consumeCandidates(
@@ -170,9 +180,15 @@ class MemberScopeTowerLevel(
         return if (empty) ProcessResult.SCOPE_EMPTY else ProcessResult.FOUND
     }
 
-    private data class MemberFromSmartcastScope<T : FirCallableSymbol<*>>(
+    private enum class DispatchReceiverToUse(val unwrapSmartcast: Boolean) {
+        UnwrapSmartcast(true),
+        SmartcastWithoutUnwrapping(false),
+        SmartcastIfUnwrappedInvisible(true),
+    }
+
+    private class MemberFromSmartcastScope<T : FirCallableSymbol<*>>(
         val memberWithBaseScope: MemberWithBaseScope<T>,
-        val cameFromSmartcast: Boolean
+        var dispatchReceiverToUse: DispatchReceiverToUse,
     )
 
     private fun <T : FirCallableSymbol<*>> FirTypeScope.collectCandidates(
@@ -208,13 +224,11 @@ class MemberScopeTowerLevel(
             ?: candidatesWithSmartcast?.values?.map { it.memberWithBaseScope }
             ?: error("candidatesWithoutSmartcast or candidatesWithSmartcast should be not null")
 
-        for (candidateWithScope in candidates) {
-            val (candidate, scope) = candidateWithScope
+        for ((candidate, scope) in candidates) {
             if (candidate.hasConsistentExtensionReceiver(givenExtensionReceiverOptions)) {
-                val isFromOriginalTypeInPresenceOfSmartCast = candidatesWithSmartcast != null &&
-                        !candidatesWithSmartcast.getValue(candidateWithScope.member).cameFromSmartcast
-
-                val dispatchReceiverToUse = when {
+                val dispatchReceiverToUse = candidatesWithSmartcast?.getValue(candidate)?.dispatchReceiverToUse
+                val isFromOriginalTypeInPresenceOfSmartCast = dispatchReceiverToUse?.unwrapSmartcast == true
+                val dispatchReceiver = when {
                     isFromOriginalTypeInPresenceOfSmartCast ->
                         getOriginalReceiverExpressionIfStableSmartCast()
                     // For a chain inference stub in dispatch receiver, we have to provide an explicit cast to Any
@@ -229,13 +243,23 @@ class MemberScopeTowerLevel(
                     else -> dispatchReceiverValue.receiverExpression
                 }
 
-                output.consumeCandidate(
+                val applicability = output.consumeCandidate(
                     candidate,
-                    dispatchReceiverToUse,
+                    dispatchReceiver,
                     givenExtensionReceiverOptions,
                     scope,
                     isFromOriginalTypeInPresenceOfSmartCast = isFromOriginalTypeInPresenceOfSmartCast
                 )
+
+                if (applicability == CandidateApplicability.K2_VISIBILITY_ERROR && dispatchReceiverToUse == DispatchReceiverToUse.SmartcastIfUnwrappedInvisible) {
+                    output.consumeCandidate(
+                        candidate,
+                        dispatchReceiverValue.receiverExpression,
+                        givenExtensionReceiverOptions,
+                        scope,
+                        isFromOriginalTypeInPresenceOfSmartCast = false
+                    )
+                }
             }
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.resolve.calls.inference.model.InferredEmptyIntersection
+import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.Variance
@@ -140,6 +141,18 @@ class FirCallCompletionResultsWriterTransformer(
             replaceCalleeReference(calleeReference.toResolvedReference())
             replaceDispatchReceiver(dispatchReceiver)
             replaceExtensionReceiver(extensionReceiver)
+
+            // If the explicit receiver is a smartcast expression, we can choose to use the unwrapped expression as dispatch receiver.
+            // To maintain the invariant
+            // explicitReceiver != null => explicitReceiver == dispatchReceiver || explicitReceiver == extensionReceiver
+            // we update the explicit receiver here.
+            // We only do this if the candidate is successful, otherwise, we can lose the explicit receiver node in red code like
+            // fun f(s: String, action: (String.() -> Unit)?) {
+            //    s.action?.let { it() }
+            //}
+            if (subCandidate.explicitReceiverKind == ExplicitReceiverKind.DISPATCH_RECEIVER && subCandidate.applicability.isSuccess) {
+                replaceExplicitReceiver(dispatchReceiver)
+            }
         }
 
         qualifiedAccessExpression.replaceContextReceiverArguments(subCandidate.contextReceiverArguments())