FIR DFA: smartcast variable to Nothing? on null assignment

In order to make resolution still work for members not available from
`Nothing`, we track the type without `Nothing?` and use that for
resolution instead.

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

@@ -32,10 +32,10 @@ import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByTypeRef
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.transformers.firClassLike
-import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
 import org.jetbrains.kotlin.fir.scopes.impl.importedFromObjectData
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
@@ -266,10 +266,21 @@ fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(
         } else {
             SmartcastStability.STABLE_VALUE
         }
+
     val originalType = qualifiedAccessExpression.resultType.coneType
+    val allTypes = typesFromSmartCast.also {
+        it += originalType
+    }
+    val intersectedType = ConeTypeIntersector.intersectTypes(session.inferenceComponents.ctx, allTypes)
+    if (intersectedType == originalType) return qualifiedAccessExpression
+    val intersectedTypeRef = buildResolvedTypeRef {
+        source = qualifiedAccessExpression.resultType.source?.fakeElement(FirFakeSourceElementKind.SmartCastedTypeRef)
+        type = intersectedType
+        annotations += qualifiedAccessExpression.resultType.annotations
+        delegatedTypeRef = qualifiedAccessExpression.resultType
+    }
     // For example, if (x == null) { ... },
-    //   we don't want to smartcast to Nothing?, but we want to record the nullability to its own kind of node.
-    // TODO: should we differentiate x == null v.s. x is Nothing?
+    //   we need to track the type without `Nothing?` so that resolution with this as receiver can go through properly.
     if (typesFromSmartCast.any { it.isNullableNothing }) {
         val typesFromSmartcastWithoutNullableNothing =
             typesFromSmartCast.filterTo(mutableListOf()) { !it.isNullableNothing }.also {
@@ -285,24 +296,13 @@ fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(
         }
         return buildExpressionWithSmartcastToNull {
             originalExpression = qualifiedAccessExpression
-            // TODO: Use Nothing? during resolution?
-            smartcastType = intersectedTypeRefWithoutNullableNothing
-            // NB: Nothing? in types from smartcast in DFA is recorded here (and the expression kind itself).
+            smartcastType = intersectedTypeRef
+            smartcastTypeWithoutNullableNothing = intersectedTypeRefWithoutNullableNothing
             this.typesFromSmartCast = typesFromSmartCast
             this.smartcastStability = smartcastStability
         }
     }
-    val allTypes = typesFromSmartCast.also {
-        it += originalType
-    }
-    val intersectedType = ConeTypeIntersector.intersectTypes(session.inferenceComponents.ctx, allTypes)
-    if (intersectedType == originalType) return qualifiedAccessExpression
-    val intersectedTypeRef = buildResolvedTypeRef {
-        source = qualifiedAccessExpression.resultType.source?.fakeElement(FirFakeSourceElementKind.SmartCastedTypeRef)
-        type = intersectedType
-        annotations += qualifiedAccessExpression.resultType.annotations
-        delegatedTypeRef = qualifiedAccessExpression.resultType
-    }
+
     return buildExpressionWithSmartcast {
         originalExpression = qualifiedAccessExpression
         smartcastType = intersectedTypeRef
@@ -339,7 +339,13 @@ fun FirAnnotationCall.fqName(session: FirSession): FqName? {
 }
 
 fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(nullableReceiverExpression: FirExpression, session: FirSession) {
-    val receiverType = nullableReceiverExpression.typeRef.coneTypeSafe<ConeKotlinType>() ?: return
+    // If the receiver expression is smartcast to `null`, it would have `Nothing?` as its type, which may not have members called by user
+    // code. Hence, we fallback to the type before intersecting with `Nothing?`.
+    val receiverType = ((nullableReceiverExpression as? FirExpressionWithSmartcastToNull)
+        ?.takeIf { it.isStable }
+        ?.smartcastTypeWithoutNullableNothing
+        ?: nullableReceiverExpression.typeRef)
+        .coneTypeSafe<ConeKotlinType>() ?: return
 
     val expandedReceiverType = if (receiverType is ConeClassLikeType) receiverType.fullyExpandedType(session) else receiverType
 

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

@@ -11,8 +11,8 @@ import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcastToNull
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
-import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirUnstableSmartcastTypeScope
@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirScopeWithFakeOverrideTypeCalculat
 import org.jetbrains.kotlin.fir.scopes.impl.FirStandardOverrideChecker
 import org.jetbrains.kotlin.fir.scopes.impl.FirTypeIntersectionScope
 import org.jetbrains.kotlin.fir.scopes.scopeForClass
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
@@ -33,7 +34,8 @@ fun FirExpressionWithSmartcast.smartcastScope(
     useSiteSession: FirSession,
     scopeSession: ScopeSession
 ): FirTypeScope? {
-    val smartcastType = smartcastType.coneType
+    val smartcastType =
+        if (this is FirExpressionWithSmartcastToNull) smartcastTypeWithoutNullableNothing.coneType else smartcastType.coneType
     val smartcastScope = smartcastType.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing)
     if (isStable) {
         return smartcastScope

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

@@ -7,9 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
-import org.jetbrains.kotlin.fir.expressions.FirThisReceiverExpression
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.builder.buildThisReceiverExpression
 import org.jetbrains.kotlin.fir.references.builder.buildImplicitThisReference
@@ -60,9 +58,18 @@ abstract class AbstractExplicitReceiverValue<E : FirExpression> : AbstractExplic
 class ExpressionReceiverValue(
     override val explicitReceiver: FirExpression
 ) : AbstractExplicitReceiverValue<FirExpression>(), ReceiverValue {
-    override fun scope(useSiteSession: FirSession, scopeSession: ScopeSession): FirTypeScope? =
-        (receiverExpression as? FirExpressionWithSmartcast)?.smartcastScope(useSiteSession, scopeSession)
-            ?: type.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing)
+    override fun scope(useSiteSession: FirSession, scopeSession: ScopeSession): FirTypeScope? {
+        var receiverExpr: FirExpression? = receiverExpression
+        // Unwrap `x!!` to `x` and use the resulted expression to derive receiver type. This is necessary so that smartcast types inside
+        // `!!` is handled correctly.
+        if (receiverExpr is FirCheckNotNullCall) {
+            receiverExpr = receiverExpr.arguments.firstOrNull()
+        }
+        if (receiverExpr is FirExpressionWithSmartcast) {
+            return receiverExpr.smartcastScope(useSiteSession, scopeSession)
+        }
+        return type.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing)
+    }
 }
 
 sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.FirVisibilityChecker
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isInfix
@@ -23,13 +22,14 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.visibilityChecker
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.*
 import org.jetbrains.kotlin.types.AbstractNullabilityChecker
-import org.jetbrains.kotlin.types.SmartcastStability
 
 abstract class ResolutionStage {
     abstract suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext)
@@ -116,9 +116,8 @@ object CheckDispatchReceiver : ResolutionStage() {
             (candidate.originScope as? FirUnstableSmartcastTypeScope)?.isSymbolFromUnstableSmartcast(candidate.symbol) == true
 
         if (explicitReceiverExpression is FirExpressionWithSmartcast &&
-            explicitReceiverExpression !is FirExpressionWithSmartcastToNull &&
-            explicitReceiverExpression.smartcastStability != SmartcastStability.STABLE_VALUE &&
-            (isCandidateFromUnstableSmartcast || isReceiverNullable)
+            !explicitReceiverExpression.isStable &&
+            (isCandidateFromUnstableSmartcast || (isReceiverNullable && !explicitReceiverExpression.smartcastType.canBeNull))
         ) {
             sink.yieldDiagnostic(UnstableSmartCast(explicitReceiverExpression, explicitReceiverExpression.smartcastType.coneType))
         } else if (isReceiverNullable) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -1124,14 +1124,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
 
         if (isAssignment) {
-            if (initializer is FirConstExpression<*> && initializer.kind == ConstantValueKind.Null) {
-                flow.addTypeStatement(
-                    propertyVariable typeEq
-                            property.returnTypeRef.coneType.withNullability(ConeNullability.NULLABLE, components.session.typeContext)
-                )
-            } else {
-                flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneType)
-            }
+            flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneType)
         }
     }