[FIR] Fix smartcast after elvis with escaping lambda, ^KT-44510 Fixed

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

@@ -756,32 +756,32 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val (node, unionNode) = graphBuilder.exitCheckNotNullCall(checkNotNullCall, callCompleted)
         node.mergeIncomingFlow()
 
-        fun FirExpression.propagateNotNullInfo() {
-            val symbol = this.symbol
-            if (symbol != null) {
-                variableStorage.getOrCreateRealVariable(node.previousFlow, symbol, this)?.let { operandVariable ->
-                    node.flow.addTypeStatement(operandVariable typeEq any)
-                    logicSystem.approveStatementsInsideFlow(
-                        node.flow,
-                        operandVariable notEq null,
-                        shouldRemoveSynthetics = true,
-                        shouldForkFlow = false
-                    )
-                }
+        checkNotNullCall.argument.propagateNotNullInfo(node)
+
+        unionNode?.let { unionFlowFromArguments(it) }
+    }
+
+    fun FirExpression.propagateNotNullInfo(node: CFGNode<*>) {
+        val symbol = this.symbol
+        if (symbol != null) {
+            variableStorage.getOrCreateRealVariable(node.previousFlow, symbol, this)?.let { operandVariable ->
+                node.flow.addTypeStatement(operandVariable typeEq any)
+                logicSystem.approveStatementsInsideFlow(
+                    node.flow,
+                    operandVariable notEq null,
+                    shouldRemoveSynthetics = true,
+                    shouldForkFlow = false
+                )
             }
-            when (this) {
-                is FirSafeCallExpression -> receiver.propagateNotNullInfo()
-                is FirTypeOperatorCall -> {
-                    if (operation == FirOperation.AS || operation == FirOperation.SAFE_AS) {
-                        argument.propagateNotNullInfo()
-                    }
+        }
+        when (this) {
+            is FirSafeCallExpression -> receiver.propagateNotNullInfo(node)
+            is FirTypeOperatorCall -> {
+                if (operation == FirOperation.AS || operation == FirOperation.SAFE_AS) {
+                    argument.propagateNotNullInfo(node)
                 }
             }
         }
-
-        checkNotNullCall.argument.propagateNotNullInfo()
-
-        unionNode?.let { unionFlowFromArguments(it) }
     }
 
     // ----------------------------------- When -----------------------------------
@@ -1459,12 +1459,16 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    fun exitElvis(elvisExpression: FirElvisExpression) {
+    fun exitElvis(elvisExpression: FirElvisExpression, isLhsNotNull: Boolean) {
         val node = graphBuilder.exitElvis().mergeIncomingFlow()
+        if (isLhsNotNull) {
+            elvisExpression.lhs.propagateNotNullInfo(node)
+        }
+
         if (!components.session.languageVersionSettings.supportsFeature(LanguageFeature.BooleanElvisBoundSmartCasts)) return
-        val lhs = elvisExpression.lhs
         val rhs = elvisExpression.rhs
         if (rhs is FirConstExpression<*> && rhs.kind == ConstantValueKind.Boolean) {
+            val lhs = elvisExpression.lhs
             if (lhs.typeRef.coneType.classId != StandardClassIds.Boolean) return
 
             val flow = node.flow

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

@@ -245,11 +245,13 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             }
         }
 
+        var isLhsNotNull = false
         if (result.rhs.typeRef.coneTypeSafe<ConeKotlinType>()?.isNothing == true) {
             val lhsType = result.lhs.typeRef.coneTypeSafe<ConeKotlinType>()
             if (lhsType != null) {
                 val newReturnType = lhsType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext)
                 result.replaceTypeRef(result.typeRef.resolvedTypeFromPrototype(newReturnType))
+                isLhsNotNull = true
             }
         }
 
@@ -265,7 +267,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             }
         }
 
-        dataFlowAnalyzer.exitElvis(elvisExpression)
+        dataFlowAnalyzer.exitElvis(elvisExpression, isLhsNotNull = isLhsNotNull)
         return result
     }
 }

compiler/testData/diagnostics/tests/smartCasts/kt32358_2.fir.kt

@@ -14,7 +14,7 @@ inline fun <R> callItContracted(fn: () -> R): R {
 
 fun smartIt(p1: String?, p2: String?) {
     p1 ?: callIt { return }
-    p1<!UNSAFE_CALL!>.<!>length
+    p1.length
 
     p2 ?: callItContracted { return }
     p2.length

compiler/testData/diagnostics/tests/smartCasts/kt32358_3.fir.kt

@@ -6,5 +6,5 @@ fun <R> callIt(fn: () -> R): R = TODO()
 
 fun smartIt(p1: String?, p2: String?) {
     p1 ?: callIt { TODO() }
-    p1<!UNSAFE_CALL!>.<!>length // smartcast
+    p1.length // smartcast
 }