[FIR] Fix return expressions calculator

There was a problem that if las statement of lambda was a return
  expression from that lambda then result of return counted twice

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

@@ -103,7 +103,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         return result.takeIf { it.isNotEmpty() }
     }
 
-    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): List<FirStatement> {
+    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirStatement> {
         return graphBuilder.returnExpressionsOfAnonymousFunction(function)
     }
 

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

@@ -161,7 +161,7 @@ class ControlFlowGraphBuilder {
         return exitNode to graph
     }
 
-    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): List<FirStatement> {
+    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirStatement> {
         fun FirElement.extractArgument(): FirElement = when {
             this is FirReturnExpression && target.labeledElement.symbol == function.symbol -> result.extractArgument()
             else -> this
@@ -175,7 +175,7 @@ class ControlFlowGraphBuilder {
         }
 
         val exitNode = function.controlFlowGraphReference.controlFlowGraph?.exitNode ?: exitsOfAnonymousFunctions.getValue(function)
-        return exitNode.previousNodes.mapNotNull {
+        return exitNode.previousNodes.mapNotNullTo(mutableSetOf()) {
             it.extractArgument() as FirStatement?
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt

@@ -212,7 +212,7 @@ class ResolvedLambdaAtom(
     val typeVariableForLambdaReturnType: TypeVariableForLambdaReturnType?
 ) : PostponedResolvedAtomMarker {
     override var analyzed: Boolean = false
-    lateinit var returnStatements: List<FirStatement>
+    lateinit var returnStatements: Collection<FirStatement>
 
 //    lateinit var resultArguments: List<KotlinCallArgument>
 //        private set

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.types.model.freshTypeConstructor
 import org.jetbrains.kotlin.types.model.safeSubstitute
 
 data class ReturnArgumentsAnalysisResult(
-    val returnArguments: List<FirStatement>,
+    val returnArguments: Collection<FirStatement>,
     val inferenceSession: FirInferenceSession?
 )
 

compiler/fir/resolve/testData/resolve/cfg/lambdaAsReturnOfLambda.dot

@@ -0,0 +1,79 @@
+digraph lambdaAsReturnOfLambda_kt {
+    graph [splines=ortho nodesep=3]
+    node [shape=box penwidth=2]
+    edge [penwidth=2]
+
+    subgraph cluster_0 {
+        color=red
+        0 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+        1 [label="Jump: ^@run lambda@fun <anonymous>(foo: R|ERROR CLASS: No type for parameter|): R|kotlin/Unit| {
+    R|/bar|(R|<local>/foo|)
+}
+"];
+        2 [label="Stub" style="filled" fillcolor=gray];
+        3 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+    }
+
+    0 -> {1};
+    1 -> {3};
+    1 -> {2} [style=dotted];
+    2 -> {3} [style=dotted];
+
+    subgraph cluster_1 {
+        color=red
+        4 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+        5 [label="Access variable R|<local>/foo|"];
+        6 [label="Function call: R|/bar|(R|<local>/foo|)"];
+        7 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+    }
+
+    4 -> {5};
+    5 -> {6};
+    6 -> {7};
+
+    subgraph cluster_2 {
+        color=red
+        8 [label="Enter function getter" style="filled" fillcolor=red];
+        9 [label="Exit function getter" style="filled" fillcolor=red];
+    }
+
+    8 -> {9};
+
+    subgraph cluster_3 {
+        color=red
+        10 [label="Enter property" style="filled" fillcolor=red];
+        11 [label="Postponed enter to lambda"];
+        12 [label="Postponed exit from lambda"];
+        13 [label="Function call: R|/run|<R|(ERROR CLASS: No type for parameter) -> kotlin/Unit|>(<L> = run@fun <anonymous>(): R|(ERROR CLASS: No type for parameter) -> kotlin/Unit|)"];
+        14 [label="Exit property" style="filled" fillcolor=red];
+    }
+
+    10 -> {11};
+    11 -> {12 12} [color=green];
+    12 -> {13};
+    13 -> {14};
+
+    subgraph cluster_4 {
+        color=red
+        15 [label="Enter function bar" style="filled" fillcolor=red];
+        16 [label="Exit function bar" style="filled" fillcolor=red];
+    }
+
+    15 -> {16};
+
+    subgraph cluster_5 {
+        color=red
+        17 [label="Enter function run" style="filled" fillcolor=red];
+        18 [label="Function call: R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|R|>|()"];
+        19 [label="Jump: ^run R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|R|>|()"];
+        20 [label="Stub" style="filled" fillcolor=gray];
+        21 [label="Exit function run" style="filled" fillcolor=red];
+    }
+
+    17 -> {18};
+    18 -> {19};
+    19 -> {21};
+    19 -> {20} [style=dotted];
+    20 -> {21} [style=dotted];
+
+}

compiler/fir/resolve/testData/resolve/cfg/lambdaAsReturnOfLambda.kt

@@ -0,0 +1,10 @@
+// !DUMP_CFG
+
+val x4: (String) -> Unit = run {
+    return@run (lambda@{ foo ->
+        bar(foo)
+    })
+}
+
+fun bar(s: String) {}
+fun <R> run(block: () -> R): R = block()

compiler/fir/resolve/testData/resolve/cfg/lambdaAsReturnOfLambda.txt

@@ -0,0 +1,14 @@
+FILE: lambdaAsReturnOfLambda.kt
+    public final val x4: R|(kotlin/String) -> kotlin/Unit| = R|/run|<R|(ERROR CLASS: No type for parameter) -> kotlin/Unit|>(<L> = run@fun <anonymous>(): R|(ERROR CLASS: No type for parameter) -> kotlin/Unit| {
+        ^@run lambda@fun <anonymous>(foo: R|ERROR CLASS: No type for parameter|): R|kotlin/Unit| {
+            R|/bar|(R|<local>/foo|)
+        }
+
+    }
+    )
+        public get(): R|(kotlin/String) -> kotlin/Unit|
+    public final fun bar(s: R|kotlin/String|): R|kotlin/Unit| {
+    }
+    public final fun <R> run(block: R|() -> R|): R|R| {
+        ^run R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|R|>|()
+    }

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

@@ -576,6 +576,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/cfg/jumps.kt");
         }
 
+        @TestMetadata("lambdaAsReturnOfLambda.kt")
+        public void testLambdaAsReturnOfLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/cfg/lambdaAsReturnOfLambda.kt");
+        }
+
         @TestMetadata("lambdas.kt")
         public void testLambdas() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/cfg/lambdas.kt");

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

@@ -576,6 +576,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/cfg/jumps.kt");
         }
 
+        @TestMetadata("lambdaAsReturnOfLambda.kt")
+        public void testLambdaAsReturnOfLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/cfg/lambdaAsReturnOfLambda.kt");
+        }
+
         @TestMetadata("lambdas.kt")
         public void testLambdas() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/cfg/lambdas.kt");