psi2ir: Fix return insertion

We should only insert a return statement at the end of a lambda or
function if the final statement is used as an expression (slice
USED_AS_RESULT_OF_LAMBDA and USED_AS_EXPRESSION).

compiler/testData/ir/irText/expressions/argumentMappedWithError.txt

@@ -3,8 +3,7 @@ FILE fqName:<root> fileName:/argumentMappedWithError.kt
     TYPE_PARAMETER name:R index:0 variance: superTypes:[kotlin.Number]
     $receiver: VALUE_PARAMETER name:<this> type:kotlin.Number
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun convert <R> (): R of <root>.convert declared in <root>'
-        CALL 'public final fun TODO (): kotlin.Nothing [inline] declared in kotlin' type=kotlin.Nothing origin=null
+      CALL 'public final fun TODO (): kotlin.Nothing [inline] declared in kotlin' type=kotlin.Nothing origin=null
   FUN name:foo visibility:public modality:FINAL <> (arg:kotlin.Number) returnType:kotlin.Unit
     VALUE_PARAMETER name:arg index:0 type:kotlin.Number
     BLOCK_BODY

compiler/testData/ir/irText/expressions/coercionToUnit.txt

@@ -5,9 +5,8 @@ FILE fqName:<root> fileName:/coercionToUnit.kt
         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>.test1'
-                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                  CONST Int type=kotlin.Int value=42
+              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                CONST Int type=kotlin.Int value=42
     FUN DEFAULT_PROPERTY_ACCESSOR name:<get-test1> visibility:public modality:FINAL <> () returnType:kotlin.Function0<kotlin.Unit>
       correspondingProperty: PROPERTY name:test1 visibility:public modality:FINAL [val]
       BLOCK_BODY

compiler/testData/ir/irText/expressions/enumEntryReferenceFromEnumEntryClass.txt

@@ -53,9 +53,8 @@ FILE fqName:<root> fileName:/enumEntryReferenceFromEnumEntryClass.kt
                     CALL 'public final fun <set-counter> (<set-?>: kotlin.Int): kotlin.Unit declared in <root>.MyEnum.Z' type=kotlin.Unit origin=EQ
                       $this: GET_ENUM 'ENUM_ENTRY name:Z' type=<root>.MyEnum.Z
                       <set-?>: CONST Int type=kotlin.Int value=1
-                    RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.MyEnum.Z.aLambda'
-                      CALL 'public final fun foo (): kotlin.Unit declared in <root>.MyEnum.Z' type=kotlin.Unit origin=null
-                        $this: GET_ENUM 'ENUM_ENTRY name:Z' type=<root>.MyEnum.Z
+                    CALL 'public final fun foo (): kotlin.Unit declared in <root>.MyEnum.Z' type=kotlin.Unit origin=null
+                      $this: GET_ENUM 'ENUM_ENTRY name:Z' type=<root>.MyEnum.Z
           FUN DEFAULT_PROPERTY_ACCESSOR name:<get-aLambda> visibility:public modality:FINAL <> ($this:<root>.MyEnum.Z) returnType:kotlin.Function0<kotlin.Unit>
             correspondingProperty: PROPERTY name:aLambda visibility:public modality:FINAL [val]
             $this: VALUE_PARAMETER name:<this> type:<root>.MyEnum.Z

compiler/testData/ir/irText/expressions/objectReference.txt

@@ -98,9 +98,8 @@ FILE fqName:<root> fileName:/objectReference.kt
                 CALL 'public final fun <set-counter> (<set-?>: kotlin.Int): kotlin.Unit declared in <root>.Z' type=kotlin.Unit origin=EQ
                   $this: GET_OBJECT 'CLASS OBJECT name:Z modality:FINAL visibility:public superTypes:[kotlin.Any]' type=<root>.Z
                   <set-?>: CONST Int type=kotlin.Int value=1
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.Z.aLambda'
-                  CALL 'public final fun foo (): kotlin.Unit declared in <root>.Z' type=kotlin.Unit origin=null
-                    $this: GET_OBJECT 'CLASS OBJECT name:Z modality:FINAL visibility:public superTypes:[kotlin.Any]' type=<root>.Z
+                CALL 'public final fun foo (): kotlin.Unit declared in <root>.Z' type=kotlin.Unit origin=null
+                  $this: GET_OBJECT 'CLASS OBJECT name:Z modality:FINAL visibility:public superTypes:[kotlin.Any]' type=<root>.Z
       FUN DEFAULT_PROPERTY_ACCESSOR name:<get-aLambda> visibility:public modality:FINAL <> ($this:<root>.Z) returnType:kotlin.Function0<kotlin.Unit>
         correspondingProperty: PROPERTY name:aLambda visibility:public modality:FINAL [val]
         $this: VALUE_PARAMETER name:<this> type:<root>.Z

compiler/testData/ir/irText/expressions/sam/samConversions.txt

@@ -15,8 +15,7 @@ FILE fqName:<root> fileName:/samConversions.kt
           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>.test1'
