[FIR] Add fir node for qualified access with smartcast

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

@@ -11,15 +11,14 @@ import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 
 class DummyFirDataFlowAnalyzer : FirDataFlowAnalyzer() {
     companion object {
         private const val DUMMY = "<DUMMY>"
     }
 
-    override fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): ConeKotlinType? {
-        return qualifiedAccessExpression.typeRef.coneTypeUnsafe()
+    override fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): Collection<ConeKotlinType>? {
+        return null
     }
 
     override fun enterFunction(function: FirFunction<*>) {}

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

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 
 abstract class FirDataFlowAnalyzer {
-    abstract fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): ConeKotlinType?
+    abstract fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): Collection<ConeKotlinType>?
 
     // ----------------------------------- Named function -----------------------------------
 

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

@@ -39,13 +39,10 @@ class FirDataFlowAnalyzerImpl(transformer: FirBodyResolveTransformer) : FirDataF
     private val variableStorage = DataFlowVariableStorage()
     private val edges = mutableMapOf<CFGNode<*>, Flow>().withDefault { Flow.EMPTY }
 
-    override fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): ConeKotlinType? {
+    override fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): Collection<ConeKotlinType>? {
         val symbol: FirBasedSymbol<*> = qualifiedAccessExpression.resolvedSymbol ?: return null
         val variable = variableStorage[symbol]?.real ?: return null
-        val smartCastTypes = graphBuilder.lastNode.flow.approvedFacts(variable)?.exactType ?: return null
-        val smartCastType = context.myIntersectTypes(smartCastTypes.toList()) ?: return null
-        val originalType = qualifiedAccessExpression.typeRef.coneTypeSafe<ConeKotlinType>() ?: return null
-        return ConeTypeIntersector.intersectTypesFromSmartcasts(context, originalType, smartCastType)
+        return graphBuilder.lastNode.flow.approvedFacts(variable)?.exactType ?: return null
     }
 
     // ----------------------------------- Named function -----------------------------------

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.impl.FirExpressionWithSmartcastImpl
 import org.jetbrains.kotlin.fir.references.FirExplicitThisReference
 import org.jetbrains.kotlin.fir.references.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
@@ -246,11 +247,19 @@ open class FirBodyResolveTransformer(
                     transformedCallee
                 }
                 if (result is FirQualifiedAccessExpression) {
-                    dataFlowAnalyzer.getTypeUsingSmartcastInfo(result)?.let { typeWithSmartCasts ->
-                        result.resultType = FirResolvedTypeRefImpl(result.psi, typeWithSmartCasts, result.resultType.annotations)
-                    }
+                    dataFlowAnalyzer.getTypeUsingSmartcastInfo(result)?.let { typesFromSmartCast ->
+                        val allTypes = typesFromSmartCast.toMutableList().also {
+                            it += result.resultType.coneTypeUnsafe<ConeKotlinType>()
+                        }
+                        val intersectedType = ConeTypeIntersector.intersectTypes(inferenceComponents.ctx as ConeInferenceContext, allTypes)
+                        // TODO: add check that intersectedType is not equal to original type
+                        FirExpressionWithSmartcastImpl(result, typesFromSmartCast).also {
+                            it.resultType = FirResolvedTypeRefImpl(result.psi, intersectedType, result.resultType.annotations)
+                        }
+                    } ?: result
+                } else {
+                    result
                 }
-                result
             }
         }
         if (result is FirQualifiedAccessExpression) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionWithSmartcast.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions
