[NI] Fix hierarchy of resolution atoms for lambda with non-local return

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt

@@ -123,18 +123,17 @@ class KotlinResolutionCallbacksImpl(
             }
         }
 
-        if (!hasReturnWithoutExpression) {
-            val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
-                if (expectedReturnType?.isUnit() == true) return@let null // coercion to Unit
+        val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
+            if (expectedReturnType?.isUnit() == true) return@let null // coercion to Unit
 
-                // todo lastExpression can be if without else
-                val lastExpressionType = trace.getType(lastExpression)
-                val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
-                createCallArgument(lastExpression, lastExpressionTypeInfo)
-            }
-            returnArguments.addIfNotNull(lastExpressionArgument)
+            // todo lastExpression can be if without else
+            val lastExpressionType = if (hasReturnWithoutExpression) null else trace.getType(lastExpression)
+            val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
+            createCallArgument(lastExpression, lastExpressionTypeInfo)
         }
 
+        returnArguments.addIfNotNull(lastExpressionArgument)
+
         return returnArguments
     }
 

compiler/testData/codegen/box/nonLocalReturns/returnInsideTwoLambdas.kt

@@ -0,0 +1,28 @@
+fun <T> outer(command: () -> T) : T = command()
+
+inline fun <K> inner(action: () -> K): K = action()
+
+fun test1(): String {
+    outer {
+        inner {
+            return@outer
+        }
+    }
+
+    return "O"
+}
+
+fun test2(): String {
+    outer {
+        return@outer
+        inner {
+        }
+    }
+
+    return "K"
+}
+
+
+fun box(): String {
+    return test1() + test2()
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -12422,6 +12422,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("returnInsideTwoLambdas.kt")
+        public void testReturnInsideTwoLambdas() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/returnInsideTwoLambdas.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("use.kt")
         public void testUse() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/use.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -12422,6 +12422,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("returnInsideTwoLambdas.kt")
+        public void testReturnInsideTwoLambdas() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/returnInsideTwoLambdas.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("use.kt")
         public void testUse() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/use.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -12422,6 +12422,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("returnInsideTwoLambdas.kt")
+        public void testReturnInsideTwoLambdas() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/returnInsideTwoLambdas.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("use.kt")
         public void testUse() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/use.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -13484,6 +13484,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("returnInsideTwoLambdas.kt")
+        public void testReturnInsideTwoLambdas() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/returnInsideTwoLambdas.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("use.kt")
         public void testUse() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/nonLocalReturns/use.kt");