FIR: Support safe-calls new format in body resolution

^KT-38444 In Progress

compiler/fir/analysis-tests/testData/resolve/callResolution/extensionInvokeAfterSafeCall.kt

+interface A
+
+fun test_1(a: A?, convert: A.() -> String) {
+    val s = a?.convert()
+}
+
+fun test_2(a: A, convert: A.() -> String) {
+    val s = a.convert()
+}

compiler/fir/analysis-tests/testData/resolve/callResolution/extensionInvokeAfterSafeCall.txt

+FILE: extensionInvokeAfterSafeCall.kt
+    public abstract interface A : R|kotlin/Any| {
+    }
+    public final fun test_1(a: R|A?|, convert: R|A.() -> kotlin/String|): R|kotlin/Unit| {
+        lval s: R|kotlin/String?| = R|<local>/a|?.{ R|<local>/convert|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/String|>|($subj$) }
+    }
+    public final fun test_2(a: R|A|, convert: R|A.() -> kotlin/String|): R|kotlin/Unit| {
+        lval s: R|kotlin/String| = R|<local>/convert|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/String|>|(R|<local>/a|)
+    }

compiler/fir/analysis-tests/testData/resolve/callResolution/safeCallOnTypeAlias.kt

+typealias MyTypeAlias = (() -> String?)?
+fun foo(x: MyTypeAlias) {
+
+    x?.let { y -> y()?.let { result -> bar(result) } }
+}
+
+fun bar(x: String) = x

compiler/fir/analysis-tests/testData/resolve/callResolution/safeCallOnTypeAlias.txt

+FILE: safeCallOnTypeAlias.kt
+    public final typealias MyTypeAlias = R|() -> kotlin/String?|
+    public final fun foo(x: R|MyTypeAlias|): R|kotlin/Unit| {
+        R|<local>/x|?.{ $subj$.R|kotlin/let|<R|() -> kotlin/String?|, R|kotlin/String|>(<L> = let@fun <anonymous>(y: R|() -> kotlin/String?|): R|kotlin/String| <kind=EXACTLY_ONCE>  {
+            ^ R|<local>/y|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/String?|>|()?.{ $subj$.R|kotlin/let|<R|kotlin/String|, R|kotlin/String|>(<L> = let@fun <anonymous>(result: R|kotlin/String|): R|kotlin/String| <kind=EXACTLY_ONCE>  {
+                ^ R|/bar|(R|<local>/result|)
+            }
+            ) }
+        }
+        ) }
+    }
+    public final fun bar(x: R|kotlin/String|): R|kotlin/String| {
+        ^bar R|<local>/x|
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -604,6 +604,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/debugInfoCall.kt");
         }
 
+        @TestMetadata("extensionInvokeAfterSafeCall.kt")
+        public void testExtensionInvokeAfterSafeCall() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/extensionInvokeAfterSafeCall.kt");
+        }
+
         @TestMetadata("invokeAmbiguity.kt")
         public void testInvokeAmbiguity() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/invokeAmbiguity.kt");
@@ -613,6 +618,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         public void testObjectInvoke() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/objectInvoke.kt");
         }
+
+        @TestMetadata("safeCallOnTypeAlias.kt")
+        public void testSafeCallOnTypeAlias() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/safeCallOnTypeAlias.kt");
+        }
     }
 
     @TestMetadata("compiler/fir/analysis-tests/testData/resolve/cfg")

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -604,6 +604,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/debugInfoCall.kt");
         }
 
+        @TestMetadata("extensionInvokeAfterSafeCall.kt")
+        public void testExtensionInvokeAfterSafeCall() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/extensionInvokeAfterSafeCall.kt");
+        }
+
         @TestMetadata("invokeAmbiguity.kt")
         public void testInvokeAmbiguity() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/invokeAmbiguity.kt");
@@ -613,6 +618,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         public void testObjectInvoke() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/objectInvoke.kt");
         }