+
+import org.jetbrains.kotlin.fir.FirReference
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+
+abstract class FirExpressionWithSmartcast(
+    open val originalExpression: FirQualifiedAccessExpression,
+    val typesFromSmartcast: Collection<ConeKotlinType>
+) : FirQualifiedAccessExpression(originalExpression.psi) {
+    val originalType: FirTypeRef get() = originalExpression.typeRef
+
+    final override val safe: Boolean
+        get() = originalExpression.safe
+
+    final override val explicitReceiver: FirExpression?
+        get() = originalExpression.explicitReceiver
+
+    final override val dispatchReceiver: FirExpression
+        get() = originalExpression.dispatchReceiver
+
+    final override val extensionReceiver: FirExpression
+        get() = originalExpression.extensionReceiver
+
+    final override val calleeReference: FirReference
+        get() = originalExpression.calleeReference
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        originalExpression.acceptChildren(visitor, data)
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirExpressionWithSmartcastImpl.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions.impl
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.transformSingle
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+
+class FirExpressionWithSmartcastImpl(
+    override var originalExpression: FirQualifiedAccessExpression,
+    typesFromSmartcast: Collection<ConeKotlinType>
+) : FirExpressionWithSmartcast(originalExpression, typesFromSmartcast) {
+    init {
+        assert(originalExpression.typeRef is FirResolvedTypeRef)
+    }
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        originalExpression = originalExpression.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformCalleeReference(transformer: FirTransformer<D>, data: D): FirQualifiedAccess {
+        throw IllegalStateException()
+    }
+
+    override fun <D> transformExplicitReceiver(transformer: FirTransformer<D>, data: D): FirQualifiedAccess {
+        throw IllegalStateException()
+    }
+
+    override fun <D> transformDispatchReceiver(transformer: FirTransformer<D>, data: D): FirQualifiedAccess {
+        throw IllegalStateException()
+    }
+
+    override fun <D> transformExtensionReceiver(transformer: FirTransformer<D>, data: D): FirQualifiedAccess {
+        throw IllegalStateException()
+    }
+}

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirTransformerGenerated.kt

@@ -314,6 +314,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformUnknownTypeExpression(errorExpression, data)
     }
 
+    open fun transformExpressionWithSmartcast(expressionWithSmartcast: FirExpressionWithSmartcast, data: D): CompositeTransformResult<FirStatement> {
+        return transformUnknownTypeExpression(expressionWithSmartcast, data)
+    }
+
     open fun transformQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression, data: D): CompositeTransformResult<FirStatement> {
         return transformUnknownTypeExpression(qualifiedAccessExpression, data)
     }
@@ -630,6 +634,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformExpression(expression, data)
     }
 
+    final override fun visitExpressionWithSmartcast(expressionWithSmartcast: FirExpressionWithSmartcast, data: D): CompositeTransformResult<FirElement> {
+        return transformExpressionWithSmartcast(expressionWithSmartcast, data)
+    }
+
     final override fun visitField(field: FirField, data: D): CompositeTransformResult<FirElement> {
         return transformField(field, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorGenerated.kt

@@ -314,6 +314,10 @@ abstract class FirVisitor<out R, in D> {
         return visitUnknownTypeExpression(errorExpression, data)
     }
 
+    open fun visitExpressionWithSmartcast(expressionWithSmartcast: FirExpressionWithSmartcast, data: D): R {
+        return visitUnknownTypeExpression(expressionWithSmartcast, data)
+    }
+
     open fun visitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression, data: D): R {
         return visitUnknownTypeExpression(qualifiedAccessExpression, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorVoidGenerated.kt

@@ -314,6 +314,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitUnknownTypeExpression(errorExpression, null)
     }
 
+    open fun visitExpressionWithSmartcast(expressionWithSmartcast: FirExpressionWithSmartcast) {
+        visitUnknownTypeExpression(expressionWithSmartcast, null)
+    }
+
     open fun visitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {
         visitUnknownTypeExpression(qualifiedAccessExpression, null)
     }
@@ -630,6 +634,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitExpression(expression)
     }
 
+    final override fun visitExpressionWithSmartcast(expressionWithSmartcast: FirExpressionWithSmartcast, data: Nothing?) {
+        visitExpressionWithSmartcast(expressionWithSmartcast)
+    }
+
     final override fun visitField(field: FirField, data: Nothing?) {
         visitField(field)
     }