[FIR] Return statements need to forward postponed lambdas DFA edges

When there is a return statement within a lambda, and the return
statement contains a lambda, data-flow information from the inner lambda
is not being passed correctly to the surrounding lambda. This is because
return statements which had a lambda target would drop all postponed
lambda exits, but should instead forward those exits to the surrounding
lambda exit.

^KT-59729 Fixed
^KT-64268 Fixed

compiler/fir/analysis-tests/testData/resolve/cfg/postponedLambdaInReturn.dot

@@ -221,6 +221,7 @@ digraph postponedLambdaInReturn_kt {
     58 -> {60};
     59 -> {61};
     60 -> {61} [color=green];
+    60 -> {71} [color=red label="Postponed"];
     61 -> {62};
     62 -> {68};
     62 -> {63} [style=dotted];
@@ -379,6 +380,7 @@ digraph postponedLambdaInReturn_kt {
     99 -> {101};
     100 -> {102};
     101 -> {102} [color=green];
+    101 -> {125} [color=red label="Postponed"];
     102 -> {103};
     103 -> {122};
     103 -> {104} [style=dotted];
@@ -505,11 +507,10 @@ digraph postponedLambdaInReturn_kt {
             180 [label="Function call: R|kotlin/run|<R|kotlin/String|>(...)" style="filled" fillcolor=yellow];
             181 [label="Variable declaration: lval x: R|kotlin/String|"];
             182 [label="Access variable R|<local>/y|"];
-            183 [label="Smart cast: R|<local>/y|"];
-            184 [label="Access variable R|kotlin/String.length|"];
-            185 [label="Exit block"];
+            183 [label="Access variable R|kotlin/String.length<Inapplicable(UNSAFE_CALL): kotlin/String.length>#|"];
+            184 [label="Exit block"];
         }
-        186 [label="Exit function test3" style="filled" fillcolor=red];
+        185 [label="Exit function test3" style="filled" fillcolor=red];
     }
     131 -> {132};
     132 -> {133};
@@ -555,6 +556,7 @@ digraph postponedLambdaInReturn_kt {
     167 -> {169};
     168 -> {170};
     169 -> {170} [color=green];
+    169 -> {180} [color=red label="Postponed"];
     170 -> {171};
     171 -> {177};
     171 -> {172} [style=dotted];
@@ -571,6 +573,5 @@ digraph postponedLambdaInReturn_kt {
     182 -> {183};
     183 -> {184};
     184 -> {185};
-    185 -> {186};
 
 }

compiler/fir/analysis-tests/testData/resolve/cfg/postponedLambdaInReturn.fir.txt

@@ -84,5 +84,5 @@ FILE: postponedLambdaInReturn.kt
 
         }
         )
-        R|<local>/y|.R|kotlin/String.length|
+        R|<local>/y|.R|kotlin/String.length<Inapplicable(UNSAFE_CALL): kotlin/String.length>#|
     }

compiler/fir/analysis-tests/testData/resolve/cfg/postponedLambdaInReturn.kt

@@ -46,5 +46,5 @@ fun test3() {
         else
             return@run ""
     }
-    y.length // bad
+    y<!UNSAFE_CALL!>.<!>length // bad
 }

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

@@ -885,9 +885,8 @@ abstract class FirDataFlowAnalyzer(
     }
 
     fun enterCallArguments(call: FirStatement, arguments: List<FirExpression>) {
-        graphBuilder.enterCall()
-
         val lambdas = arguments.mapNotNull { it.unwrapAnonymousFunctionExpression() }
+        graphBuilder.enterCall(lambdas.mapTo(mutableSetOf()) { it.symbol })
         context.variableAssignmentAnalyzer.enterFunctionCall(lambdas)
         graphBuilder.enterCallArguments(call, lambdas)?.mergeIncomingFlow()
     }

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