+
+        @TestMetadata("safeCallOnTypeAlias.kt")
+        public void testSafeCallOnTypeAlias() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/safeCallOnTypeAlias.kt");
+        }
     }
 
     @TestMetadata("compiler/fir/analysis-tests/testData/resolve/cfg")

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -55,7 +55,11 @@ class Fir2IrVisitor(
     fakeOverrideMode: FakeOverrideMode
 ) : Fir2IrComponents by components, FirDefaultVisitor<IrElement, Any?>(), IrGeneratorContextInterface {
 
-    private val integerApproximator = IntegerLiteralTypeApproximationTransformer(session.firSymbolProvider, session.inferenceContext)
+    private val integerApproximator = IntegerLiteralTypeApproximationTransformer(
+        session.firSymbolProvider,
+        session.inferenceContext,
+        session
+    )
 
     private val conversionScope = Fir2IrConversionScope()
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -439,4 +439,38 @@ fun CallableId.isIterator() =
 fun FirAnnotationCall.fqName(session: FirSession): FqName? {
     val symbol = session.firSymbolProvider.getSymbolByTypeRef<FirRegularClassSymbol>(annotationTypeRef) ?: return null
     return symbol.classId.asSingleFqName()
-}
\ No newline at end of file
+}
+
+fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(nullableReceiverExpression: FirExpression, session: FirSession) {
+    val receiverType = nullableReceiverExpression.typeRef.coneTypeSafe<ConeKotlinType>() ?: return
+
+    val expandedReceiverType = if (receiverType is ConeClassLikeType) receiverType.fullyExpandedType(session) else receiverType
+
+    replaceTypeRef(typeRef.resolvedTypeFromPrototype(expandedReceiverType.makeConeTypeDefinitelyNotNullOrNotNull()))
+}
+
+fun FirSafeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(
+    nullableReceiverExpression: FirExpression,
+    session: FirSession,
+) {
+    val receiverType = nullableReceiverExpression.typeRef.coneTypeSafe<ConeKotlinType>()
+    val typeAfterNullCheck = regularQualifiedAccess.expressionTypeOrUnitForAssignment() ?: return
+    val isReceiverActuallyNullable = receiverType != null && session.inferenceContext.run { receiverType.isNullableType() }
+
+    val resultingType =
+        if (isReceiverActuallyNullable)
+            typeAfterNullCheck.withNullability(ConeNullability.NULLABLE, session.inferenceContext)
+        else
+            typeAfterNullCheck
+
+    replaceTypeRef(typeRef.resolvedTypeFromPrototype(resultingType))
+}
+
+private fun FirQualifiedAccess.expressionTypeOrUnitForAssignment(): ConeKotlinType? {
+    if (this is FirExpression) return typeRef.coneTypeSafe()
+
+    require(this is FirVariableAssignment) {
+        "The only non-expression FirQualifiedAccess is FirVariableAssignment, but ${this::class} was found"
+    }
+    return StandardClassIds.Unit.constructClassLikeType(emptyArray(), isNullable = false)
+}

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

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.scopes.impl.FirILTTypeRefPlaceHolder
 import org.jetbrains.kotlin.fir.scopes.impl.FirIntegerOperator
 import org.jetbrains.kotlin.fir.scopes.impl.FirIntegerOperatorCall
