Refine 'handleResult' calls generation within coroutines

Before this change everything works just fine for 'handleResult' methods
accepting non-Unit parameters

For other cases the same coercion-to-unit strategy is in plain lambdas:
- if last statement is not Unit type, execute it, pop from the stack, then put Unit instance
- for 'return@label' (no expression) just put Unit on the stack

 #KT-13531 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -1839,27 +1839,54 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     @Nullable
-    private StackValue genControllerHandleResultCallIfNeeded(@NotNull KtExpression callOwner, @Nullable KtExpression valueToReturn) {
+    private StackValue genControllerHandleResultCallIfNeeded(@NotNull KtExpression callOwner, @Nullable KtExpression returnValue) {
         ResolvedCall<FunctionDescriptor> resolvedCall = bindingContext.get(RETURN_HANDLE_RESULT_RESOLVED_CALL, callOwner);
 
         if (resolvedCall != null) {
             assert resolvedCall.getValueArgumentsByIndex() != null : "Arguments were not resolved for call element: " + callOwner.getText();
             KtExpression argumentExpression =
                     resolvedCall.getValueArgumentsByIndex().get(0).getArguments().get(0).getArgumentExpression();
-            if (valueToReturn == null
+
-                    && KotlinBuiltIns.isUnit(resolvedCall.getResultingDescriptor().getValueParameters().get(0).getType())) {
+            final StackValue putValueBeforeCall;
+            // This condition may be true in cases like return-statement without value or for last statement in a lambda block that
+            // has a type different from Unit, while 'handleResult' method accepts exactly the latter
+            // (see org.jetbrains.kotlin.coroutines.resolveCoroutineHandleResultCallIfNeeded for clarifications)
+            if (argumentExpression != returnValue) {
+                assert KotlinBuiltIns.isUnit(resolvedCall.getResultingDescriptor().getValueParameters().get(0).getType())
+                        : "If handleResult argument is different from returnValue, handleResult's first parameter must accept Unit, but " +
+                          resolvedCall.getResultingDescriptor() + " was found";
+
+                // generate last statement in the coroutine lambda
+                putValueBeforeCall = returnValue != null ? genStatement(returnValue) : null;
+
+                // Here 'argumentExpression' is a special fake one that used as an expression of Unit type
+                // when 'handleResult' call was resolved
                 tempVariables.put(argumentExpression, StackValue.unit());
             }
             else {
-                assert valueToReturn != null : "valueReturn expected to be not null for non-unit types";
+                putValueBeforeCall = null;
-                tempVariables.put(argumentExpression, gen(valueToReturn));
             }
 
             tempVariables.put(
                     resolvedCall.getValueArgumentsByIndex().get(1).getArguments().get(0).getArgumentExpression(),
                     genCoroutineInstanceValueFromResolvedCall(resolvedCall));
 
-            return invokeFunction(resolvedCall, StackValue.none());
+
+            final StackValue handleResultCallValue = invokeFunction(resolvedCall, StackValue.none());
+            if (putValueBeforeCall == null) return handleResultCallValue;
+
+            return new StackValue(handleResultCallValue.type) {
+                @Override
+                public void putSelector(@NotNull Type type, @NotNull InstructionAdapter v) {
+                    putValueBeforeCall.put(type, v);
+                    handleResultCallValue.putSelector(type, v);
+                }
+
+                @Override
+                public void putReceiver(@NotNull InstructionAdapter v, boolean isRead) {
+                    handleResultCallValue.putReceiver(v, isRead);
+                }
+            };
         }
         return null;
     }

compiler/testData/codegen/box/coroutines/coercionToUnit.kt

@@ -0,0 +1,46 @@
+class Controller {
+    var result = "fail"
+    operator fun handleResult(u: Unit, c: Continuation<Nothing>) {
+        result = "OK"
+    }
+
+    suspend fun <T> await(t: T, c: Continuation<T>) {
+        c.resume(t)
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>): String {
+    val controller = Controller()
+    c(controller).resume(Unit)
+    return controller.result
+}
+
+var TRUE = true
+var FALSE = false
+fun box(): String {
+    val r1 = builder { await(Unit) }
+    if (r1 != "OK") return "fail 1"
+
+    val r2 = builder {
+        if (await(1) != 1) throw RuntimeException("fail1")
+
+        if (TRUE) return@builder
+    }
+    if (r2 != "OK") return "fail 2"
+
+    val r3 = builder {
+        if (await(1) != 1) throw RuntimeException("fail2")
+
+        if (FALSE) return@builder
+    }
+    if (r3 != "OK") return "fail 3"
+
+    val r4 = builder {
+        if (await(1) != 1) throw RuntimeException("fail3")
+
+        return@builder
+    }
+    if (r4 != "OK") return "fail 4"
+
+    return builder { await(1) }
+}

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

@@ -4324,6 +4324,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/coroutines"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("coercionToUnit.kt")
+        public void testCoercionToUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coercionToUnit.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("controllerAccessFromInnerLambda.kt")
         public void testControllerAccessFromInnerLambda() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controllerAccessFromInnerLambda.kt");