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

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

@@ -756,7 +756,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val (node, unionNode) = graphBuilder.exitCheckNotNullCall(checkNotNullCall, callCompleted)
         node.mergeIncomingFlow()
 
-        fun FirExpression.propagateNotNullInfo() {
+        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 ->
@@ -770,20 +775,15 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             }
         }
         when (this) {
-                is FirSafeCallExpression -> receiver.propagateNotNullInfo()
+            is FirSafeCallExpression -> receiver.propagateNotNullInfo(node)
             is FirTypeOperatorCall -> {
                 if (operation == FirOperation.AS || operation == FirOperation.SAFE_AS) {
-                        argument.propagateNotNullInfo()
+                    argument.propagateNotNullInfo(node)
                 }
             }
         }
     }
 
-        checkNotNullCall.argument.propagateNotNullInfo()
-
-        unionNode?.let { unionFlowFromArguments(it) }
-    }
-
     // ----------------------------------- When -----------------------------------
 
     fun enterWhenExpression(whenExpression: FirWhenExpression) {
@@ -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
 }