[FIR] Fix cfg for safe call inside elvis

compiler/fir/analysis-tests/testData/resolve/smartcasts/controlStructures/elvis.dot

@@ -61,7 +61,9 @@ digraph elvis_kt {
     5 -> {6};
     6 -> {7};
     7 -> {8};
-    8 -> {9 11};
+    8 -> {9};
+    8 -> {11} [color=red];
+    8 -> {13} [color=green];
     9 -> {10};
     10 -> {11};
     11 -> {12};

compiler/fir/analysis-tests/testData/resolve/smartcasts/controlStructures/returns.dot

@@ -469,7 +469,9 @@ digraph returns_kt {
     151 -> {152};
     152 -> {153};
     153 -> {154};
-    154 -> {155 157};
+    154 -> {155};
+    154 -> {157} [color=red];
+    154 -> {159} [color=green];
     155 -> {156};
     156 -> {157};
     157 -> {158};

compiler/fir/analysis-tests/testData/resolve/smartcasts/nullability.dot

@@ -583,7 +583,9 @@ digraph nullability_kt {
     194 -> {195 197};
     195 -> {196};
     196 -> {197};
-    197 -> {198 200};
+    197 -> {198};
+    197 -> {200} [color=red];
+    197 -> {202} [color=green];
     198 -> {199};
     199 -> {200};
     200 -> {201};

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

@@ -1288,6 +1288,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     // ----------------------------------- Elvis -----------------------------------
 
+    fun enterElvis(elvisExpression: FirElvisExpression) {
+        graphBuilder.enterElvis(elvisExpression)
+    }
+
     fun exitElvisLhs(elvisExpression: FirElvisExpression) {
         val (lhsExitNode, lhsIsNotNullNode, rhsEnterNode) = graphBuilder.exitElvisLhs(elvisExpression)
         lhsExitNode.mergeIncomingFlow()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNode.kt

@@ -119,7 +119,9 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
         protected set
 
     internal fun updateDeadStatus() {
-        isDead = incomingEdges.size == previousNodes.size && incomingEdges.values.all { it.kind == EdgeKind.DeadForward }
+        isDead = incomingEdges.size == previousNodes.size && incomingEdges.values.all {
+            it.kind == EdgeKind.DeadForward || !it.kind.usedInCfa
+        }
     }
 
     abstract fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -95,6 +95,7 @@ class ControlFlowGraphBuilder {
 
     private val exitSafeCallNodes: Stack<ExitSafeCallNode> = stackOf()
     private val exitElvisExpressionNodes: Stack<ElvisExitNode> = stackOf()
+    private val elvisRhsEnterNodes: Stack<ElvisRhsEnterNode> = stackOf()
 
     /*
      * ignoredFunctionCalls is needed for resolve of += operator:
@@ -1165,7 +1166,15 @@ class ControlFlowGraphBuilder {
         val exitNode = createExitSafeCallNode(safeCall)
         exitSafeCallNodes.push(exitNode)
         addEdge(lastNode, enterNode)
-        addEdge(lastNode, exitNode)
+        if (elvisRhsEnterNodes.topOrNull()?.fir?.lhs === safeCall) {
+            //if this is safe call in lhs of elvis, we make two edges
+            // 1. Df-only edge to exit node, to get not null implications there
+            // 2. Cf-only edge to elvis rhs
+            addEdge(lastNode, exitNode, preferredKind = EdgeKind.DfgForward)
+            addEdge(lastNode, elvisRhsEnterNodes.top(), preferredKind = EdgeKind.CfgForward)
+        } else {
+            addEdge(lastNode, exitNode)
+        }
         return enterNode
     }
 
@@ -1184,6 +1193,10 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Elvis -----------------------------------
 
+    fun enterElvis(elvisExpression: FirElvisExpression) {
+        elvisRhsEnterNodes.push(createElvisRhsEnterNode(elvisExpression))
+    }
+
     fun exitElvisLhs(elvisExpression: FirElvisExpression): Triple<ElvisLhsExitNode, ElvisLhsIsNotNullNode, ElvisRhsEnterNode> {
         val exitNode = createElvisExitNode(elvisExpression).also {
             exitElvisExpressionNodes.push(it)
@@ -1198,7 +1211,7 @@ class ControlFlowGraphBuilder {
             addEdge(it, exitNode)
         }
 
-        val rhsEnterNode = createElvisRhsEnterNode(elvisExpression).also {
+        val rhsEnterNode = elvisRhsEnterNodes.pop().also {
             addEdge(lhsExitNode, it)
         }
         lastNodes.push(rhsEnterNode)

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

@@ -224,6 +224,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
                 withExpectedType(expectedType, mayBeCoercionToUnitApplied = true)
             else
                 withExpectedType(expectedType?.withNullability(ConeNullability.NULLABLE, session.typeContext))
+        dataFlowAnalyzer.enterElvis(elvisExpression)
         elvisExpression.transformLhs(transformer, resolutionModeForLhs)
         dataFlowAnalyzer.exitElvisLhs(elvisExpression)
 

compiler/testData/diagnostics/testsWithStdLib/contracts/controlflow/flowInlining/safeCallAndInPlaceReturn.fir.kt

@@ -10,6 +10,13 @@ inline fun <T> Any?.myRun(block: () -> T): T {
     return block()
 }
 
+inline fun <T> directRun(block: () -> T): T {
+    contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+    return block()
+}
+
 fun bad(): String {
     val x: String? = null
 
@@ -20,4 +27,16 @@ fun ok(): String {
     val x: String? = null
 
     x?.run { return "non-null" } ?: return "null"
-<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
+}
+
+fun ok2(): String {
+    directRun {
+        return "nonNull"
+    }
+}
+
+fun ok3(arg: Any?): String {
+    arg?.myRun {
+        return "nonNull"
+    } ?: error("null")
+}

compiler/testData/diagnostics/testsWithStdLib/contracts/controlflow/flowInlining/safeCallAndInPlaceReturn.kt

@@ -10,6 +10,13 @@ inline fun <T> Any?.myRun(block: () -> T): T {
     return block()
 }
 
+inline fun <T> directRun(block: () -> T): T {
+    contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+    return block()
+}
+
 fun bad(): String {
     val x: String? = null
 
@@ -20,4 +27,16 @@ fun ok(): String {
     val x: String? = null
 
     x?.run { return "non-null" } ?: return "null"
+}
+
+fun ok2(): String {
+    directRun {
+        return "nonNull"
+    }
+}
+
+fun ok3(arg: Any?): String {
+    arg?.myRun {
+        return "nonNull"
+    } ?: error("null")
 }

compiler/testData/diagnostics/testsWithStdLib/contracts/controlflow/flowInlining/safeCallAndInPlaceReturn.txt

@@ -1,7 +1,12 @@
 package
 
 public fun bad(): kotlin.String
+public inline fun </*0*/ T> directRun(/*0*/ block: () -> T): T
+    CallsInPlace(block, EXACTLY_ONCE)
+
 public fun ok(): kotlin.String
+public fun ok2(): kotlin.String
+public fun ok3(/*0*/ arg: kotlin.Any?): kotlin.String
 public inline fun </*0*/ T> kotlin.Any?.myRun(/*0*/ block: () -> T): T
     CallsInPlace(block, EXACTLY_ONCE)