[FIR] Add property type narrowing

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

@@ -24,10 +24,7 @@ import org.jetbrains.kotlin.fir.expressions.FirConstExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.impl.FirExpressionStub
-import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyClass
+import org.jetbrains.kotlin.fir.lazy.*
-import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyConstructor
-import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyProperty
-import org.jetbrains.kotlin.fir.lazy.Fir2IrLazySimpleFunction
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
@@ -613,7 +610,13 @@ class Fir2IrDeclarationStorage(
             signature,
             containerSource
         ) { symbol ->
-            val accessorReturnType = if (isSetter) irBuiltIns.unitType else propertyType
+            val accessorReturnType = if (isSetter) {
+                irBuiltIns.unitType
+            } else if (property.canNarrowDownGetterType) {
+                property.backingField?.returnTypeRef?.toIrType() ?: propertyType
+            } else {
+                propertyType
+            }
             val visibility = propertyAccessor?.visibility?.let {
                 components.visibilityConverter.convertToDescriptorVisibility(it)
             }

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

@@ -30,12 +30,36 @@ sealed class ResolutionMode {
     class WithExpectedTypeFromCast(
         val expectedTypeRef: FirTypeRef,
     ) : ResolutionMode()
+
+    /**
+     * This resolution mode is similar to
+     * WithExpectedType, but it's ok if the
+     * types turn out to be incompatible.
+     * Consider the following examples with
+     * properties and their backing fields:
+     *
+     * val items: List<T>
+     *     field = mutableListOf()
+     *
+     * val s: String
+     *     field = 10
+     *     get() = ...
+     *
+     * In these examples we should try using
+     * the property type information while
+     * resolving the initializer, but it's ok
+     * if it's not applicable.
+     */
+    class WithSuggestedType(
+        val suggestedTypeRef: FirTypeRef,
+    ) : ResolutionMode()
 }
 
 fun ResolutionMode.expectedType(components: BodyResolveComponents, allowFromCast: Boolean = false): FirTypeRef? = when (this) {
     is ResolutionMode.WithExpectedType -> expectedTypeRef
     is ResolutionMode.ContextIndependent -> components.noExpectedType
     is ResolutionMode.WithExpectedTypeFromCast -> expectedTypeRef.takeIf { allowFromCast }
+    is ResolutionMode.WithSuggestedType -> suggestedTypeRef
     else -> null
 }
 

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

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.FirPropertyWithExplicitBackingFieldResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
 import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
@@ -45,6 +46,7 @@ import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
 import org.jetbrains.kotlin.types.SmartcastStability
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 fun List<FirQualifierPart>.toTypeProjections(): Array<ConeTypeProjection> =
     asReversed().flatMap { it.typeArgumentList.typeArguments.map { typeArgument -> typeArgument.toConeTypeProjection() } }.toTypedArray()
@@ -195,6 +197,18 @@ fun <T : FirResolvable> BodyResolveComponents.typeFromCallee(access: T): FirReso
         is FirNamedReferenceWithCandidate -> {
             typeFromSymbol(newCallee.candidateSymbol, false)
         }
+        is FirPropertyWithExplicitBackingFieldResolvedNamedReference -> {
+            val backingField = newCallee.resolvedSymbol
+                .safeAs<FirPropertySymbol>()
+                ?.fir
+                ?.backingField
+
+            if (newCallee.hasVisibleBackingField && backingField != null) {
+                typeFromSymbol(backingField.symbol, false)
+            } else {
+                typeFromSymbol(newCallee.resolvedSymbol, false)
+            }
+        }
         is FirResolvedNamedReference -> {
             typeFromSymbol(newCallee.resolvedSymbol, false)
         }

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

@@ -22,11 +22,8 @@ import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
-import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.types.SmartcastStability
 
 interface Receiver
@@ -55,7 +52,7 @@ abstract class AbstractExplicitReceiverValue<E : FirExpression> : AbstractExplic
         get() = explicitReceiver
 }
 
