FIR. Refactor smart-cast representation in FIR tree

Make smart-casts non-transparent expression without delegation to underlying FirQualifiedAccessExpression, as children delegation in fir tree has unclear semantics Remove two different kinds of tree nodes for smart-casts
2022-08-02 00:49:24 +02:00
parent bc9db58b3c
commit 513af2dfbc
154 changed files with 9573 additions and 9320 deletions
@@ -322,24 +322,24 @@ private fun BodyResolveComponents.typeFromSymbol(symbol: FirBasedSymbol<*>, make

 fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(
    qualifiedAccessExpression: FirQualifiedAccessExpression
-): FirQualifiedAccessExpression {
+): FirExpression {
+    val (stability, typesFromSmartCast) =
+        dataFlowAnalyzer.getTypeUsingSmartcastInfo(qualifiedAccessExpression)
+            ?: return qualifiedAccessExpression
    val builder = transformExpressionUsingSmartcastInfo(
        qualifiedAccessExpression,
-        dataFlowAnalyzer::getTypeUsingSmartcastInfo,
-        ::FirExpressionWithSmartcastBuilder,
-        ::FirExpressionWithSmartcastToNothingBuilder
+        stability, typesFromSmartCast
    ) ?: return qualifiedAccessExpression
    return builder.build()
 }

 fun BodyResolveComponents.transformWhenSubjectExpressionUsingSmartcastInfo(
    whenSubjectExpression: FirWhenSubjectExpression
-): FirWhenSubjectExpression {
+): FirExpression {
+    val (stability, typesFromSmartCast) = dataFlowAnalyzer.getTypeUsingSmartcastInfo(whenSubjectExpression) ?: return whenSubjectExpression
    val builder = transformExpressionUsingSmartcastInfo(
        whenSubjectExpression,
-        dataFlowAnalyzer::getTypeUsingSmartcastInfo,
-        ::FirWhenSubjectExpressionWithSmartcastBuilder,
-        ::FirWhenSubjectExpressionWithSmartcastToNothingBuilder
+        stability, typesFromSmartCast
    ) ?: return whenSubjectExpression
    return builder.build()
 }
@@ -347,13 +347,18 @@ fun BodyResolveComponents.transformWhenSubjectExpressionUsingSmartcastInfo(
 private val ConeKotlinType.isKindOfNothing
    get() = lowerBoundIfFlexible().let { it.isNothing || it.isNullableNothing }

-private inline fun <T : FirExpression> BodyResolveComponents.transformExpressionUsingSmartcastInfo(
+private fun FirSmartCastExpressionBuilder.applyResultTypeRef() {
+    typeRef =
+        if (smartcastStability == SmartcastStability.STABLE_VALUE)
+            smartcastType.copyWithNewSourceKind(KtFakeSourceElementKind.ImplicitTypeRef)
+        else
+            originalExpression.typeRef.copyWithNewSourceKind(KtFakeSourceElementKind.ImplicitTypeRef)
+}
+private fun <T : FirExpression> BodyResolveComponents.transformExpressionUsingSmartcastInfo(
    expression: T,
-    smartcastExtractor: (T) -> Pair<PropertyStability, MutableList<ConeKotlinType>>?,
-    smartcastBuilder: () -> FirWrappedExpressionWithSmartcastBuilder<T>,
-    smartcastToNothingBuilder: () -> FirWrappedExpressionWithSmartcastToNothingBuilder<T>
-): FirWrappedExpressionWithSmartcastBuilder<T>? {
-    val (stability, typesFromSmartCast) = smartcastExtractor(expression) ?: return null
+    stability: PropertyStability,
+    typesFromSmartCast: MutableList<ConeKotlinType>
+): FirSmartCastExpressionBuilder? {
    val smartcastStability = stability.impliedSmartcastStability
        ?: if (dataFlowAnalyzer.isAccessToUnstableLocalVariable(expression)) {
            SmartcastStability.CAPTURED_VARIABLE
@@ -395,20 +400,24 @@ private inline fun <T : FirExpression> BodyResolveComponents.transformExpression
            annotations += expression.resultType.annotations
            delegatedTypeRef = expression.resultType
        }
-        return smartcastToNothingBuilder().apply {
+        return FirSmartCastExpressionBuilder().apply {
            originalExpression = expression
+            source = originalExpression.source?.fakeElement(KtFakeSourceElementKind.SmartCastExpression)
            smartcastType = intersectedTypeRef
            smartcastTypeWithoutNullableNothing = reducedIntersectedTypeRef
            this.typesFromSmartCast = typesFromSmartCast
            this.smartcastStability = smartcastStability
+            applyResultTypeRef()
        }
    }

-    return smartcastBuilder().apply {
+    return FirSmartCastExpressionBuilder().apply {
        originalExpression = expression
+        source = originalExpression.source?.fakeElement(KtFakeSourceElementKind.SmartCastExpression)
        smartcastType = intersectedTypeRef
        this.typesFromSmartCast = typesFromSmartCast
        this.smartcastStability = smartcastStability
+        applyResultTypeRef()
    }
 }

@@ -419,7 +428,7 @@ fun FirCheckedSafeCallSubject.propagateTypeFromOriginalReceiver(
 ) {
    // 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? FirExpressionWithSmartcastToNothing)
+    val receiverType = ((nullableReceiverExpression as? FirSmartCastExpression)
        ?.takeIf { it.isStable }
        ?.smartcastTypeWithoutNullableNothing
        ?: nullableReceiverExpression.typeRef)
@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.lookupTracker
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
+import org.jetbrains.kotlin.fir.resolve.dfa.unwrapSmartcastExpression
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.inference.preprocessCallableReference
 import org.jetbrains.kotlin.fir.resolve.inference.preprocessLambdaArgument
@@ -66,7 +67,7 @@ fun Candidate.resolveArgumentExpression(
        // 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.selector
+            val nestedQualifier = (argument.selector as? FirExpression)?.unwrapSmartcastExpression()
            if (nestedQualifier is FirQualifiedAccessExpression) {
                resolveSubCallArgument(
                    csBuilder,
@@ -403,7 +404,7 @@ private fun checkApplicabilityForArgumentType(
    }

    if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, expectedType, position)) {
-        val smartcastExpression = argument as? FirExpressionWithSmartcast
+        val smartcastExpression = argument as? FirSmartCastExpression
        if (smartcastExpression != null && !smartcastExpression.isStable) {
            val unstableType = smartcastExpression.smartcastType.coneType
            if (csBuilder.addSubtypeConstraintIfCompatible(unstableType, expectedType, position)) {
@@ -166,7 +166,7 @@ object CheckDispatchReceiver : ResolutionStage() {
        val smartcastedReceiver = when (explicitReceiverExpression) {
            is FirCheckNotNullCall -> explicitReceiverExpression.argument
            else -> explicitReceiverExpression
-        } as? FirExpressionWithSmartcast
+        } as? FirSmartCastExpression

        if (smartcastedReceiver != null &&
            !smartcastedReceiver.isStable &&
@@ -184,7 +184,7 @@ object CheckDispatchReceiver : ResolutionStage() {
                UnstableSmartCast(
                    smartcastedReceiver,
                    targetType,
-                    context.session.typeContext.isTypeMismatchDueToNullability(smartcastedReceiver.originalType.coneType, targetType)
+                    context.session.typeContext.isTypeMismatchDueToNullability(smartcastedReceiver.originalExpression.typeRef.coneType, targetType)
                )
            )
        } else if (isReceiverNullable) {
@@ -5,7 +5,9 @@

 package org.jetbrains.kotlin.fir.resolve.calls

+import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fakeElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.FirBackingField
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
@@ -13,9 +15,10 @@ import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.expressions.FirSmartCastExpression
 import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
-import org.jetbrains.kotlin.fir.expressions.builder.buildExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.expressions.builder.buildSmartCastExpression
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.isNullableNothing
 import org.jetbrains.kotlin.fir.types.makeConeTypeDefinitelyNotNullOrNotNull
@@ -78,12 +81,14 @@ fun FirVisibilityChecker.isVisible(

 private fun removeSmartCastTypeForAttemptToFitVisibility(dispatchReceiverValue: ReceiverValue?, session: FirSession): ReceiverValue? {
    val expressionWithSmartcastIfStable =
-        (dispatchReceiverValue?.receiverExpression as? FirExpressionWithSmartcast)?.takeIf { it.isStable } ?: return null
+        (dispatchReceiverValue?.receiverExpression as? FirSmartCastExpression)?.takeIf { it.isStable } ?: return null

    if (dispatchReceiverValue.type.isNullableNothing) return null

+    val originalExpression = expressionWithSmartcastIfStable.originalExpression
+    val originalType = originalExpression.typeRef.coneType
    val originalTypeNotNullable =
-        expressionWithSmartcastIfStable.originalType.coneType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext)
+        originalType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext)

    // Basically, this `if` is just for sake of optimizaton
    // We have only nullability enhancement, here, so return initial smart cast receiver value
@@ -91,15 +96,19 @@ private fun removeSmartCastTypeForAttemptToFitVisibility(dispatchReceiverValue:

    val expressionForReceiver = with(session.typeContext) {
        when {
-            expressionWithSmartcastIfStable.originalType.coneType.isNullableType() && !dispatchReceiverValue.type.isNullableType() ->
-                buildExpressionWithSmartcast {
-                    originalExpression = expressionWithSmartcastIfStable.originalExpression
-                    smartcastType =
-                        expressionWithSmartcastIfStable.originalExpression.typeRef.resolvedTypeFromPrototype(originalTypeNotNullable)
+            originalType.isNullableType() && !dispatchReceiverValue.type.isNullableType() ->
+                buildSmartCastExpression {
+                    source = originalExpression.source?.fakeElement(KtFakeSourceElementKind.SmartCastExpression)
+                    this.originalExpression = originalExpression
+                    smartcastType = buildResolvedTypeRef {
+                        source = originalExpression.typeRef.source?.fakeElement(KtFakeSourceElementKind.SmartCastedTypeRef)
+                        type = originalTypeNotNullable
+                    }
                    typesFromSmartCast = listOf(originalTypeNotNullable)
                    smartcastStability = expressionWithSmartcastIfStable.smartcastStability
+                    typeRef = smartcastType.copyWithNewSourceKind(KtFakeSourceElementKind.ImplicitTypeRef)
                }
-            else -> expressionWithSmartcastIfStable.originalExpression
+            else -> originalExpression
        }
    }

@@ -318,7 +318,9 @@ private fun BodyResolveComponents.createExplicitReceiverForInvokeByCallable(
        if (candidate.currentApplicability == CandidateApplicability.K2_PROPERTY_AS_OPERATOR) {
            nonFatalDiagnostics.add(ConePropertyAsOperator(candidate.symbol as FirPropertySymbol))
        }
-    }.build().let(::transformQualifiedAccessUsingSmartcastInfo)
+    }.build().let {
+        transformQualifiedAccessUsingSmartcastInfo(it)
+    }
 }

 private class InvokeReceiverResolveTask(
@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.ContextReceiverGroup
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.expressions.FirSmartCastExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildResolvedQualifier
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.*
@@ -84,9 +84,9 @@ class MemberScopeTowerLevel(
        val scope = dispatchReceiverValue.scope(session, scopeSession) ?: return ProcessResult.SCOPE_EMPTY
        var (empty, candidates) = scope.collectCandidates(processScopeMembers)

-        val scopeWithoutSmartcast = (dispatchReceiverValue.receiverExpression as? FirExpressionWithSmartcast)
+        val scopeWithoutSmartcast = (dispatchReceiverValue.receiverExpression as? FirSmartCastExpression)
            ?.takeIf { it.isStable }
-            ?.originalType
+            ?.originalExpression?.typeRef
            ?.coneType
            ?.scope(session, scopeSession, bodyResolveComponents.returnTypeCalculator.fakeOverrideTypeCalculator)
        if (scopeWithoutSmartcast == null) {
@@ -168,6 +168,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(

    fun isAccessToUnstableLocalVariable(expression: FirExpression): Boolean {
        val qualifiedAccessExpression = when (expression) {
+            is FirSmartCastExpression -> expression.originalExpression as FirQualifiedAccessExpression
            is FirQualifiedAccessExpression -> expression
            is FirWhenSubjectExpression -> {
                val whenExpression = expression.whenRef.value
@@ -1002,6 +1003,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
        processConditionalContract(qualifiedAccessExpression)
    }

+    fun exitSmartCastExpression(smartCastExpression: FirSmartCastExpression) {
+        graphBuilder.exitSmartCastExpression(smartCastExpression).mergeIncomingFlow()
+    }
+
    fun enterSafeCallAfterNullCheck(safeCall: FirSafeCallExpression) {
        val node = graphBuilder.enterSafeCall(safeCall).mergeIncomingFlow()
        val previousNode = node.firstPreviousNode
@@ -1239,7 +1244,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
            logicSystem.recordNewAssignment(flow, propertyVariable, context.newAssignmentIndex())
        }

-        variableStorage.getOrCreateRealVariable(flow, initializer.symbol, initializer)
+        variableStorage.getOrCreateRealVariable(flow, initializer.symbol, initializer.unwrapSmartcastExpression())
            ?.let { initializerVariable ->
                val isInitializerStable =
                    initializerVariable.isStable || (initializerVariable.hasLocalStability && initializer.isAccessToStableVariable())
@@ -1274,7 +1279,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
    }

    private fun FirExpression.isAccessToStableVariable(): Boolean =
-        this is FirQualifiedAccessExpression && !isAccessToUnstableLocalVariable(this)
+        !isAccessToUnstableLocalVariable(this)

    private val RealVariable.isStable get() = stability == PropertyStability.STABLE_VALUE
    private val RealVariable.hasLocalStability get() = stability == PropertyStability.LOCAL_VAR
@@ -1087,6 +1087,17 @@ class ControlFlowGraphBuilder {
        return node
    }

+    fun exitSmartCastExpression(smartCastExpression: FirSmartCastExpression): SmartCastExpressionExitNode {
+        val returnsNothing = smartCastExpression.resultType.isNothing
+        val node = createSmartCastExitNode(smartCastExpression)
+        if (returnsNothing) {
+            addNodeThatReturnsNothing(node)
+        } else {
+            addNewSimpleNode(node)
+        }
+        return node
+    }
+
    fun exitResolvedQualifierNode(resolvedQualifier: FirResolvedQualifier): ResolvedQualifierNode {
        return createResolvedQualifierNode(resolvedQualifier).also(this::addNewSimpleNode)
    }
@@ -264,3 +264,6 @@ fun ControlFlowGraphBuilder.createExitDefaultArgumentsNode(fir: FirValueParamete

 fun ControlFlowGraphBuilder.createComparisonExpressionNode(fir: FirComparisonExpression): ComparisonExpressionNode =
    ComparisonExpressionNode(currentGraph, fir, levelCounter, createId())
+
+fun ControlFlowGraphBuilder.createSmartCastExitNode(fir: FirSmartCastExpression): SmartCastExpressionExitNode =
+    SmartCastExpressionExitNode(currentGraph, fir, levelCounter, createId())
@@ -429,6 +429,10 @@ class FirCallCompletionResultsWriterTransformer(
        )
    }

+    override fun transformSmartCastExpression(smartCastExpression: FirSmartCastExpression, data: ExpectedArgumentType?): FirStatement {
+        return smartCastExpression.transformOriginalExpression(this, data)
+    }
+
    private inner class TypeUpdaterForDelegateArguments : FirTransformer<Any?>() {
        override fun <E : FirElement> transformElement(element: E, data: Any?): E {
            return element
@@ -160,9 +160,13 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
        }
        when (result) {
            is FirQualifiedAccessExpression -> {
+                // TODO: Is it really needed?
                dataFlowAnalyzer.enterQualifiedAccessExpression()
-                result = components.transformQualifiedAccessUsingSmartcastInfo(result)
                dataFlowAnalyzer.exitQualifiedAccessExpression(result)
+                result = components.transformQualifiedAccessUsingSmartcastInfo(result)
+                if (result is FirSmartCastExpression) {
+                    dataFlowAnalyzer.exitSmartCastExpression(result)
+                }
            }
            is FirResolvedQualifier -> {
                dataFlowAnalyzer.exitResolvedQualifierNode(result)
@@ -375,7 +379,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
            try {
                val initialExplicitReceiver = functionCall.explicitReceiver
                val resultExpression = callResolver.resolveCallAndSelectCandidate(functionCall)
-                val resultExplicitReceiver = resultExpression.explicitReceiver
+                val resultExplicitReceiver = resultExpression.explicitReceiver?.unwrapSmartcastExpression()
                if (initialExplicitReceiver !== resultExplicitReceiver && resultExplicitReceiver is FirQualifiedAccess) {
                    // name.invoke() case
                    callCompleter.completeCall(resultExplicitReceiver, noExpectedType)
@@ -111,18 +111,8 @@ class ConeEffectExtractor(
        return ConeIsNullPredicate(arg, isNegated)
    }

-    override fun visitExpressionWithSmartcast(
-        expressionWithSmartcast: FirExpressionWithSmartcast,
-        data: Nothing?
-    ): ConeContractDescriptionElement? {
-        return expressionWithSmartcast.originalExpression.accept(this, data)
-    }
-
-    override fun visitExpressionWithSmartcastToNothing(
-        expressionWithSmartcastToNothing: FirExpressionWithSmartcastToNothing,
-        data: Nothing?
-    ): ConeContractDescriptionElement? {
-        return expressionWithSmartcastToNothing.originalExpression.accept(this, data)
+    override fun visitSmartCastExpression(smartCastExpression: FirSmartCastExpression, data: Nothing?): ConeContractDescriptionElement? {
+        return smartCastExpression.originalExpression.accept(this, data)
    }

    override fun visitQualifiedAccessExpression(