[FIR] Fix binding return expression to function

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -1037,7 +1037,7 @@ class DeclarationsConverter(
         val parentNode = functionDeclaration.getParent()
         val isLocal = !(parentNode?.tokenType == KT_FILE || parentNode?.tokenType == CLASS_BODY)
         val firFunction = if (identifier == null) {
-            FirAnonymousFunctionImpl(null, session, returnType!!, receiverType, FirAnonymousFunctionSymbol())
+            FirAnonymousFunctionImpl(null, session, returnType!!, receiverType, FirAnonymousFunctionSymbol(), isLambda = false)
         } else {
             val functionName = identifier.nameAsSafeName()
             val status = FirDeclarationStatusImpl(

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -127,7 +127,7 @@ class ExpressionsConverter(
             }
         }
 
-        return FirAnonymousFunctionImpl(null, session, implicitType, implicitType, FirAnonymousFunctionSymbol()).apply {
+        return FirAnonymousFunctionImpl(null, session, implicitType, implicitType, FirAnonymousFunctionSymbol(), isLambda = true).apply {
             context.firFunctions += this
             var destructuringBlock: FirExpression? = null
             for (valueParameter in valueParameterList) {

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -105,8 +105,8 @@ abstract class BaseFirBuilder<T>(val session: FirSession, val context: Context =
             this
         ).apply {
             target = FirFunctionTarget(labelName)
-            val lastFunction = context.firFunctions.lastOrNull()
             if (labelName == null) {
+                val lastFunction = context.firFunctions.lastOrNull { !(it is FirAnonymousFunction && it.isLambda) }
                 if (lastFunction != null) {
                     target.bind(lastFunction)
                 } else {

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -592,7 +592,7 @@ class RawFirBuilder(session: FirSession, val stubMode: Boolean) : BaseFirBuilder
             }
             val receiverType = function.receiverTypeReference.convertSafe<FirTypeRef>()
             val firFunction = if (function.name == null) {
-                FirAnonymousFunctionImpl(function.toFirSourceElement(), session, returnType, receiverType, FirAnonymousFunctionSymbol())
+                FirAnonymousFunctionImpl(function.toFirSourceElement(), session, returnType, receiverType, FirAnonymousFunctionSymbol(), isLambda = false)
             } else {
                 val status = FirDeclarationStatusImpl(
                     if (function.isLocal) Visibilities.LOCAL else function.visibility,
@@ -636,7 +636,7 @@ class RawFirBuilder(session: FirSession, val stubMode: Boolean) : BaseFirBuilder
             val literalSource = literal.toFirSourceElement()
             val returnType = FirImplicitTypeRefImpl(literalSource)
             val receiverType = FirImplicitTypeRefImpl(literalSource)
-            return FirAnonymousFunctionImpl(literalSource, session, returnType, receiverType, FirAnonymousFunctionSymbol()).apply {
+            return FirAnonymousFunctionImpl(literalSource, session, returnType, receiverType, FirAnonymousFunctionSymbol(), isLambda = true).apply {
                 context.firFunctions += this
                 var destructuringBlock: FirExpression? = null
                 for (valueParameter in literal.valueParameters) {

compiler/fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt

@@ -0,0 +1,10 @@
+fun <T> run(block: () -> T): T = block()
+
+fun test_1() {
+    run { return@run }
+    run { return }
+}
+
+fun test_2() {
+    run(fun (): Int { return 1 })
+}

compiler/fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.txt

@@ -0,0 +1,20 @@
+FILE: lambdaAndAnonymousFunction.kt
+    public? final? fun <T> run(block: ( () -> T )): T {
+        ^run block#()
+    }
+    public? final? fun test_1(): R|kotlin/Unit| {
+        run#(<L> = run@fun <implicit>.<anonymous>(): <implicit> {
+            ^@run Unit
+        }
+        )
+        run#(<L> = run@fun <implicit>.<anonymous>(): <implicit> {
+            ^test_1 Unit
+        }
+        )
+    }
+    public? final? fun test_2(): R|kotlin/Unit| {
+        run#(fun <anonymous>(): Int {
+            ^ Int(1)
+        }
+        )
+    }

compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTestCaseGenerated.java

@@ -223,6 +223,11 @@ public class RawFirBuilderTestCaseGenerated extends AbstractRawFirBuilderTestCas
             runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");
         }
 
+        @TestMetadata("lambdaAndAnonymousFunction.kt")
+        public void testLambdaAndAnonymousFunction() throws Exception {
+            runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
+        }
+
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {
             runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/locals.kt");

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/CopyUtils.kt

@@ -62,7 +62,7 @@ fun FirAnonymousFunction.copy(
     controlFlowGraphReference: FirControlFlowGraphReference = this.controlFlowGraphReference,
     invocationKind: InvocationKind? = this.invocationKind
 ): FirAnonymousFunction {
-    return FirAnonymousFunctionImpl(source, session, returnTypeRef, receiverTypeRef, symbol).apply {
+    return FirAnonymousFunctionImpl(source, session, returnTypeRef, receiverTypeRef, symbol, isLambda).apply {
         this.valueParameters.addAll(valueParameters)
         this.body = body
         this.annotations.addAll(annotations)

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

@@ -120,25 +120,27 @@ digraph returnValuesFromLambda_kt {
     subgraph cluster_10 {
       color=blue
       34 [label="Enter function anonymousFunction"];
-      35 [label="Jump: ^ Unit"];
+      35 [label="Jump: ^test_3 Unit"];
       36 [label="Stub" style="filled" fillcolor=gray];
-      37 [label="Exit function anonymousFunction"];
+      37 [label="Exit function anonymousFunction" style="filled" fillcolor=gray];
     }
-    38 [label="Function call: R|kotlin/run|<R|kotlin/Unit|>(<L> = run@fun <anonymous>(): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    38 [label="Function call: R|kotlin/run|<R|kotlin/Nothing|>(<L> = run@fun <anonymous>(): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
-    ^ Unit
+    ^test_3 Unit
 }
-)"];
+)" style="filled" fillcolor=gray];
-    39 [label="Variable declaration: lval x: R|kotlin/Unit|"];
+    39 [label="Stub" style="filled" fillcolor=gray];
-    40 [label="Exit function test_3" style="filled" fillcolor=red];
+    40 [label="Variable declaration: lval x: R|kotlin/Nothing|" style="filled" fillcolor=gray];
+    41 [label="Exit function test_3" style="filled" fillcolor=red];
   }
 
   33 -> {34};
   34 -> {35};
-  35 -> {37};
+  35 -> {41};
   35 -> {36} [style=dotted];
   36 -> {37} [style=dotted];
-  37 -> {38};
+  37 -> {38} [style=dotted];
-  38 -> {39};
+  38 -> {41 39} [style=dotted];
-  39 -> {40};
+  39 -> {40} [style=dotted];
+  40 -> {41} [style=dotted];
 
 }

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

@@ -32,8 +32,8 @@ FILE: returnValuesFromLambda.kt
         )
     }
     public final fun test_3(): R|kotlin/Unit| {
-        lval x: R|kotlin/Unit| = R|kotlin/run|<R|kotlin/Unit|>(<L> = run@fun <anonymous>(): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+        lval x: R|kotlin/Nothing| = R|kotlin/run|<R|kotlin/Nothing|>(<L> = run@fun <anonymous>(): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
-            ^ Unit
+            ^test_3 Unit
         }
         )
     }

compiler/fir/resolve/testData/resolve/smartcasts/safeCalls.dot

@@ -195,44 +195,45 @@ digraph safeCalls_kt {
     subgraph cluster_12 {
       color=blue
       71 [label="Enter function anonymousFunction"];
-      72 [label="Jump: ^ Unit"];
+      72 [label="Jump: ^test_5 Unit"];
       73 [label="Stub" style="filled" fillcolor=gray];
-      74 [label="Exit function anonymousFunction"];
+      74 [label="Exit function anonymousFunction" style="filled" fillcolor=gray];
     }
-    75 [label="Function call: R|<local>/x|?.R|kotlin/let|<R|A|, R|kotlin/Unit|>(<L> = let@fun <anonymous>(it: R|A|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    75 [label="Function call: R|<local>/x|?.R|kotlin/let|<R|A|, R|kotlin/Nothing|>(<L> = let@fun <anonymous>(it: R|A|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
-    ^ Unit
+    ^test_5 Unit
 }
-)"];
+)" style="filled" fillcolor=gray];
-    76 [label="Exit safe call"];
+    76 [label="Exit safe call" style="filled" fillcolor=gray];
-    77 [label="Enter safe call"];
+    77 [label="Enter safe call" style="filled" fillcolor=gray];
-    78 [label="Access variable R|<local>/x|"];
+    78 [label="Access variable R|<local>/x|" style="filled" fillcolor=gray];
-    79 [label="Function call: R|<local>/x|.R|/A.bool|()"];
+    79 [label="Function call: R|<local>/x|.R|/A.bool|()" style="filled" fillcolor=gray];
-    80 [label="Function call: R|<local>/x|?.R|kotlin/let|<R|A|, R|kotlin/Unit|>(<L> = let@fun <anonymous>(it: R|A|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    80 [label="Function call: R|<local>/x|?.R|kotlin/let|<R|A|, R|kotlin/Nothing|>(<L> = let@fun <anonymous>(it: R|A|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
-    ^ Unit
+    ^test_5 Unit
 }
-)?.R|/boo|(R|<local>/x|.R|/A.bool|())"];
+)?.R|/boo|(R|<local>/x|.R|/A.bool|())" style="filled" fillcolor=gray];
-    81 [label="Exit safe call"];
+    81 [label="Exit safe call" style="filled" fillcolor=gray];
-    82 [label="Access variable R|<local>/x|"];
+    82 [label="Access variable R|<local>/x|" style="filled" fillcolor=gray];
-    83 [label="Function call: R|<local>/x|.<Inapplicable(WRONG_RECEIVER): [/A.id]>#()"];
+    83 [label="Function call: R|<local>/x|.<Inapplicable(WRONG_RECEIVER): [/A.id]>#()" style="filled" fillcolor=gray];
     84 [label="Exit function test_5" style="filled" fillcolor=red];
   }
 
   68 -> {69};
-  69 -> {70 76};
+  69 -> {70};
+  69 -> {76} [style=dotted];
   70 -> {71};
   71 -> {72};
-  72 -> {74};
+  72 -> {84};
   72 -> {73} [style=dotted];
   73 -> {74} [style=dotted];
-  74 -> {75};
+  74 -> {75} [style=dotted];
-  75 -> {76};
+  75 -> {76} [style=dotted];
-  76 -> {77 81};
+  76 -> {77 81} [style=dotted];
-  77 -> {78};
+  77 -> {78} [style=dotted];
-  78 -> {79};
+  78 -> {79} [style=dotted];
-  79 -> {80};
+  79 -> {80} [style=dotted];
-  80 -> {81};
+  80 -> {81} [style=dotted];
-  81 -> {82};
+  81 -> {82} [style=dotted];
-  82 -> {83};
+  82 -> {83} [style=dotted];
-  83 -> {84};
+  83 -> {84} [style=dotted];
 
 }

compiler/fir/resolve/testData/resolve/smartcasts/safeCalls.txt

@@ -32,8 +32,8 @@ FILE: safeCalls.kt
     }
     public final fun R|kotlin/Any?|.boo(b: R|kotlin/Boolean|): R|kotlin/Unit|
     public final fun test_5(x: R|A?|): R|kotlin/Unit| {
-        R|<local>/x|?.R|kotlin/let|<R|A|, R|kotlin/Unit|>(<L> = let@fun <anonymous>(it: R|A|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+        R|<local>/x|?.R|kotlin/let|<R|A|, R|kotlin/Nothing|>(<L> = let@fun <anonymous>(it: R|A|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
-            ^ Unit
+            ^test_5 Unit
         }
         )?.R|/boo|(R|<local>/x|.R|/A.bool|())
         R|<local>/x|.<Inapplicable(WRONG_RECEIVER): [/A.id]>#()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirAnonymousFunction.kt

@@ -38,6 +38,7 @@ abstract class FirAnonymousFunction : FirPureAbstractElement(), FirFunction<FirA
     abstract override val symbol: FirAnonymousFunctionSymbol
     abstract val label: FirLabel?
     abstract val invocationKind: InvocationKind?
+    abstract val isLambda: Boolean
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitAnonymousFunction(this, data)
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/impl/FirAnonymousFunctionImpl.kt

@@ -33,7 +33,8 @@ class FirAnonymousFunctionImpl(
     override val session: FirSession,
     override var returnTypeRef: FirTypeRef,
     override var receiverTypeRef: FirTypeRef?,
-    override val symbol: FirAnonymousFunctionSymbol
+    override val symbol: FirAnonymousFunctionSymbol,
+    override val isLambda: Boolean
 ) : FirAnonymousFunction(), FirModifiableFunction<FirAnonymousFunction>, FirAbstractAnnotatedElement {
     override var resolvePhase: FirResolvePhase = FirResolvePhase.DECLARATIONS
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -253,6 +253,7 @@ object NodeConfigurator : AbstractFieldConfigurator() {
             +field(invocationKindType, nullable = true, withReplace = true).apply {
                 isMutable = true
             }
+            +booleanField("isLambda")
         }
 
         typeParameter.configure {

compiler/testData/ir/irText/lambdas/nonLocalReturn.fir.txt

@@ -5,7 +5,7 @@ FILE fqName:<root> fileName:/nonLocalReturn.kt
         block: FUN_EXPR type=kotlin.Function0<kotlin.Unit> origin=LAMBDA
           FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
             BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.test0'
+              RETURN type=kotlin.Nothing from='public final fun test0 (): kotlin.Unit declared in <root>'
                 GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
   FUN name:test1 visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/fir/FirVisualizerForRawFirDataGenerated.java

@@ -223,6 +223,11 @@ public class FirVisualizerForRawFirDataGenerated extends AbstractFirVisualizer {
             runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");
         }
 
+        @TestMetadata("lambdaAndAnonymousFunction.kt")
+        public void testLambdaAndAnonymousFunction() throws Exception {
+            runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
+        }
+
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {
             runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/locals.kt");

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/psi/PsiVisualizerForRawFirDataGenerated.java

@@ -223,6 +223,11 @@ public class PsiVisualizerForRawFirDataGenerated extends AbstractPsiVisualizer {
             runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");
         }
 
+        @TestMetadata("lambdaAndAnonymousFunction.kt")
+        public void testLambdaAndAnonymousFunction() throws Exception {
+            runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
+        }
+
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {
             runTest("compiler/fir/psi2fir/testData/rawBuilder/expressions/locals.kt");