-class ExpressionReceiverValue(
+open class ExpressionReceiverValue(
     override val explicitReceiver: FirExpression
 ) : AbstractExplicitReceiverValue<FirExpression>(), ReceiverValue {
     override fun scope(useSiteSession: FirSession, scopeSession: ScopeSession): FirTypeScope? {
@@ -72,6 +69,11 @@ class ExpressionReceiverValue(
     }
 }
 
+class NarrowedExpressionReceiverValue(
+    explicitReceiver: FirExpression,
+    override val type: ConeKotlinType,
+) : ExpressionReceiverValue(explicitReceiver)
+
 sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
     val boundSymbol: S,
     type: ConeKotlinType,

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

@@ -5,21 +5,33 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls.tower
 
+import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.asReversedFrozen
+import org.jetbrains.kotlin.fir.containingClass
+import org.jetbrains.kotlin.fir.declarations.utils.canNarrowDownGetterType
+import org.jetbrains.kotlin.fir.declarations.utils.isFinal
+import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.builder.buildExpressionStub
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.DoubleColonLHS
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.descriptorUtil.HIDES_MEMBERS_NAME_LIST
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 
 internal class TowerDataElementsForName(
@@ -224,12 +236,64 @@ internal open class FirTowerResolveTask(
         )
     }
 
+    private fun FirPropertySymbol.hasAccessibleBackingField(info: CallInfo): Boolean {
+        val receiverContainer = containingClass()?.toSymbol(info.session)
+
+        if (!this.isFinal) {
+            return false
+        }
+
+        return when (fir.backingField?.visibility) {
+            Visibilities.Private -> {
+                info.containingDeclarations.any {
+                    it.symbol == receiverContainer
+                }
+            }
+            Visibilities.Internal -> {
+                receiverContainer?.moduleData == info.containingFile.moduleData
+            }
+            else -> {
+                false
+            }
+        }
+    }
+
+    private fun FirExpression.getNarrowedDownReturnType(info: CallInfo): ConeKotlinType? {
+        val propertyReceiver = this.safeAs<FirQualifiedAccessExpression>()
+            ?.calleeReference
+            ?.safeAs<FirResolvedNamedReference>()
+            ?.resolvedSymbol
+            ?.safeAs<FirPropertySymbol>()
+            ?: return null
+
+        // This can happen in case of 2 properties referencing
+        // each other recursively. See: Jet81.fir.kt
+        if (
+            propertyReceiver.fir.returnTypeRef is FirImplicitTypeRef ||
+            propertyReceiver.fir.backingField?.returnTypeRef is FirImplicitTypeRef
+        ) {
+            return null
+        }
+
+        if (propertyReceiver.hasAccessibleBackingField(info) && propertyReceiver.canNarrowDownGetterType) {
+            return propertyReceiver.fir.backingField?.returnTypeRef?.coneType
+        }
+
+        return null
+    }
+
     suspend fun runResolverForExpressionReceiver(
         info: CallInfo,
         receiver: FirExpression,
         parentGroup: TowerGroup = TowerGroup.EmptyRoot
     ) {
-        val explicitReceiverValue = ExpressionReceiverValue(receiver)
+        val narrowedType = receiver.getNarrowedDownReturnType(info)
+
+        val explicitReceiverValue = if (narrowedType != null) {
+            NarrowedExpressionReceiverValue(receiver, narrowedType)
+        } else {
+            ExpressionReceiverValue(receiver)
+        }
 
         processExtensionsThatHideMembers(info, explicitReceiverValue, parentGroup)
 

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

@@ -57,7 +57,14 @@ class FirCallCompleter(
         expectedTypeRef: FirTypeRef?,
         expectedTypeMismatchIsReportedInChecker: Boolean = false,
     ): CompletionResult<T> where T : FirResolvable, T : FirStatement =
-        completeCall(call, expectedTypeRef, mayBeCoercionToUnitApplied = false, expectedTypeMismatchIsReportedInChecker, isFromCast = false)
+        completeCall(
+            call,
+            expectedTypeRef,
+            mayBeCoercionToUnitApplied = false,
+            expectedTypeMismatchIsReportedInChecker,
+            isFromCast = false,
+            shouldEnforceExpectedType = true,
+        )
 
     fun <T> completeCall(call: T, data: ResolutionMode): CompletionResult<T> where T : FirResolvable, T : FirStatement =
         completeCall(
@@ -66,6 +73,7 @@ class FirCallCompleter(
             (data as? ResolutionMode.WithExpectedType)?.mayBeCoercionToUnitApplied == true,
             (data as? ResolutionMode.WithExpectedType)?.expectedTypeMismatchIsReportedInChecker == true,
             isFromCast = data is ResolutionMode.WithExpectedTypeFromCast,
+            shouldEnforceExpectedType = data !is ResolutionMode.WithSuggestedType,
         )
 
     private fun <T> completeCall(
@@ -73,6 +81,7 @@ class FirCallCompleter(
         mayBeCoercionToUnitApplied: Boolean,
         expectedTypeMismatchIsReportedInChecker: Boolean,
         isFromCast: Boolean,
+        shouldEnforceExpectedType: Boolean,
     ): CompletionResult<T>
             where T : FirResolvable, T : FirStatement {
         val typeRef = components.typeFromCallee(call)
@@ -92,25 +101,32 @@ class FirCallCompleter(
             session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, call.source, null)
         }
 
-        if (expectedTypeRef is FirResolvedTypeRef) {
+        val expectedTypeConstraintPosition = ConeExpectedTypeConstraintPosition(expectedTypeMismatchIsReportedInChecker)
-            if (isFromCast) {
+
-                if (candidate.isFunctionForExpectTypeFromCastFeature()) {
+        when {
+            expectedTypeRef !is FirResolvedTypeRef -> {
+                // nothing
+            }
+            !shouldEnforceExpectedType -> {
+                candidate.system.addSubtypeConstraintIfCompatible(
+                    initialType, expectedTypeRef.type, expectedTypeConstraintPosition
+                )
+            }
+            isFromCast -> when {
+                candidate.isFunctionForExpectTypeFromCastFeature() -> {
                     candidate.system.addSubtypeConstraint(
                         initialType, expectedTypeRef.type,
                         ConeExpectedTypeConstraintPosition(expectedTypeMismatchIsReportedInChecker = false),
                     )
                 }
-            } else {
+            }
-                val expectedTypeConstraintPosition = ConeExpectedTypeConstraintPosition(expectedTypeMismatchIsReportedInChecker)
+            !expectedTypeRef.coneType.isUnitOrFlexibleUnit || !mayBeCoercionToUnitApplied -> {
-                if (expectedTypeRef.coneType.isUnitOrFlexibleUnit && mayBeCoercionToUnitApplied) {
+                candidate.system.addSubtypeConstraint(initialType, expectedTypeRef.type, expectedTypeConstraintPosition)
-                    if (candidate.system.notFixedTypeVariables.isNotEmpty()) {
+            }
-                        candidate.system.addSubtypeConstraintIfCompatible(
+            candidate.system.notFixedTypeVariables.isNotEmpty() -> {
-                            initialType, expectedTypeRef.type, expectedTypeConstraintPosition
+                candidate.system.addSubtypeConstraintIfCompatible(
-                        )
+                    initialType, expectedTypeRef.type, expectedTypeConstraintPosition
-                    }
+                )
-                } else {
-                    candidate.system.addSubtypeConstraint(initialType, expectedTypeRef.type, expectedTypeConstraintPosition)
-                }
             }
         }
 

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

@@ -45,8 +45,6 @@ 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.name.Name
-import org.jetbrains.kotlin.types.AbstractTypeChecker
-import org.jetbrains.kotlin.name.SpecialNames
 
 open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
     private val statusResolver: FirStatusResolver = FirStatusResolver(session, scopeSession)
@@ -128,7 +126,12 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         val returnTypeRef = property.returnTypeRef
         val bodyResolveState = property.bodyResolveState
         if (bodyResolveState == FirPropertyBodyResolveState.EVERYTHING_RESOLVED) return property
-        if (returnTypeRef !is FirImplicitTypeRef && implicitTypeOnly) return property
+
+        val canHaveDeepImplicitTypeRefs = property.hasExplicitBackingField
+
+        if (returnTypeRef !is FirImplicitTypeRef && implicitTypeOnly && !canHaveDeepImplicitTypeRefs) {
+            return property
+        }
 
         property.transformReceiverTypeRef(transformer, ResolutionMode.ContextIndependent)
         dataFlowAnalyzer.enterProperty(property)
@@ -727,8 +730,20 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     }
                 }
             }
-            is ResolutionMode.WithExpectedType, is ResolutionMode.ContextIndependent -> {
+            is ResolutionMode.WithExpectedType,
-                val expectedTypeRef = (data as? ResolutionMode.WithExpectedType)?.expectedTypeRef ?: buildImplicitTypeRef()
+            is ResolutionMode.ContextIndependent,
+            is ResolutionMode.WithSuggestedType -> {
+                val expectedTypeRef = when (data) {
+                    is ResolutionMode.WithExpectedType -> {
+                        data.expectedTypeRef
+                    }
+                    is ResolutionMode.WithSuggestedType -> {
+                        data.suggestedTypeRef
+                    }
+                    else -> {
+                        buildImplicitTypeRef()
+                    }
+                }
                 val resolvedLambdaAtom = (expectedTypeRef as? FirResolvedTypeRef)?.let {
                     extractLambdaInfoFromFunctionalType(
                         it.type, it, anonymousFunction, returnTypeVariable = null, components, candidate = null
@@ -883,10 +898,15 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         backingField: FirBackingField,
         data: ResolutionMode,
     ): FirStatement {
-        backingField.transformInitializer(
+        val propertyType = data.expectedType
-            transformer,
+        val initializerData = if (backingField.returnTypeRef is FirResolvedTypeRef) {
             withExpectedType(backingField.returnTypeRef)
-        )
+        } else if (propertyType != null) {
+            ResolutionMode.WithSuggestedType(propertyType)
+        } else {
+            ResolutionMode.ContextDependent
+        }
+        backingField.transformInitializer(transformer, initializerData)
         if (
             backingField.returnTypeRef is FirErrorTypeRef ||
             backingField.returnTypeRef is FirResolvedTypeRef

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
+import org.jetbrains.kotlin.fir.declarations.utils.hasExplicitBackingField
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
@@ -166,7 +167,9 @@ open class FirImplicitAwareBodyResolveTransformer(
             return transform()
         }
 
-        if (member.returnTypeRef is FirResolvedTypeRef) return member
+        val canHaveDeepImplicitTypeRefs = member is FirProperty && member.hasExplicitBackingField
+
+        if (member.returnTypeRef is FirResolvedTypeRef && !canHaveDeepImplicitTypeRefs) return member
         val symbol = member.symbol
         val status = implicitBodyResolveComputationSession.getStatus(symbol)
         if (status is ImplicitBodyResolveComputationStatus.Computed) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/utils/declarationAttributes.kt

@@ -64,6 +64,15 @@ fun FirPropertySymbol.getExplicitBackingField(): FirBackingField? {
     return fir.getExplicitBackingField()
 }
 
+val FirProperty.canNarrowDownGetterType: Boolean
+    get() {
+        val backingFieldHasDifferentType = backingField != null && backingField?.returnTypeRef?.coneType != returnTypeRef.coneType
+        return backingFieldHasDifferentType && getter is FirDefaultPropertyGetter
+    }
+
+val FirPropertySymbol.canNarrowDownGetterType: Boolean
+    get() = fir.canNarrowDownGetterType
+
 // See [BindingContext.BACKING_FIELD_REQUIRED]
 val FirProperty.hasBackingField: Boolean
     get() {