@@ -65,7 +65,7 @@ class ControlFlowGraphBuilder {
     private val argumentListSplitNodes: Stack<SplitPostponedLambdasNode?> = stackOf()
     private val postponedAnonymousFunctionNodes =
         mutableMapOf<FirFunctionSymbol<*>, Pair<CFGNode<*>, PostponedLambdaExitNode?>>()
-    private val postponedLambdaExits: Stack<MutableList<Pair<CFGNode<*>, EdgeKind>>> = stackOf()
+    private val postponedLambdaExits: Stack<PostponedLambdas> = stackOf()
 
     private val loopConditionEnterNodes: MutableMap<FirLoop, LoopConditionEnterNode> = mutableMapOf()
     private val loopExitNodes: MutableMap<FirLoop, LoopExitNode> = mutableMapOf()
@@ -266,7 +266,7 @@ class ControlFlowGraphBuilder {
         // So we need an edge right now to enforce ordering, and mark it as dead later if needed.
         addEdge(enterNode, exitNode)
         postponedAnonymousFunctionNodes[symbol] = enterNode to exitNode
-        postponedLambdaExits.top().add(exitNode to EdgeKind.Forward)
+        postponedLambdaExits.top().exits.add(exitNode to EdgeKind.Forward)
         return null
     }
 
@@ -310,15 +310,32 @@ class ControlFlowGraphBuilder {
         return Triple(exitNode, postponedExitNode, graph)
     }
 
-    private fun splitDataFlowForPostponedLambdas() {
-        postponedLambdaExits.push(mutableListOf())
+    private fun splitDataFlowForPostponedLambdas(lambdas: Set<FirFunctionSymbol<*>> = emptySet()) {
+        postponedLambdaExits.push(PostponedLambdas(lambdas))
+    }
+
+    /**
+     * Pop and add the current level exits (if any) to the exit corresponding with the specified
+     * lambda function symbol. This is used when a postponed lambda is present within a return
+     * statement for an outer lambda and data-flow information needs to be preserved.
+     */
+    private fun jumpDataFlowFromPostponedLambdas(symbol: FirFunctionSymbol<*>) {
+        val currentLevelExits = postponedLambdaExits.pop().exits
+        if (currentLevelExits.isEmpty()) return
+
+        for ((lambdas, exits) in postponedLambdaExits.all()) {
+            if (symbol in lambdas) {
+                exits.addAll(currentLevelExits)
+                break
+            }
+        }
     }
 
     private fun unifyDataFlowFromPostponedLambdas(node: CFGNode<*>, callCompleted: Boolean) {
-        val currentLevelExits = postponedLambdaExits.pop()
+        val currentLevelExits = postponedLambdaExits.pop().exits
         if (currentLevelExits.isEmpty()) return
 
-        val nextLevelExits = postponedLambdaExits.topOrNull().takeIf { !callCompleted }
+        val nextLevelExits = postponedLambdaExits.topOrNull()?.exits.takeIf { !callCompleted }
         if (nextLevelExits != null) {
             // Call is incomplete, don't pass data flow from lambdas inside it to lambdas in the outer call.
             for ((exit, kind) in currentLevelExits) {
@@ -363,10 +380,10 @@ class ControlFlowGraphBuilder {
     //   until the entire "when" is resolved; then either unify each branch's lambdas into its
     //   exit node, or create N union nodes (1/branch) and point them into the merge node. KT-59730
     private fun mergeDataFlowFromPostponedLambdas(node: CFGNode<*>, callCompleted: Boolean) {
-        val currentLevelExits = postponedLambdaExits.pop()
+        val currentLevelExits = postponedLambdaExits.pop().exits
         if (currentLevelExits.isEmpty()) return
 
-        val nextLevelExits = postponedLambdaExits.topOrNull().takeIf { !callCompleted }
+        val nextLevelExits = postponedLambdaExits.topOrNull()?.exits.takeIf { !callCompleted }
         if (nextLevelExits != null) {
             node.updateDeadStatus()
             nextLevelExits += createMergePostponedLambdaExitsNode(node.fir).also {
@@ -843,9 +860,6 @@ class ControlFlowGraphBuilder {
     // ----------------------------------- Jump -----------------------------------
 
     fun enterJump(jump: FirJump<*>) {
-        // Data flow from anonymous functions in return values does not merge with any enclosing calls.
-        // For named functions, the return value has to be a completed call anyway, so there should
-        // be no postponed lambdas in it.
         if (jump is FirReturnExpression && jump.target.labeledElement is FirAnonymousFunction) {
             splitDataFlowForPostponedLambdas()
         }
@@ -856,10 +870,7 @@ class ControlFlowGraphBuilder {
         addNonSuccessfullyTerminatingNode(node)
 
         if (jump is FirReturnExpression && jump.target.labeledElement is FirAnonymousFunction) {
-            // TODO: these should be DFA-only edges; they should be pointed into the postponed function exit node?
-            //  With builder inference, lambdas are not necessarily resolved starting from the innermost one...
-            //  See analysis test cfg/postponedLambdaInReturn.kt. KT-59729
-            postponedLambdaExits.pop()
+            jumpDataFlowFromPostponedLambdas(jump.target.labeledElement.symbol)
         }
 
         val nextNode = when (jump) {
@@ -1257,8 +1268,8 @@ class ControlFlowGraphBuilder {
         return createResolvedQualifierNode(resolvedQualifier).also(this::addNewSimpleNode)
     }
 
-    fun enterCall() {
-        splitDataFlowForPostponedLambdas()
+    fun enterCall(lambdas: Set<FirFunctionSymbol<*>> = emptySet()) {
+        splitDataFlowForPostponedLambdas(lambdas)
     }
 
     fun enterCallArguments(call: FirStatement, anonymousFunctions: List<FirAnonymousFunction>): FunctionCallArgumentsEnterNode? {
@@ -1552,6 +1563,11 @@ class ControlFlowGraphBuilder {
             }
         }
     }
+
+    private data class PostponedLambdas(
+        val lambdas: Set<FirFunctionSymbol<*>>,
+        val exits: MutableList<Pair<CFGNode<*>, EdgeKind>> = mutableListOf(),
+    )
 }
 
 fun FirDeclaration?.isLocalClassOrAnonymousObject() = ((this as? FirRegularClass)?.isLocal == true) || this is FirAnonymousObject