[FIR] Properly support smartcasts on stable when subjects in when conditions

^KT-49860 Fixed
2021-11-24 17:14:21 +03:00
parent 20425fb458
commit 330574cab6
41 changed files with 1225 additions and 791 deletions
@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.canNarrowDownGetterType
 import org.jetbrains.kotlin.fir.declarations.utils.expandedConeType
 import org.jetbrains.kotlin.fir.declarations.utils.isFinal
-import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
@@ -31,6 +30,8 @@ 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.dfa.FirDataFlowAnalyzer
+import org.jetbrains.kotlin.fir.resolve.dfa.PropertyStability
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
@@ -291,26 +292,52 @@ private fun BodyResolveComponents.typeFromSymbol(symbol: FirBasedSymbol<*>, make
 fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(
    qualifiedAccessExpression: FirQualifiedAccessExpression
 ): FirQualifiedAccessExpression {
-    val (stability, typesFromSmartCast) = dataFlowAnalyzer.getTypeUsingSmartcastInfo(qualifiedAccessExpression)
-        ?: return qualifiedAccessExpression
+    val builder = transformExpressionUsingSmartcastInfo(
+        qualifiedAccessExpression,
+        dataFlowAnalyzer::getTypeUsingSmartcastInfo,
+        ::FirExpressionWithSmartcastBuilder,
+        ::FirExpressionWithSmartcastToNullBuilder
+    ) ?: return qualifiedAccessExpression
+    return builder.build()
+}
+
+fun BodyResolveComponents.transformWhenSubjectExpressionUsingSmartcastInfo(
+    whenSubjectExpression: FirWhenSubjectExpression
+): FirWhenSubjectExpression {
+    val builder = transformExpressionUsingSmartcastInfo(
+        whenSubjectExpression,
+        dataFlowAnalyzer::getTypeUsingSmartcastInfo,
+        ::FirWhenSubjectExpressionWithSmartcastBuilder,
+        ::FirWhenSubjectExpressionWithSmartcastToNullBuilder
+    ) ?: return whenSubjectExpression
+    return builder.build()
+}
+
+private inline fun <T : FirExpression> BodyResolveComponents.transformExpressionUsingSmartcastInfo(
+    expression: T,
+    smartcastExtractor: (T) -> Pair<PropertyStability, MutableList<ConeKotlinType>>?,
+    smartcastBuilder: () -> FirWrappedExpressionWithSmartcastBuilder<T>,
+    smartcastToNullBuilder: () -> FirWrappedExpressionWithSmartcastToNullBuilder<T>
+): FirWrappedExpressionWithSmartcastBuilder<T>? {
+    val (stability, typesFromSmartCast) = smartcastExtractor(expression) ?: return null
    val smartcastStability = stability.impliedSmartcastStability
-        ?: if (dataFlowAnalyzer.isAccessToUnstableLocalVariable(qualifiedAccessExpression)) {
+        ?: if (dataFlowAnalyzer.isAccessToUnstableLocalVariable(expression)) {
            SmartcastStability.CAPTURED_VARIABLE
        } else {
            SmartcastStability.STABLE_VALUE
        }

-    val originalType = qualifiedAccessExpression.resultType.coneType
+    val originalType = expression.resultType.coneType
    val allTypes = typesFromSmartCast.also {
        it += originalType
    }
    val intersectedType = ConeTypeIntersector.intersectTypes(session.typeContext, allTypes)
-    if (intersectedType == originalType) return qualifiedAccessExpression
+    if (intersectedType == originalType) return null
    val intersectedTypeRef = buildResolvedTypeRef {
-        source = qualifiedAccessExpression.resultType.source?.fakeElement(KtFakeSourceElementKind.SmartCastedTypeRef)
+        source = expression.resultType.source?.fakeElement(KtFakeSourceElementKind.SmartCastedTypeRef)
        type = intersectedType
-        annotations += qualifiedAccessExpression.resultType.annotations
-        delegatedTypeRef = qualifiedAccessExpression.resultType
+        annotations += expression.resultType.annotations
+        delegatedTypeRef = expression.resultType
    }
    // For example, if (x == null) { ... },
    //   we need to track the type without `Nothing?` so that resolution with this as receiver can go through properly.
@@ -322,13 +349,13 @@ fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(
        val intersectedTypeWithoutNullableNothing =
            ConeTypeIntersector.intersectTypes(session.typeContext, typesFromSmartcastWithoutNullableNothing)
        val intersectedTypeRefWithoutNullableNothing = buildResolvedTypeRef {
-            source = qualifiedAccessExpression.resultType.source?.fakeElement(KtFakeSourceElementKind.SmartCastedTypeRef)
+            source = expression.resultType.source?.fakeElement(KtFakeSourceElementKind.SmartCastedTypeRef)
            type = intersectedTypeWithoutNullableNothing
-            annotations += qualifiedAccessExpression.resultType.annotations
-            delegatedTypeRef = qualifiedAccessExpression.resultType
+            annotations += expression.resultType.annotations
+            delegatedTypeRef = expression.resultType
        }
-        return buildExpressionWithSmartcastToNull {
-            originalExpression = qualifiedAccessExpression
+        return smartcastToNullBuilder().apply {
+            originalExpression = expression
            smartcastType = intersectedTypeRef
            smartcastTypeWithoutNullableNothing = intersectedTypeRefWithoutNullableNothing
            this.typesFromSmartCast = typesFromSmartCast
@@ -336,8 +363,8 @@ fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(
        }
    }

-    return buildExpressionWithSmartcast {
-        originalExpression = qualifiedAccessExpression
+    return smartcastBuilder().apply {
+        originalExpression = expression
        smartcastType = intersectedTypeRef
        this.typesFromSmartCast = typesFromSmartCast
        this.smartcastStability = smartcastStability
@@ -6,8 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa

 import org.jetbrains.kotlin.config.LanguageFeature
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.PrivateForInline
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.description.ConeBooleanConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeConditionalEffectDeclaration
@@ -17,7 +16,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
@@ -35,7 +33,6 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.visibilityChecker
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.types.ConstantValueKind
@@ -165,18 +162,41 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(

    // ----------------------------------- Requests -----------------------------------

-    fun isAccessToUnstableLocalVariable(qualifiedAccessExpression: FirQualifiedAccessExpression): Boolean {
+    fun isAccessToUnstableLocalVariable(expression: FirExpression): Boolean {
+        val qualifiedAccessExpression = when (expression) {
+            is FirQualifiedAccessExpression -> expression
+            is FirWhenSubjectExpression -> {
+                val whenExpression = expression.whenRef.value
+                when {
+                    whenExpression.subjectVariable != null -> return true
+                    else -> whenExpression.subject as? FirQualifiedAccessExpression
+                }
+            }
+            else -> null
+        }  ?: return false
        return context.firLocalVariableAssignmentAnalyzer?.isAccessToUnstableLocalVariable(qualifiedAccessExpression) == true
    }

+    fun getTypeUsingSmartcastInfo(whenSubjectExpression: FirWhenSubjectExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
+        val symbol = whenSubjectExpression.symbol ?: return null
+        return getTypeUsingSmartcastInfo(symbol, whenSubjectExpression)
+    }
+
    fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
        /*
         * DataFlowAnalyzer holds variables only for declarations that have some smartcast (or can have)
         * If there is no useful information there is no data flow variable also
         */
        val symbol: FirBasedSymbol<*> = qualifiedAccessExpression.symbol ?: return null
+        return getTypeUsingSmartcastInfo(symbol, qualifiedAccessExpression)
+    }
+
+    private fun getTypeUsingSmartcastInfo(
+        symbol: FirBasedSymbol<*>,
+        expression: FirExpression
+    ): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
        val flow = graphBuilder.lastNode.flow
-        var variable = variableStorage.getRealVariableWithoutUnwrappingAlias(symbol, qualifiedAccessExpression, flow) ?: return null
+        var variable = variableStorage.getRealVariableWithoutUnwrappingAlias(symbol, expression, flow) ?: return null
        val stability = variable.stability
        val result = mutableListOf<ConeKotlinType>()
        flow.directAliasMap[variable]?.let {
@@ -740,6 +760,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
        whenExitNode.mergeIncomingFlow()
    }

+    fun exitWhenSubjectExpression(expression: FirWhenSubjectExpression) {
+        graphBuilder.exitWhenSubjectExpression(expression).mergeIncomingFlow()
+    }
+
    // ----------------------------------- While Loop -----------------------------------

    private fun exitCommonLoop(exitNode: LoopExitNode) {
@@ -1178,6 +1178,9 @@ class ControlFlowGraphBuilder {
        return Pair(kind, unionNode)
    }

+    fun exitWhenSubjectExpression(expression: FirWhenSubjectExpression): WhenSubjectExpressionExitNode {
+        return createWhenSubjectExpressionExitNode(expression).also { addNewSimpleNode(it) }
+    }

    // ----------------------------------- Annotations -----------------------------------

@@ -157,6 +157,9 @@ fun ControlFlowGraphBuilder.createElvisRhsEnterNode(fir: FirElvisExpression): El
 fun ControlFlowGraphBuilder.createElvisLhsExitNode(fir: FirElvisExpression): ElvisLhsExitNode =
    ElvisLhsExitNode(currentGraph, fir, levelCounter, createId())

+fun ControlFlowGraphBuilder.createWhenSubjectExpressionExitNode(fir: FirWhenSubjectExpression): WhenSubjectExpressionExitNode =
+    WhenSubjectExpressionExitNode(currentGraph, fir, levelCounter, createId())
+
 fun ControlFlowGraphBuilder.createElvisExitNode(fir: FirElvisExpression): ElvisExitNode =
    ElvisExitNode(currentGraph, fir, levelCounter, createId())

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.isUnitOrFlexibleUnit
+import org.jetbrains.kotlin.fir.resolve.transformWhenSubjectExpressionUsingSmartcastInfo
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.fir.resolve.withExpectedType
@@ -126,7 +127,8 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
        if (subjectType != null) {
            whenSubjectExpression.resultType = subjectType
        }
-        return whenSubjectExpression
+        dataFlowAnalyzer.exitWhenSubjectExpression(whenSubjectExpression)
+        return components.transformWhenSubjectExpressionUsingSmartcastInfo(whenSubjectExpression)
    }

    // ------------------------------- Try/catch expressions -------------------------------