[FIR] Allow nullable receiver type in case of safe-calls

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -32,6 +32,7 @@ fun resolveArgumentExpression(
     expectedTypeRef: FirTypeRef,
     sink: CheckerSink,
     isReceiver: Boolean,
+    isSafeCall: Boolean,
     acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
@@ -42,6 +43,7 @@ fun resolveArgumentExpression(
             expectedType,
             sink,
             isReceiver,
+            isSafeCall,
             typeProvider
         )
         // TODO:!
@@ -57,6 +59,7 @@ fun resolveArgumentExpression(
             expectedTypeRef,
             sink,
             isReceiver,
+            isSafeCall,
             acceptLambdaAtoms,
             typeProvider
         )
@@ -67,10 +70,11 @@ fun resolveArgumentExpression(
             expectedTypeRef,
             sink,
             isReceiver,
+            isSafeCall,
             acceptLambdaAtoms,
             typeProvider
         )
-        else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, typeProvider)
+        else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isSafeCall, typeProvider)
     }
 }
 
@@ -80,11 +84,12 @@ fun resolvePlainExpressionArgument(
     expectedType: ConeKotlinType?,
     sink: CheckerSink,
     isReceiver: Boolean,
+    isSafeCall: Boolean,
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
     if (expectedType == null) return
     val argumentType = typeProvider(argument)?.coneTypeSafe<ConeKotlinType>() ?: return
-    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver)
+    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isSafeCall)
 }
 
 fun resolvePlainArgumentType(
@@ -92,12 +97,20 @@ fun resolvePlainArgumentType(
     argumentType: ConeKotlinType,
     expectedType: ConeKotlinType,
     sink: CheckerSink,
-    isReceiver: Boolean
+    isReceiver: Boolean,
+    isSafeCall: Boolean
 ) {
     val position = SimpleConstraintSystemConstraintPosition //TODO
 
+    val nullableExpectedType = expectedType.withNullability(ConeNullability.NULLABLE)
+    if (isReceiver && isSafeCall) {
+        if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, nullableExpectedType, position)) {
+            sink.reportApplicability(CandidateApplicability.WRONG_RECEIVER) // TODO
+        }
+        return
+    }
+
     if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, expectedType, position)) {
-        val nullableExpectedType = expectedType.withNullability(ConeNullability.NULLABLE)
         if (!isReceiver) {
             if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, nullableExpectedType, position)) {
                 csBuilder.addSubtypeConstraint(argumentType, expectedType, position)
@@ -119,6 +132,7 @@ internal fun Candidate.resolveArgument(
     argument: FirExpression,
     parameter: FirValueParameter,
     isReceiver: Boolean,
+    isSafeCall: Boolean,
     typeProvider: (FirExpression) -> FirTypeRef?,
     sink: CheckerSink
 ) {
@@ -131,6 +145,7 @@ internal fun Candidate.resolveArgument(
         parameter.returnTypeRef,
         sink,
         isReceiver,
+        isSafeCall,
         { this.postponedAtoms += it },
         typeProvider
     )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CallResolver.kt

@@ -36,6 +36,7 @@ class CallInfo(
 
     val explicitReceiver: FirExpression?,
     val arguments: List<FirExpression>,
+    val isSafeCall: Boolean,
 
     val typeArguments: List<FirTypeProjection>,
     val session: FirSession,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArgumentsAnalyzer.kt

@@ -109,6 +109,7 @@ class PostponedArgumentsAnalyzer(
                 lambda.atom.returnTypeRef, // TODO: proper ref
                 checkerSink,
                 false,
+                false,
                 { atom = it },
                 typeProvider
             )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolverParts.kt

@@ -103,10 +103,15 @@ internal sealed class CheckReceivers : ResolutionStage() {
         if (receiverParameterValue != null) {
             if (explicitReceiverExpression != null && explicitReceiverKind.shouldBeResolvedAsExplicit()) {
                 resolveArgumentExpression(
-                    candidate.csBuilder, explicitReceiverExpression,
+                    candidate.csBuilder,
+                    explicitReceiverExpression,
                     candidate.substitutor.substituteOrSelf(receiverParameterValue.type),
                     explicitReceiverExpression.typeRef,
-                    sink, isReceiver = true, typeProvider = callInfo.typeProvider, acceptLambdaAtoms = { candidate.postponedAtoms += it }
+                    sink,
+                    isReceiver = true,
+                    isSafeCall = callInfo.isSafeCall,
+                    typeProvider = callInfo.typeProvider,
+                    acceptLambdaAtoms = { candidate.postponedAtoms += it }
                 )
             }
         }
@@ -132,7 +137,14 @@ internal object CheckArguments : CheckerStage() {
         val argumentMapping =
             candidate.argumentMapping ?: throw IllegalStateException("Argument should be already mapped while checking arguments!")
         for ((argument, parameter) in argumentMapping) {
-            candidate.resolveArgument(argument, parameter, isReceiver = false, typeProvider = callInfo.typeProvider, sink = sink)
+            candidate.resolveArgument(
+                argument,
+                parameter,
+                isReceiver = false,
+                isSafeCall = false,
+                typeProvider = callInfo.typeProvider,
+                sink = sink
+            )
         }
 
         if (candidate.system.hasContradiction) {

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

@@ -248,7 +248,16 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
 
         val receiver = qualifiedAccess.explicitReceiver?.transformSingle(this, noExpectedType)
 
-        val info = CallInfo(CallKind.VariableAccess, receiver, emptyList(), emptyList(), session, file, container!!) { it.resultType }
+        val info = CallInfo(
+            CallKind.VariableAccess,
+            receiver,
+            emptyList(),
+            qualifiedAccess.safe,
+            emptyList(),
+            session,
+            file,
+            container!!
+        ) { it.resultType }
         val resolver = CallResolver(jump, inferenceComponents)
         resolver.callInfo = info
         resolver.scopes = (scopes + localScopes).asReversed()
@@ -358,7 +367,16 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         val arguments = functionCall.arguments
         val typeArguments = functionCall.typeArguments
 
-        val info = CallInfo(CallKind.Function, explicitReceiver, arguments, typeArguments, session, file, container!!) { it.resultType }
+        val info = CallInfo(
+            CallKind.Function,
+            explicitReceiver,
+            arguments,
+            functionCall.safe,
+            typeArguments,
+            session,
+            file,
+            container!!
+        ) { it.resultType }
         val resolver = CallResolver(jump, inferenceComponents)
         resolver.callInfo = info
         resolver.scopes = (scopes + localScopes).asReversed()