From a420fda825dbb45022c90e69c553a2f095e588a4 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Mon, 4 Jul 2016 18:07:16 +0300
Subject: [PATCH] Fix problem with missing handleResult call

At first we try to resolve 'handleResult' just as last expression
in a lambda is first argument, then if results are unsuccessful
try resolve 'handleResult' with fake Unit expression

 #KT-12969 Fixed
---
 .../kotlin/coroutines/coroutineUtil.kt        | 33 +++++++++++++------
 .../handleResultNonUnitExpression.kt          | 28 ++++++++++++++++
 .../codegen/BlackBoxCodegenTestGenerated.java |  6 ++++
 3 files changed, 57 insertions(+), 10 deletions(-)
 create mode 100644 compiler/testData/codegen/box/coroutines/handleResultNonUnitExpression.kt

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/coroutines/coroutineUtil.kt b/compiler/frontend/src/org/jetbrains/kotlin/coroutines/coroutineUtil.kt
index d675495b831..fd462ae59ca 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/coroutines/coroutineUtil.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/coroutines/coroutineUtil.kt
@@ -72,21 +72,34 @@ fun FakeCallResolver.resolveCoroutineHandleResultCallIfNeeded(
                     // should be Continuation<Nothing>
                     functionDescriptor.builtIns.nothingType)
 
+    fun tryToResolveCall(firstArgument: KtExpression): Boolean {
+        val resolutionResults = resolveFakeCall(
+                context.replaceBindingTrace(temporaryBindingTrace), functionDescriptor.extensionReceiverParameter!!.value,
+                OperatorNameConventions.COROUTINE_HANDLE_RESULT, callElement, callElement, FakeCallKind.OTHER,
+                listOf(firstArgument, continuation))
+
+        if (resolutionResults.isSuccess && resolutionResults.resultingDescriptor.isOperator) {
+            context.trace.record(BindingContext.RETURN_HANDLE_RESULT_RESOLVED_CALL, callElement, resolutionResults.resultingCall)
+            return true
+        }
+
+        return false
+    }
+
+    val unitExpression = ExpressionTypingUtils.createFakeExpressionOfType(
+            callElement.project, temporaryBindingTrace, "unit", functionDescriptor.builtIns.unitType)
     val firstArgument =
             if (expressionToReturn == null || info != null && info.type != null && KotlinBuiltIns.isUnit(info.type))
-                ExpressionTypingUtils.createFakeExpressionOfType(
-                        callElement.project, temporaryBindingTrace, "unit", functionDescriptor.builtIns.unitType)
-            else expressionToReturn
+                unitExpression
+            else
+                expressionToReturn
 
-    val resolutionResults = resolveFakeCall(
-            context.replaceBindingTrace(temporaryBindingTrace), functionDescriptor.extensionReceiverParameter!!.value,
-            OperatorNameConventions.COROUTINE_HANDLE_RESULT, callElement, callElement, FakeCallKind.OTHER,
-            listOf(firstArgument, continuation))
-
-    if (resolutionResults.isSuccess && resolutionResults.resultingDescriptor.isOperator) {
-        context.trace.record(BindingContext.RETURN_HANDLE_RESULT_RESOLVED_CALL, callElement, resolutionResults.resultingCall)
+    if (!tryToResolveCall(firstArgument) && firstArgument === expressionToReturn) {
+        tryToResolveCall(unitExpression)
     }
 }
 
+
+
 fun KotlinType.isValidContinuation() =
         (constructor.declarationDescriptor as? ClassDescriptor)?.fqNameUnsafe == DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME.toUnsafe()
diff --git a/compiler/testData/codegen/box/coroutines/handleResultNonUnitExpression.kt b/compiler/testData/codegen/box/coroutines/handleResultNonUnitExpression.kt
new file mode 100644
index 00000000000..9da50876350
--- /dev/null
+++ b/compiler/testData/codegen/box/coroutines/handleResultNonUnitExpression.kt
@@ -0,0 +1,28 @@
+class Controller {
+    var isCompleted = false
+    suspend fun suspendHere(x: Continuation<String>) {
+        x.resume("OK")
+    }
+
+    operator fun handleResult(x: Unit, y: Continuation<Nothing>) {
+        isCompleted = true
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    val controller = Controller()
+    c(controller).resume(Unit)
+    if (!controller.isCompleted) throw java.lang.RuntimeException("fail")
+}
+
+fun box(): String {
+    builder {
+        "OK"
+    }
+
+    builder {
+        suspendHere()
+    }
+
+    return "OK"
+}
diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java
index 31b54d0ba9a..67bd9ba19d5 100644
--- a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java
+++ b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java
@@ -4189,6 +4189,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("handleResultNonUnitExpression.kt")
+        public void testHandleResultNonUnitExpression() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/handleResultNonUnitExpression.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("illegalState.kt")
         public void testIllegalState() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/illegalState.kt");