[FIR] Fix smartcasts on receivers of implict invoke call

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

@@ -13,12 +13,15 @@ import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.FirEmptyDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.FirSimpleDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.expressions.impl.FirExpressionWithSmartcastImpl
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.calls.ConeInferenceContext
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
@@ -345,3 +348,16 @@ private fun BodyResolveComponents.typeFromSymbol(symbol: AbstractFirBasedSymbol<
         else -> error("WTF ! $symbol")
     }
 }
+
+fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): FirQualifiedAccessExpression {
+    val typesFromSmartCast = dataFlowAnalyzer.getTypeUsingSmartcastInfo(qualifiedAccessExpression) ?: return qualifiedAccessExpression
+    val allTypes = typesFromSmartCast.toMutableList().also {
+        it += qualifiedAccessExpression.resultType.coneTypeUnsafe<ConeKotlinType>()
+    }
+    val intersectedType = ConeTypeIntersector.intersectTypes(inferenceComponents.ctx as ConeInferenceContext, allTypes)
+    // TODO: add check that intersectedType is not equal to original type
+    val intersectedTypeRef = FirResolvedTypeRefImpl(qualifiedAccessExpression.resultType.source, intersectedType).apply {
+        annotations += qualifiedAccessExpression.resultType.annotations
+    }
+    return FirExpressionWithSmartcastImpl(qualifiedAccessExpression, intersectedTypeRef, typesFromSmartCast)
+}

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.calls.InferenceComponents
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
+import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.phasedFir
@@ -37,6 +38,7 @@ interface BodyResolveComponents : SessionHolder {
     val callResolver: FirCallResolver
     val doubleColonExpressionResolver: FirDoubleColonExpressionResolver
     val syntheticCallGenerator: FirSyntheticCallGenerator
+    val dataFlowAnalyzer: FirDataFlowAnalyzer
 
     val <D> AbstractFirBasedSymbol<D>.phasedFir: D where D : FirDeclaration, D : FirSymbolOwner<D>
         get() = phasedFir(session, FirResolvePhase.DECLARATIONS)

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.expressions.impl.FirQualifiedAccessExpressionImpl
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.transformQualifiedAccessUsingSmartcastInfo
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -81,18 +82,21 @@ class InvokeReceiverCandidateCollector(
         if (applicability >= CandidateApplicability.SYNTHETIC_RESOLVED) {
 
             val session = components.session
+            val explicitReceiver = FirQualifiedAccessExpressionImpl(null).apply {
+                calleeReference = FirNamedReferenceWithCandidate(
+                    null,
+                    (candidate.symbol as FirCallableSymbol<*>).callableId.callableName,
+                    candidate
+                )
+                dispatchReceiver = candidate.dispatchReceiverExpression()
+                extensionReceiver = candidate.extensionReceiverExpression()
+                typeRef = towerResolver.typeCalculator.tryCalculateReturnType(candidate.symbol.firUnsafe())
+            }.let {
+                components.transformQualifiedAccessUsingSmartcastInfo(it)
+            }
             val boundInvokeCallInfo = CallInfo(
                 invokeCallInfo.callKind,
-                FirQualifiedAccessExpressionImpl(null).apply {
+                explicitReceiver,
-                    calleeReference = FirNamedReferenceWithCandidate(
-                        null,
-                        (candidate.symbol as FirCallableSymbol<*>).callableId.callableName,
-                        candidate
-                    )
-                    dispatchReceiver = candidate.dispatchReceiverExpression()
-                    extensionReceiver = candidate.extensionReceiverExpression()
-                    typeRef = towerResolver.typeCalculator.tryCalculateReturnType(candidate.symbol.firUnsafe())
-                },
                 invokeCallInfo.arguments,
                 invokeCallInfo.isSafeCall,
                 invokeCallInfo.typeArguments,

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.references.impl.FirExplicitThisReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStackImpl
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.dfa.Condition.*
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.dfa.contracts.buildContractFir
@@ -542,10 +543,11 @@ class FirDataFlowAnalyzer(private val components: FirAbstractBodyResolveTransfor
 
     private val FirResolvable.resolvedSymbol: AbstractFirBasedSymbol<*>?
         get() = calleeReference.let {
-            if (it is FirExplicitThisReference) {
+            when (it) {
-                it.boundSymbol
+                is FirExplicitThisReference -> it.boundSymbol
-            } else {
+                is FirResolvedNamedReference -> it.resolvedSymbol
-                (it as? FirResolvedNamedReference)?.resolvedSymbol
+                is FirNamedReferenceWithCandidate -> it.candidateSymbol
+                else -> null
             }
         }
 
@@ -778,7 +780,7 @@ class FirDataFlowAnalyzer(private val components: FirAbstractBodyResolveTransfor
 
     private val FirElement.realVariable: RealDataFlowVariable?
         get() {
-            val symbol: AbstractFirBasedSymbol<*> = if (this is FirThisReceiverExpressionImpl) {
+            val symbol: AbstractFirBasedSymbol<*> = if (this is FirThisReceiverExpression) {
                 calleeReference.boundSymbol
             } else {
                 resolvedSymbol

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

@@ -120,7 +120,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             FirTypeResolveScopeForBodyResolve(topLevelScopes, implicitReceiverStack, localScopes), session
         )
         val callCompleter: FirCallCompleter = FirCallCompleter(transformer, this)
-        val dataFlowAnalyzer: FirDataFlowAnalyzer = FirDataFlowAnalyzer(this)
+        override val dataFlowAnalyzer: FirDataFlowAnalyzer = FirDataFlowAnalyzer(this)
         override val syntheticCallGenerator: FirSyntheticCallGenerator = FirSyntheticCallGenerator(this, callCompleter)
 
         internal var containerIfAny: FirDeclaration? = null

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

@@ -29,7 +29,6 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.FirOperatorAmbiguityError
 import org.jetbrains.kotlin.fir.resolve.diagnostics.FirVariableExpectedError
 import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreReceiver
-import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
@@ -122,25 +121,12 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
             }
         }
         if (result is FirQualifiedAccessExpression) {
-            result = transformQualifiedAccessUsingSmartcastInfo(result)
+            result = components.transformQualifiedAccessUsingSmartcastInfo(result)
             dataFlowAnalyzer.exitQualifiedAccessExpression(result)
         }
         return result.compose()
     }
 
-    private fun transformQualifiedAccessUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): FirQualifiedAccessExpression {
-        val typesFromSmartCast = dataFlowAnalyzer.getTypeUsingSmartcastInfo(qualifiedAccessExpression) ?: return qualifiedAccessExpression
-        val allTypes = typesFromSmartCast.toMutableList().also {
-            it += qualifiedAccessExpression.resultType.coneTypeUnsafe<ConeKotlinType>()
-        }
-        val intersectedType = ConeTypeIntersector.intersectTypes(inferenceComponents.ctx as ConeInferenceContext, allTypes)
-        // TODO: add check that intersectedType is not equal to original type
-        val intersectedTypeRef = FirResolvedTypeRefImpl(qualifiedAccessExpression.resultType.source, intersectedType).apply {
-            annotations += qualifiedAccessExpression.resultType.annotations
-        }
-        return FirExpressionWithSmartcastImpl(qualifiedAccessExpression, intersectedTypeRef, typesFromSmartCast)
-    }
-
     override fun transformFunctionCall(functionCall: FirFunctionCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         if (functionCall.calleeReference is FirResolvedNamedReference && functionCall.resultType is FirImplicitTypeRef) {
             storeTypeFromCallee(functionCall)

compiler/fir/resolve/testData/resolve/smartcasts/smartcastOnLambda.dot

@@ -0,0 +1,56 @@
+digraph smartcastOnLambda_kt {
+  graph [splines=ortho nodesep=3]
+  node [shape=box penwidth=2]
+  edge [penwidth=2]
+
+  subgraph cluster_0 {
+    color=red
+    0 [label="Enter function test" style="filled" fillcolor=red];
+    subgraph cluster_1 {
+      color=blue
+      1 [label="Enter block"];
+      subgraph cluster_2 {
+        color=blue
+        2 [label="Enter when"];
+        subgraph cluster_3 {
+          color=blue
+          3 [label="Enter when branch condition "];
+          4 [label="Access variable R|<local>/func|"];
+          5 [label="Const: Null(null)"];
+          6 [label="Operator !="];
+          7 [label="Exit when branch condition"];
+        }
+        8 [label="Synthetic else branch"];
+        9 [label="Enter when branch result"];
+        subgraph cluster_4 {
+          color=blue
+          10 [label="Enter block"];
+          11 [label="Function call: R|<local>/func|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
+          12 [label="Exit block"];
+        }
+        13 [label="Exit when branch result"];
+        14 [label="Exit when"];
+      }
+      15 [label="Exit block"];
+    }
+    16 [label="Exit function test" style="filled" fillcolor=red];
+  }
+
+  0 -> {1};
+  1 -> {2};
+  2 -> {3};
+  3 -> {4};
+  4 -> {5};
+  5 -> {6};
+  6 -> {7};
+  7 -> {9 8};
+  8 -> {14};
+  9 -> {10};
+  10 -> {11};
+  11 -> {12};
+  12 -> {13};
+  13 -> {14};
+  14 -> {15};
+  15 -> {16};
+
+}

compiler/fir/resolve/testData/resolve/smartcasts/smartcastOnLambda.kt

@@ -0,0 +1,5 @@
+fun test(func: (() -> Unit)?) {
+    if (func != null) {
+        func()
+    }
+}

compiler/fir/resolve/testData/resolve/smartcasts/smartcastOnLambda.txt

@@ -0,0 +1,9 @@
+FILE: smartcastOnLambda.kt
+    public final fun test(func: R|() -> kotlin/Unit|): R|kotlin/Unit| {
+        when () {
+            !=(R|<local>/func|, Null(null)) ->  {
+                R|<local>/func|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+            }
+        }
+
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithCfgTestGenerated.java

@@ -193,6 +193,11 @@ public class FirDiagnosticsWithCfgTestGenerated extends AbstractFirDiagnosticsWi
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/smartcastAfterReassignment.kt");
         }
 
+        @TestMetadata("smartcastOnLambda.kt")
+        public void testSmartcastOnLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/smartcasts/smartcastOnLambda.kt");
+        }
+
         @TestMetadata("when.kt")
         public void testWhen() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/when.kt");