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/fir/fir2ir/tests/org/jetbrains/kotlin/fir/Fir2IrTextTestGenerated.java

@@ -1286,6 +1286,11 @@ public class Fir2IrTextTestGenerated extends AbstractFir2IrTextTest {
             runTest("compiler/testData/ir/irText/expressions/whenReturn.kt");
         }
 
+        @TestMetadata("whenReturnUnit.kt")
+        public void testWhenReturnUnit() throws Exception {
+            runTest("compiler/testData/ir/irText/expressions/whenReturnUnit.kt");
+        }
+
         @TestMetadata("whenUnusedExpression.kt")
         public void testWhenUnusedExpression() throws Exception {
             runTest("compiler/testData/ir/irText/expressions/whenUnusedExpression.kt");

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/BodyGenerator.kt

@@ -31,6 +31,8 @@ import org.jetbrains.kotlin.psi.psiUtil.startOffsetSkippingComments
 import org.jetbrains.kotlin.psi.synthetics.findClassDescriptor
 import org.jetbrains.kotlin.psi2ir.intermediate.VariableLValue
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
+import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsResultOfLambda
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassOrAny
 import org.jetbrains.kotlin.types.KotlinType
@@ -56,7 +58,13 @@ class BodyGenerator(
         if (ktBody is KtBlockExpression) {
             statementGenerator.generateStatements(ktBody.statements, irBlockBody)
         } else {
-            statementGenerator.generateReturnExpression(ktBody, irBlockBody)
+            val irBody = statementGenerator.generateStatement(ktBody)
+            irBlockBody.statements.add(
+                if (ktBody.isUsedAsExpression(context.bindingContext) && irBody is IrExpression)
+                    generateReturnExpression(irBody.startOffset, irBody.endOffset, irBody)
+                else
+                    irBody
+            )
         }
 
         return irBlockBody
@@ -90,7 +98,14 @@ class BodyGenerator(
                 irBlockBody.statements.add(statementGenerator.generateStatement(ktStatement))
             }
             val ktReturnedValue = ktBodyStatements.last()
-            statementGenerator.generateReturnExpression(ktReturnedValue, irBlockBody)
+            val irReturnedValue = statementGenerator.generateStatement(ktReturnedValue)
+            irBlockBody.statements.add(
+                if (ktReturnedValue.isUsedAsResultOfLambda(context.bindingContext) && irReturnedValue is IrExpression) {
+                    generateReturnExpression(irReturnedValue.startOffset, irReturnedValue.endOffset, irReturnedValue)
+                } else {
+                    irReturnedValue
+                }
+            )
         } else {
             irBlockBody.statements.add(
                 generateReturnExpression(
@@ -106,23 +121,6 @@ class BodyGenerator(
         return irBlockBody
     }
 
-    private fun StatementGenerator.generateReturnExpression(ktExpression: KtExpression, irBlockBody: IrBlockBodyImpl) {
-        val irReturnExpression = generateStatement(ktExpression)
-        if (irReturnExpression is IrExpression) {
-            irBlockBody.statements.add(irReturnExpression.wrapWithReturn())
-        } else {
-            irBlockBody.statements.add(irReturnExpression)
-        }
-    }
-
-    private fun IrExpression.wrapWithReturn() =
-        if (this is IrReturn || this is IrErrorExpression || this is IrThrow)
-            this
-        else {
-            generateReturnExpression(startOffset, endOffset, this)
-        }
-
-
     private fun generateReturnExpression(startOffset: Int, endOffset: Int, returnValue: IrExpression): IrReturnImpl {
         val returnTarget = (scopeOwner as? CallableDescriptor) ?: throw AssertionError("'return' in a non-callable: $scopeOwner")
         return IrReturnImpl(
@@ -132,7 +130,6 @@ class BodyGenerator(
         )
     }
 
-
     fun generateSecondaryConstructorBody(ktConstructor: KtSecondaryConstructor): IrBody {
         val irBlockBody = IrBlockBodyImpl(ktConstructor.startOffsetSkippingComments, ktConstructor.endOffset)
 
@@ -319,7 +316,7 @@ class BodyGenerator(
         }
 
         val enumDefaultConstructorCall = getResolvedCall(ktEnumEntry)
-                ?: throw AssertionError("No default constructor call for enum entry $enumClassDescriptor")
+            ?: throw AssertionError("No default constructor call for enum entry $enumClassDescriptor")
         return statementGenerator.generateEnumConstructorCall(enumDefaultConstructorCall, ktEnumEntry)
     }
 

compiler/testData/codegen/boxInline/smap/anonymous/lambda.kt

@@ -1,4 +1,4 @@
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
+
 // FILE: 1.kt
 
 package builders

compiler/testData/codegen/boxInline/smap/anonymous/object.kt

@@ -1,6 +1,6 @@
 // FILE: 1.kt
 
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
+
 package builders
 
 inline fun call(crossinline init: () -> Unit) {

compiler/testData/codegen/boxInline/smap/newsmap/mappingInInlineFunLambda.kt

@@ -1,4 +1,4 @@
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
+
 // FILE: 1.kt
 
 package test

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"

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

@@ -33,14 +33,13 @@ 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>.test3'
-                CALL 'public final fun run <R> (block: kotlin.Function0<R of kotlin.run>): R of kotlin.run [inline] declared in kotlin' type=kotlin.Nothing origin=null
-                  <R>: kotlin.Nothing
-                  block: FUN_EXPR type=kotlin.Function0<kotlin.Nothing> origin=LAMBDA
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Nothing
-                      BLOCK_BODY
-                        RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.test3'
-                          GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+              CALL 'public final fun run <R> (block: kotlin.Function0<R of kotlin.run>): R of kotlin.run [inline] declared in kotlin' type=kotlin.Nothing origin=null
+                <R>: kotlin.Nothing
+                block: FUN_EXPR type=kotlin.Function0<kotlin.Nothing> origin=LAMBDA
+                  FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Nothing
+                    BLOCK_BODY
+                      RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.test3'
+                        GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
   FUN name:testLrmFoo1 visibility:public modality:FINAL <> (ints:kotlin.collections.List<kotlin.Int>) returnType:kotlin.Unit
     VALUE_PARAMETER name:ints index:0 type:kotlin.collections.List<kotlin.Int>
     BLOCK_BODY
@@ -58,9 +57,8 @@ FILE fqName:<root> fileName:/nonLocalReturn.kt
                     arg1: CONST Int type=kotlin.Int value=0
                   then: RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testLrmFoo1'
                     GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testLrmFoo1'
-                CALL 'public final fun print (message: kotlin.Int): kotlin.Unit [inline] declared in kotlin.io' type=kotlin.Unit origin=null
-                  message: GET_VAR 'it: kotlin.Int declared in <root>.testLrmFoo1.<anonymous>' type=kotlin.Int origin=null
+              CALL 'public final fun print (message: kotlin.Int): kotlin.Unit [inline] declared in kotlin.io' type=kotlin.Unit origin=null
+                message: GET_VAR 'it: kotlin.Int declared in <root>.testLrmFoo1.<anonymous>' type=kotlin.Int origin=null
   FUN name:testLrmFoo2 visibility:public modality:FINAL <> (ints:kotlin.collections.List<kotlin.Int>) returnType:kotlin.Unit
     VALUE_PARAMETER name:ints index:0 type:kotlin.collections.List<kotlin.Int>
     BLOCK_BODY
@@ -78,6 +76,5 @@ FILE fqName:<root> fileName:/nonLocalReturn.kt
                     arg1: CONST Int type=kotlin.Int value=0
                   then: RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testLrmFoo2'
                     GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testLrmFoo2'
-                CALL 'public final fun print (message: kotlin.Int): kotlin.Unit [inline] declared in kotlin.io' type=kotlin.Unit origin=null
-                  message: GET_VAR 'it: kotlin.Int declared in <root>.testLrmFoo2.<anonymous>' type=kotlin.Int origin=null
+              CALL 'public final fun print (message: kotlin.Int): kotlin.Unit [inline] declared in kotlin.io' type=kotlin.Unit origin=null
+                message: GET_VAR 'it: kotlin.Int declared in <root>.testLrmFoo2.<anonymous>' type=kotlin.Int origin=null

compiler/testData/ir/irText/lambdas/samAdapter.txt

@@ -6,8 +6,7 @@ FILE fqName:<root> fileName:/samAdapter.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 println (message: kotlin.Any?): kotlin.Unit [inline] declared in kotlin.io' type=kotlin.Unit origin=null
-                    message: CONST String type=kotlin.String value="Hello, world!"
+                CALL 'public final fun println (message: kotlin.Any?): kotlin.Unit [inline] declared in kotlin.io' type=kotlin.Unit origin=null
+                  message: CONST String type=kotlin.String value="Hello, world!"
       CALL 'public abstract fun run (): kotlin.Unit declared in java.lang.Runnable' type=kotlin.Unit origin=null
         $this: GET_VAR 'val hello: java.lang.Runnable [val] declared in <root>.test1' type=java.lang.Runnable origin=null

compiler/testData/ir/irText/regressions/integerCoercionToT.txt

@@ -18,8 +18,7 @@ FILE fqName:<root> fileName:/integerCoercionToT.kt
     TYPE_PARAMETER name:T index:0 variance: superTypes:[<root>.CPointed]
     $receiver: VALUE_PARAMETER name:<this> type:<root>.CPointed
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun reinterpret <T> (): T of <root>.reinterpret [inline] 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
   CLASS CLASS name:CInt32VarX modality:FINAL visibility:public superTypes:[<root>.CPointed]
     $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CInt32VarX<T of <root>.CInt32VarX>
     TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
@@ -47,8 +46,7 @@ FILE fqName:<root> fileName:/integerCoercionToT.kt
       TYPE_PARAMETER name:T_INT index:0 variance: superTypes:[kotlin.Int]
       $receiver: VALUE_PARAMETER name:<this> type:<root>.CInt32VarX<T_INT of <root>.<get-value>>
       BLOCK_BODY
-        RETURN type=kotlin.Nothing from='public final fun <get-value> <T_INT> (): T_INT of <root>.<get-value> 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:<set-value> visibility:public modality:FINAL <T_INT> ($receiver:<root>.CInt32VarX<T_INT of <root>.<set-value>>, value:T_INT of <root>.<set-value>) returnType:kotlin.Unit
       correspondingProperty: PROPERTY name:value visibility:public modality:FINAL [var]
       TYPE_PARAMETER name:T_INT index:0 variance: superTypes:[kotlin.Int]

compiler/testData/ir/irText/regressions/newInference/fixationOrder1.txt

@@ -3,8 +3,7 @@ FILE fqName:<root> fileName:/fixationOrder1.kt
     TYPE_PARAMETER name:X index:0 variance: superTypes:[kotlin.Any?]
     TYPE_PARAMETER name:Y index:1 variance: superTypes:[kotlin.Any?]
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun foo <X, Y> (): kotlin.Function1<X of <root>.foo, Y of <root>.foo> 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
   CLASS INTERFACE name:Inv2 modality:ABSTRACT visibility:public superTypes:[kotlin.Any]
     $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.Inv2<A of <root>.Inv2, B of <root>.Inv2>
     TYPE_PARAMETER name:A index:0 variance: superTypes:[kotlin.Any?]
@@ -29,8 +28,7 @@ FILE fqName:<root> fileName:/fixationOrder1.kt
     VALUE_PARAMETER name:y index:1 type:R of <root>.check
     VALUE_PARAMETER name:f index:2 type:kotlin.Function1<T of <root>.check, R of <root>.check>
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun check <T, R> (x: T of <root>.check, y: R of <root>.check, f: kotlin.Function1<T of <root>.check, R of <root>.check>): <root>.Inv2<T of <root>.check, R of <root>.check> 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:test visibility:public modality:FINAL <> () returnType:<root>.Inv2<kotlin.String, kotlin.Int>
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun test (): <root>.Inv2<kotlin.String, kotlin.Int> declared in <root>'

compiler/testData/ir/irText/stubs/builtinMap.txt

@@ -26,11 +26,10 @@ FILE fqName:<root> fileName:/builtinMap.kt
                 FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:java.util.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?>{ kotlin.collections.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?> }) returnType:kotlin.Unit
                   $receiver: VALUE_PARAMETER name:<this> type:java.util.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?>{ kotlin.collections.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?> }
                   BLOCK_BODY
-                    RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.plus'
-                      TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                        CALL 'public open fun put (key: K of java.util.LinkedHashMap, value: V of java.util.LinkedHashMap): V of java.util.LinkedHashMap? declared in java.util.LinkedHashMap' type=V1 of <root>.plus? origin=null
-                          $this: GET_VAR '<this>: java.util.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?>{ kotlin.collections.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?> } declared in <root>.plus.<anonymous>' type=java.util.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?>{ kotlin.collections.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?> } origin=null
-                          key: CALL 'public final fun <get-first> (): A of kotlin.Pair declared in kotlin.Pair' type=K1 of <root>.plus origin=GET_PROPERTY
-                            $this: GET_VAR 'pair: kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> declared in <root>.plus' type=kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> origin=null
-                          value: CALL 'public final fun <get-second> (): B of kotlin.Pair declared in kotlin.Pair' type=V1 of <root>.plus origin=GET_PROPERTY
-                            $this: GET_VAR 'pair: kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> declared in <root>.plus' type=kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> origin=null
+                    TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                      CALL 'public open fun put (key: K of java.util.LinkedHashMap, value: V of java.util.LinkedHashMap): V of java.util.LinkedHashMap? declared in java.util.LinkedHashMap' type=V1 of <root>.plus? origin=null
+                        $this: GET_VAR '<this>: java.util.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?>{ kotlin.collections.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?> } declared in <root>.plus.<anonymous>' type=java.util.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?>{ kotlin.collections.LinkedHashMap<K1 of <root>.plus?, V1 of <root>.plus?> } origin=null
+                        key: CALL 'public final fun <get-first> (): A of kotlin.Pair declared in kotlin.Pair' type=K1 of <root>.plus origin=GET_PROPERTY
+                          $this: GET_VAR 'pair: kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> declared in <root>.plus' type=kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> origin=null
+                        value: CALL 'public final fun <get-second> (): B of kotlin.Pair declared in kotlin.Pair' type=V1 of <root>.plus origin=GET_PROPERTY
+                          $this: GET_VAR 'pair: kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> declared in <root>.plus' type=kotlin.Pair<K1 of <root>.plus, V1 of <root>.plus> origin=null

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -1286,6 +1286,11 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             runTest("compiler/testData/ir/irText/expressions/whenReturn.kt");
         }
 
+        @TestMetadata("whenReturnUnit.kt")
+        public void testWhenReturnUnit() throws Exception {
+            runTest("compiler/testData/ir/irText/expressions/whenReturnUnit.kt");
+        }
+
         @TestMetadata("whenUnusedExpression.kt")
         public void testWhenUnusedExpression() throws Exception {
             runTest("compiler/testData/ir/irText/expressions/whenUnusedExpression.kt");