-                  CALL 'public final fun test1 (): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                CALL 'public final fun test1 (): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
   FUN name:test2 visibility:public modality:FINAL <> ($receiver:<root>.J) returnType:kotlin.Unit
     $receiver: VALUE_PARAMETER name:<this> type:<root>.J
     BLOCK_BODY
@@ -26,8 +25,7 @@ FILE fqName:<root> fileName:/samConversions.kt
           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>.test2'
-                  CALL 'public final fun test1 (): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                CALL 'public final fun test1 (): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
   FUN name:test3 visibility:public modality:FINAL <> ($receiver:<root>.J, a:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
     $receiver: VALUE_PARAMETER name:<this> type:<root>.J
     VALUE_PARAMETER name:a index:0 type:kotlin.Function0<kotlin.Unit>

compiler/testData/ir/irText/expressions/whenReturnUnit.fir.txt

@@ -0,0 +1,27 @@
+FILE fqName:<root> fileName:/whenReturnUnit.kt
+  FUN name:run visibility:public modality:FINAL <> (block:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
+    VALUE_PARAMETER name:block index:0 type:kotlin.Function0<kotlin.Unit>
+    BLOCK_BODY
+  FUN name:branch visibility:public modality:FINAL <> (x:kotlin.Int) returnType:kotlin.Unit
+    VALUE_PARAMETER name:x index:0 type:kotlin.Int
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun branch (x: kotlin.Int): kotlin.Unit declared in <root>'
+        CALL 'public final fun run (block: kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+          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
+                BLOCK type=kotlin.Unit origin=WHEN
+                  VAR IR_TEMPORARY_VARIABLE name:tmp0_subject type:kotlin.Int [val]
+                  WHEN type=kotlin.Unit origin=WHEN
+                    BRANCH
+                      if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                        arg0: GET_VAR 'val tmp0_subject: kotlin.Int [val] declared in <root>.branch.<anonymous>' type=kotlin.Int origin=null
+                        arg1: CONST Int type=kotlin.Int value=1
+                      then: CALL 'public final fun TODO (reason: kotlin.String): kotlin.Nothing [inline] declared in kotlin' type=kotlin.Nothing origin=null
+                        reason: CONST String type=kotlin.String value="1"
+                    BRANCH
+                      if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                        arg0: GET_VAR 'val tmp0_subject: kotlin.Int [val] declared in <root>.branch.<anonymous>' type=kotlin.Int origin=null
+                        arg1: CONST Int type=kotlin.Int value=2
+                      then: CALL 'public final fun TODO (reason: kotlin.String): kotlin.Nothing [inline] declared in kotlin' type=kotlin.Nothing origin=null
+                        reason: CONST String type=kotlin.String value="2"

compiler/testData/ir/irText/expressions/whenReturnUnit.kt

@@ -0,0 +1,8 @@
+fun run(block: () -> Unit) {}
+
+fun branch(x: Int) = run {
+    when (x) {
+        1 -> TODO("1")
+        2 -> TODO("2")
+    }
+}

compiler/testData/ir/irText/expressions/whenReturnUnit.txt

@@ -0,0 +1,28 @@
+FILE fqName:<root> fileName:/whenReturnUnit.kt
+  FUN name:run visibility:public modality:FINAL <> (block:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
+    VALUE_PARAMETER name:block index:0 type:kotlin.Function0<kotlin.Unit>
+    BLOCK_BODY
+  FUN name:branch visibility:public modality:FINAL <> (x:kotlin.Int) returnType:kotlin.Unit
+    VALUE_PARAMETER name:x index:0 type:kotlin.Int
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun branch (x: kotlin.Int): kotlin.Unit declared in <root>'
+        CALL 'public final fun run (block: kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+          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
+                BLOCK type=kotlin.Nothing origin=WHEN
+                  VAR IR_TEMPORARY_VARIABLE name:tmp0_subject type:kotlin.Int [val]
+                    GET_VAR 'x: kotlin.Int declared in <root>.branch' type=kotlin.Int origin=null
+                  WHEN type=kotlin.Nothing origin=WHEN
+                    BRANCH
+                      if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                        arg0: GET_VAR 'val tmp0_subject: kotlin.Int [val] declared in <root>.branch.<anonymous>' type=kotlin.Int origin=null
+                        arg1: CONST Int type=kotlin.Int value=1
+                      then: CALL 'public final fun TODO (reason: kotlin.String): kotlin.Nothing [inline] declared in kotlin' type=kotlin.Nothing origin=null
+                        reason: CONST String type=kotlin.String value="1"
+                    BRANCH
+                      if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                        arg0: GET_VAR 'val tmp0_subject: kotlin.Int [val] declared in <root>.branch.<anonymous>' type=kotlin.Int origin=null
+                        arg1: CONST Int type=kotlin.Int value=2
+                      then: CALL 'public final fun TODO (reason: kotlin.String): kotlin.Nothing [inline] declared in kotlin' type=kotlin.Nothing origin=null
+                        reason: CONST String type=kotlin.String value="2"