+import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
@@ -32,21 +33,13 @@ import org.jetbrains.kotlin.types.model.CaptureStatus
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 fun Candidate.resolveArgumentExpression(
-    /*
-    csBuilder: ConstraintSystemBuilder,
-    argument: KotlinCallArgument,
-    expectedType: UnwrappedType?,
-    diagnosticsHolder: KotlinDiagnosticsHolder,
-    isReceiver: Boolean
-     */
     csBuilder: ConstraintSystemBuilder,
     argument: FirExpression,
     expectedType: ConeKotlinType,
     expectedTypeRef: FirTypeRef,
     sink: CheckerSink,
     isReceiver: Boolean,
-    isDispatch: Boolean,
-    isSafeCall: Boolean
+    isDispatch: Boolean
 ) {
     when (argument) {
         is FirFunctionCall, is FirWhenExpression, is FirTryExpression, is FirCheckNotNullCall -> resolveSubCallArgument(
@@ -55,9 +48,39 @@ fun Candidate.resolveArgumentExpression(
             expectedType,
             sink,
             isReceiver,
-            isDispatch,
-            isSafeCall
+            isDispatch
         )
+        // x?.bar() is desugared to `x SAFE-CALL-OPERATOR { $not-null-receiver$.bar() }`
+        //
+        // If we have a safe-call as argument like in a call "foo(x SAFE-CALL-OPERATOR { $not-null-receiver$.bar() })"
+        // we obtain argument type (and argument's constraint system) from "$not-null-receiver$.bar()" (argument.regularQualifiedAccess)
+        // and then add constraint: typeOf(`$not-null-receiver$.bar()`).makeNullable() <: EXPECTED_TYPE
+        // NB: argument.regularQualifiedAccess is either a call or a qualified access
+        is FirSafeCallExpression -> {
+            val nestedQualifier = argument.regularQualifiedAccess
+            if (nestedQualifier is FirExpression) {
+                resolveSubCallArgument(
+                    csBuilder,
+                    nestedQualifier,
+                    expectedType,
+                    sink,
+                    isReceiver,
+                    isDispatch,
+                    useNullableArgumentType = true
+                )
+            } else {
+                // Assignment
+                checkApplicabilityForArgumentType(
+                    csBuilder,
+                    StandardClassIds.Unit.constructClassLikeType(emptyArray(), isNullable = false),
+                    expectedType.type,
+                    SimpleConstraintSystemConstraintPosition,
+                    isReceiver = false,
+                    isDispatch = false,
+                    sink = sink
+                )
+            }
+        }
         is FirCallableReferenceAccess ->
             if (argument.calleeReference is FirResolvedNamedReference)
                 resolvePlainExpressionArgument(
@@ -66,8 +89,7 @@ fun Candidate.resolveArgumentExpression(
                     expectedType,
                     sink,
                     isReceiver,
-                    isDispatch,
-                    isSafeCall
+                    isDispatch
                 )
             else
                 preprocessCallableReference(argument, expectedType)
@@ -82,8 +104,7 @@ fun Candidate.resolveArgumentExpression(
             expectedTypeRef,
             sink,
             isReceiver,
-            isDispatch,
-            isSafeCall
+            isDispatch
         )
         is FirBlock -> resolveBlockArgument(
             csBuilder,
@@ -92,10 +113,9 @@ fun Candidate.resolveArgumentExpression(
             expectedTypeRef,
             sink,
             isReceiver,
-            isDispatch,
-            isSafeCall
+            isDispatch
         )
-        else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isDispatch, isSafeCall)
+        else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isDispatch)
     }
 }
 
@@ -106,8 +126,7 @@ private fun Candidate.resolveBlockArgument(
     expectedTypeRef: FirTypeRef,
     sink: CheckerSink,
     isReceiver: Boolean,
-    isDispatch: Boolean,
-    isSafeCall: Boolean
+    isDispatch: Boolean
 ) {
     val returnArguments = block.returnExpressions()
     if (returnArguments.isEmpty()) {
@@ -118,7 +137,6 @@ private fun Candidate.resolveBlockArgument(
             SimpleConstraintSystemConstraintPosition,
             isReceiver = false,
             isDispatch = false,
-            nullableExpectedType = expectedType.type.withNullability(ConeNullability.NULLABLE, sink.components.session.inferenceContext),
             sink = sink
         )
         return
@@ -131,8 +149,7 @@ private fun Candidate.resolveBlockArgument(
             expectedTypeRef,
             sink,
             isReceiver,
-            isDispatch,
-            isSafeCall
+            isDispatch
         )
     }
 }
@@ -144,7 +161,7 @@ fun Candidate.resolveSubCallArgument(
     sink: CheckerSink,
     isReceiver: Boolean,
     isDispatch: Boolean,
-    isSafeCall: Boolean
+    useNullableArgumentType: Boolean = false
 ) {
     val candidate = argument.candidate() ?: return resolvePlainExpressionArgument(
         csBuilder,
@@ -153,7 +170,7 @@ fun Candidate.resolveSubCallArgument(
         sink,
         isReceiver,
         isDispatch,
-        isSafeCall
+        useNullableArgumentType
     )
     /*
      * It's important to extract type from argument neither from symbol, because of symbol contains
@@ -165,7 +182,7 @@ fun Candidate.resolveSubCallArgument(
         sink.components.returnTypeCalculator.tryCalculateReturnType(candidate.symbol.firUnsafe()).coneTypeUnsafe()
     }
     val argumentType = candidate.substitutor.substituteOrSelf(type)
-    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, isSafeCall)
+    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, useNullableArgumentType)
 }
 
 fun Candidate.resolvePlainExpressionArgument(
@@ -175,11 +192,11 @@ fun Candidate.resolvePlainExpressionArgument(
     sink: CheckerSink,
     isReceiver: Boolean,
     isDispatch: Boolean,
-    isSafeCall: Boolean
+    useNullableArgumentType: Boolean = false
 ) {
     if (expectedType == null) return
     val argumentType = argument.typeRef.coneTypeSafe<ConeKotlinType>() ?: return
-    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, isSafeCall)
+    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, useNullableArgumentType)
     checkApplicabilityForIntegerOperatorCall(sink, argument)
 }
 
@@ -197,22 +214,22 @@ fun Candidate.resolvePlainArgumentType(
     sink: CheckerSink,
     isReceiver: Boolean,
     isDispatch: Boolean,
-    isSafeCall: Boolean
+    useNullableArgumentType: Boolean = false
 ) {
     val position = SimpleConstraintSystemConstraintPosition //TODO
 
     val session = sink.components.session
     val capturedType = prepareCapturedType(argumentType)
 
-    val nullableExpectedType = expectedType.withNullability(ConeNullability.NULLABLE, session.inferenceContext)
-    if (isReceiver && isSafeCall) {
-        if (!isDispatch && !csBuilder.addSubtypeConstraintIfCompatible(capturedType, nullableExpectedType, position)) {
-            sink.reportApplicability(CandidateApplicability.WRONG_RECEIVER) // TODO
-        }
-        return
-    }
+    val argumentTypeForApplicabilityCheck =
+        if (useNullableArgumentType)
+            capturedType.withNullability(ConeNullability.NULLABLE, session.inferenceContext)
+        else
+            capturedType
 
-    checkApplicabilityForArgumentType(csBuilder, capturedType, expectedType, position, isReceiver, isDispatch, nullableExpectedType, sink)
+    checkApplicabilityForArgumentType(
+        csBuilder, argumentTypeForApplicabilityCheck, expectedType, position, isReceiver, isDispatch, sink
+    )
 }
 
 fun Candidate.prepareCapturedType(argumentType: ConeKotlinType): ConeKotlinType {
@@ -238,7 +255,6 @@ private fun checkApplicabilityForArgumentType(
     position: SimpleConstraintSystemConstraintPosition,
     isReceiver: Boolean,
     isDispatch: Boolean,
-    nullableExpectedType: ConeKotlinType,
     sink: CheckerSink
 ) {
     if (isReceiver && isDispatch) {
@@ -252,6 +268,9 @@ private fun checkApplicabilityForArgumentType(
             csBuilder.addSubtypeConstraint(argumentType, expectedType, position)
             return
         }
+
+        val nullableExpectedType = expectedType.withNullability(ConeNullability.NULLABLE, sink.components.session.inferenceContext)
+
         if (csBuilder.addSubtypeConstraintIfCompatible(argumentType, nullableExpectedType, position)) {
             sink.reportApplicability(CandidateApplicability.WRONG_RECEIVER) // TODO
         } else {
@@ -277,8 +296,7 @@ internal fun Candidate.resolveArgument(
         parameter.returnTypeRef,
         sink,
         isReceiver,
-        false,
-        isSafeCall
+        false
     )
 }
 

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

@@ -61,6 +61,7 @@ fun PostponedArgumentsAnalyzer.Context.addSubsystemFromExpression(statement: Fir
     when (statement) {
         is FirFunctionCall, is FirQualifiedAccessExpression, is FirWhenExpression, is FirTryExpression, is FirCheckNotNullCall, is FirCallableReferenceAccess ->
             (statement as FirResolvable).candidate()?.let { addOtherSystem(it.system.asReadOnlyStorage()) }
+        is FirSafeCallExpression -> addSubsystemFromExpression(statement.regularQualifiedAccess)
         is FirWrappedArgumentExpression -> addSubsystemFromExpression(statement.expression)
         is FirBlock -> statement.returnExpressions().forEach { addSubsystemFromExpression(it) }
         else -> {}

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

@@ -133,8 +133,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
                     expectedTypeRef = explicitReceiverExpression.typeRef,
                     sink = sink,
                     isReceiver = true,
-                    isDispatch = this is Dispatch,
-                    isSafeCall = callInfo.isSafeCall
+                    isDispatch = this is Dispatch
                 )
                 sink.yieldIfNeed()
             } else {
@@ -146,8 +145,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
                         expectedType = candidate.substitutor.substituteOrSelf(expectedReceiverType.type),
                         sink = sink,
                         isReceiver = true,
-                        isDispatch = this is Dispatch,
-                        isSafeCall = callInfo.isSafeCall
+                        isDispatch = this is Dispatch
                     )
                     sink.yieldIfNeed()
                 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceCompletion.kt

@@ -259,6 +259,10 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
                 this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
             }
 
+            is FirSafeCallExpression -> {
+                this.regularQualifiedAccess.processAllContainingCallCandidates(processBlocks, processor)
+            }
+
             is FirWhenExpression -> {
                 processCandidateIfApplicable(processor, processBlocks)
                 this.branches.forEach { it.result.processAllContainingCallCandidates(processBlocks, processor) }

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

@@ -145,8 +145,7 @@ class PostponedArgumentsAnalyzer(
                 lambda.atom.returnTypeRef, // TODO: proper ref
                 checkerSink,
                 isReceiver = false,
-                isDispatch = false,
-                isSafeCall = false
+                isDispatch = false
             )
         }
 
@@ -193,4 +192,4 @@ fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(
     ) as ResolvedLambdaAtom
     analyzed = true
     return resolvedAtom
-}
\ No newline at end of file
+}

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

@@ -5,13 +5,15 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers
 
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildVarargArgumentsExpression
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.varargElementType
@@ -19,8 +21,10 @@ import org.jetbrains.kotlin.fir.resolve.constructFunctionalTypeRef
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.inference.returnType
+import org.jetbrains.kotlin.fir.resolve.propagateTypeFromQualifiedAccessAfterNullCheck
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
+import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.scopes.impl.FirIntegerOperatorCall
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.types.*
@@ -59,7 +63,8 @@ class FirCallCompletionResultsWriterTransformer(
         val typeRef = if (declaration is FirTypedDeclaration) {
             typeCalculator.tryCalculateReturnType(declaration).let {
                 if (qualifiedAccess.safe) {
-                    val nullableType = it.coneTypeUnsafe<ConeKotlinType>().withNullability(ConeNullability.NULLABLE, session.inferenceContext)
+                    val nullableType =
+                        it.coneTypeUnsafe<ConeKotlinType>().withNullability(ConeNullability.NULLABLE, session.inferenceContext)
                     it.withReplacedConeType(nullableType)
                 } else {
                     it
@@ -205,6 +210,22 @@ class FirCallCompletionResultsWriterTransformer(
         return withReplacedConeType(finalType)
     }
 
+    override fun transformSafeCallExpression(
+        safeCallExpression: FirSafeCallExpression,
+        data: ExpectedArgumentType?
+    ): CompositeTransformResult<FirStatement> {
+        safeCallExpression.transformRegularQualifiedAccess(
+            this,
+            data?.getExpectedType(
+                safeCallExpression
+            )?.toExpectedType()
+        )
+
+        safeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(safeCallExpression.receiver, session)
+
+        return safeCallExpression.compose()
+    }
+
     override fun transformCallableReferenceAccess(
         callableReferenceAccess: FirCallableReferenceAccess,
         data: ExpectedArgumentType?,

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

@@ -7,12 +7,14 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirImplementationDetail
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.propagateTypeFromOriginalReceiver
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.scopes.impl.FirIntegerOperator
@@ -31,7 +33,8 @@ import org.jetbrains.kotlin.types.AbstractTypeChecker
 
 class IntegerLiteralTypeApproximationTransformer(
     private val symbolProvider: FirSymbolProvider,
-    private val inferenceContext: ConeInferenceContext
+    private val inferenceContext: ConeInferenceContext,
+    private val session: FirSession
 ) : FirTransformer<ConeKotlinType?>() {
     override fun <E : FirElement> transformElement(element: E, data: ConeKotlinType?): CompositeTransformResult<E> {
         return element.compose()
@@ -127,6 +130,16 @@ class IntegerLiteralTypeApproximationTransformer(
         typeOperatorCall.argumentList.transformArguments(this, null)
         return typeOperatorCall.compose()
     }
+
+    override fun transformCheckedSafeCallSubject(
+        checkedSafeCallSubject: FirCheckedSafeCallSubject,
+        data: ConeKotlinType?
+    ): CompositeTransformResult<FirStatement> {
+        val newReceiver =
+            checkedSafeCallSubject.originalReceiverReference.value.transform<FirExpression, ConeKotlinType?>(this, data).single
+        checkedSafeCallSubject.propagateTypeFromOriginalReceiver(newReceiver, session)
+        return super.transformCheckedSafeCallSubject(checkedSafeCallSubject, data)
+    }
 }
 
 fun FirFunctionCall.getOriginalFunction(): FirCallableDeclaration<*>? {
@@ -217,4 +230,4 @@ private fun FirFunctionCall.toOperatorCall(): FirIntegerOperatorCall {
         argumentList,
         calleeReference,
     )
-}
\ No newline at end of file
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt

@@ -289,7 +289,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             FirDataFlowAnalyzer.createFirDataFlowAnalyzer(this, context.dataFlowAnalyzerContext)
         override val syntheticCallGenerator: FirSyntheticCallGenerator = FirSyntheticCallGenerator(this)
         override val integerLiteralTypeApproximator: IntegerLiteralTypeApproximationTransformer =
-            IntegerLiteralTypeApproximationTransformer(symbolProvider, inferenceComponents.ctx)
+            IntegerLiteralTypeApproximationTransformer(symbolProvider, inferenceComponents.ctx, inferenceComponents.session)
         override val doubleColonExpressionResolver: FirDoubleColonExpressionResolver =
             FirDoubleColonExpressionResolver(session, integerLiteralTypeApproximator)
         override val integerOperatorsTypeUpdater: IntegerOperatorsTypeUpdater = IntegerOperatorsTypeUpdater(integerLiteralTypeApproximator)

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

@@ -189,6 +189,20 @@ open class FirBodyResolveTransformer(
         return expressionsTransformer.transformAugmentedArraySetCall(augmentedArraySetCall, data)
     }
 
+    override fun transformSafeCallExpression(
+        safeCallExpression: FirSafeCallExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return expressionsTransformer.transformSafeCallExpression(safeCallExpression, data)
+    }
+
+    override fun transformCheckedSafeCallSubject(
+        checkedSafeCallSubject: FirCheckedSafeCallSubject,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return expressionsTransformer.transformCheckedSafeCallSubject(checkedSafeCallSubject, data)
+    }
+
     // ------------------------------------- Declarations -------------------------------------
 
     override fun transformDeclaration(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -13,7 +13,9 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.expressions.builder.buildErrorExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildVariableAssignment
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference
@@ -164,6 +166,34 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         return result.compose()
     }
 
+    override fun transformSafeCallExpression(
+        safeCallExpression: FirSafeCallExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        safeCallExpression.transformReceiver(this, ResolutionMode.ContextIndependent)
+
+        val receiver = safeCallExpression.receiver
+
+        dataFlowAnalyzer.enterSafeCallAfterNullCheck(safeCallExpression)
+
+        safeCallExpression.apply {
+            checkedSubject.value.propagateTypeFromOriginalReceiver(receiver, components.session)
+            transformRegularQualifiedAccess(this@FirExpressionsResolveTransformer, data)
+            propagateTypeFromQualifiedAccessAfterNullCheck(receiver, session)
+        }
+
+        dataFlowAnalyzer.exitSafeCall(safeCallExpression)
+
+        return safeCallExpression.compose()
+    }
+
+    override fun transformCheckedSafeCallSubject(
+        checkedSafeCallSubject: FirCheckedSafeCallSubject,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return checkedSafeCallSubject.compose()
+    }
+
     override fun transformFunctionCall(functionCall: FirFunctionCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         if (functionCall.calleeReference is FirResolvedNamedReference && functionCall.resultType is FirImplicitTypeRef) {
             storeTypeFromCallee(functionCall)