Resolve all return expressions in coroutine as fake calls to 'handleResult'

compiler/frontend/src/org/jetbrains/kotlin/coroutines/coroutineUtil.kt

@@ -16,10 +16,20 @@
 
 package org.jetbrains.kotlin.coroutines
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
+import org.jetbrains.kotlin.types.expressions.FakeCallKind
+import org.jetbrains.kotlin.types.expressions.FakeCallResolver
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 /**
@@ -37,3 +47,40 @@ val CallableDescriptor.controllerTypeIfCoroutine: KotlinType?
 
         return this.extensionReceiverParameter?.returnType
     }
+
+fun FakeCallResolver.resolveCoroutineHandleResultCallIfNeeded(
+        callElement: KtExpression,
+        expressionToReturn: KtExpression?,
+        functionDescriptor: FunctionDescriptor,
+        context: ResolutionContext<*>
+) {
+    functionDescriptor.controllerTypeIfCoroutine ?: return
+
+    val info = if (expressionToReturn != null)
+        context.trace.bindingContext.get(BindingContext.EXPRESSION_TYPE_INFO, expressionToReturn)
+    else
+        null
+
+    val temporaryBindingTrace = TemporaryBindingTrace.create(context.trace, "trace to store fake argument for", "continuation")
+
+    val continuation =
+            ExpressionTypingUtils.createFakeExpressionOfType(
+                    callElement.project, temporaryBindingTrace, "continuation",
+                    // should be Continuation<Nothing>
+                    functionDescriptor.builtIns.nothingType)
+
+    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
+
+    val resolutionResults = resolveFakeCall(
+            context.replaceBindingTrace(temporaryBindingTrace), functionDescriptor.extensionReceiverParameter!!.value,
+            Name.identifier("handleResult"), callElement, callElement, FakeCallKind.OTHER,
+            listOf(firstArgument, continuation))
+
+    if (resolutionResults.isSuccess) {
+        context.trace.record(BindingContext.RETURN_HANDLE_RESULT_RESOLVED_CALL, callElement, resolutionResults.resultingCall)
+    }
+}

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -128,6 +128,7 @@ public interface BindingContext {
 
     WritableSlice<KtExpression, ResolvedCall<FunctionDescriptor>> LOOP_RANGE_HAS_NEXT_RESOLVED_CALL = Slices.createSimpleSlice();
     WritableSlice<KtExpression, ResolvedCall<FunctionDescriptor>> LOOP_RANGE_NEXT_RESOLVED_CALL = Slices.createSimpleSlice();
+    WritableSlice<KtExpression, ResolvedCall<FunctionDescriptor>> RETURN_HANDLE_RESULT_RESOLVED_CALL = Slices.createSimpleSlice();
 
     WritableSlice<VariableAccessorDescriptor, ResolvedCall<FunctionDescriptor>> DELEGATED_PROPERTY_RESOLVED_CALL = Slices.createSimpleSlice();
     WritableSlice<VariableAccessorDescriptor, Call> DELEGATED_PROPERTY_CALL = Slices.createSimpleSlice();

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallCompleter.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.resolve.calls
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.coroutines.controllerTypeIfCoroutine
+import org.jetbrains.kotlin.coroutines.resolveCoroutineHandleResultCallIfNeeded
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.*
@@ -46,6 +48,7 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.expressions.DataFlowAnalyzer
+import org.jetbrains.kotlin.types.expressions.FakeCallResolver
 import java.util.*
 
 class CallCompleter(
@@ -54,7 +57,8 @@ class CallCompleter(
         private val symbolUsageValidator: SymbolUsageValidator,
         private val dataFlowAnalyzer: DataFlowAnalyzer,
         private val callCheckers: Iterable<CallChecker>,
-        private val builtIns: KotlinBuiltIns
+        private val builtIns: KotlinBuiltIns,
+        private val fakeCallResolver: FakeCallResolver
 ) {
     fun <D : CallableDescriptor> completeCall(
             context: BasicCallResolutionContext,
@@ -86,6 +90,8 @@ class CallCompleter(
             else
                 resolvedCall.call.calleeExpression
             symbolUsageValidator.validateCall(resolvedCall, resolvedCall.resultingDescriptor, context.trace, element!!)
+
+            resolveHandleResultCallForCoroutineLambdaExpressions(context, resolvedCall)
         }
 
         if (results.isSingleResult && results.resultingCall.status.isSuccess) {
@@ -94,6 +100,27 @@ class CallCompleter(
         return results
     }
 
+    private fun <D : CallableDescriptor> resolveHandleResultCallForCoroutineLambdaExpressions(
+            context: BasicCallResolutionContext,
+            resolvedCall: ResolvedCall<D>
+    ) {
+        resolvedCall.valueArguments.values
+                .flatMap { it.arguments.map { it.getArgumentExpression() } }
+                .filterIsInstance<KtLambdaExpression>()
+                .forEach {
+                    val function = context.trace.bindingContext[BindingContext.FUNCTION, it.functionLiteral] ?: return@forEach
+
+                    function.controllerTypeIfCoroutine ?: return@forEach
+
+                    val lastBlockStatement = it.functionLiteral.bodyExpression?.statements?.lastOrNull() ?: return@forEach
+
+                    // Already resolved
+                    if (lastBlockStatement is KtReturnExpression) return@forEach
+
+                    fakeCallResolver.resolveCoroutineHandleResultCallIfNeeded(lastBlockStatement, lastBlockStatement, function, context)
+                }
+    }
+
     private fun <D : CallableDescriptor> completeAllCandidates(
             context: BasicCallResolutionContext,
             results: OverloadResolutionResultsImpl<D>

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallResolver.java

@@ -138,7 +138,7 @@ public class CallResolver {
 
     @NotNull
     public OverloadResolutionResults<FunctionDescriptor> resolveCallWithGivenName(
-            @NotNull ExpressionTypingContext context,
+            @NotNull ResolutionContext<?> context,
             @NotNull Call call,
             @NotNull KtReferenceExpression functionReference,
             @NotNull Name name

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ControlStructureTypingVisitor.java

@@ -23,6 +23,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
+import org.jetbrains.kotlin.coroutines.CoroutineUtilKt;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.BindingContext;
@@ -602,6 +603,9 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
                     context.trace.report(RETURN_NOT_ALLOWED.on(expression));
                     resultType = ErrorUtils.createErrorType(RETURN_NOT_ALLOWED_MESSAGE);
                 }
+
+                CoroutineUtilKt.resolveCoroutineHandleResultCallIfNeeded(
+                        components.fakeCallResolver, expression, expression.getReturnedExpression(), functionDescriptor, context);
             }
             else {
                 context.trace.report(NOT_A_RETURN_LABEL.on(expression, expression.getLabelName()));

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DestructuringDeclarationResolver.kt

@@ -72,7 +72,7 @@ class DestructuringDeclarationResolver(
         val expectedType = getExpectedTypeForComponent(context, entry)
         val results = fakeCallResolver.resolveFakeCall(
                 context.replaceExpectedType(expectedType), receiver, componentName,
-                entry, initializer, FakeCallKind.COMPONENT
+                entry, initializer, FakeCallKind.COMPONENT, emptyList()
         )
 
         if (!results.isSuccess) {

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingComponents.java

@@ -39,6 +39,7 @@ public class ExpressionTypingComponents {
     /*package*/ PlatformToKotlinClassMap platformToKotlinClassMap;
     /*package*/ ControlStructureTypingUtils controlStructureTypingUtils;
     /*package*/ ForLoopConventionsChecker forLoopConventionsChecker;
+    /*package*/ FakeCallResolver fakeCallResolver;
     /*package*/ ReflectionTypes reflectionTypes;
     /*package*/ SymbolUsageValidator symbolUsageValidator;
     /*package*/ DynamicTypesSettings dynamicTypesSettings;
@@ -93,6 +94,11 @@ public class ExpressionTypingComponents {
         this.forLoopConventionsChecker = forLoopConventionsChecker;
     }
 
+    @Inject
+    public void setFakeCallResolver(@NotNull FakeCallResolver fakeCallResolver) {
+        this.fakeCallResolver = fakeCallResolver;
+    }
+
     @Inject
     public void setReflectionTypes(@NotNull ReflectionTypes reflectionTypes) {
         this.reflectionTypes = reflectionTypes;

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FakeCallResolver.kt

@@ -28,8 +28,8 @@ import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
 import org.jetbrains.kotlin.resolve.calls.CallResolver
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResults
-import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResultsImpl
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.KotlinType
@@ -66,7 +66,7 @@ class FakeCallResolver(
 
     @JvmOverloads
     fun resolveFakeCall(
-            context: ExpressionTypingContext,
+            context: ResolutionContext<*>,
             receiver: ReceiverValue,
             name: Name,
             callElement: KtExpression,
@@ -79,7 +79,7 @@ class FakeCallResolver(
 
     fun makeAndResolveFakeCall(
             receiver: ReceiverValue?,
-            context: ExpressionTypingContext,
+            context: ResolutionContext<*>,
             valueArguments: List<KtExpression>,
             name: Name,
             callElement: KtExpression?,
@@ -132,7 +132,7 @@ class FakeCallResolver(
 
     @JvmOverloads fun makeAndResolveFakeCallInContext(
             receiver: ReceiverValue?,
-            context: ExpressionTypingContext,
+            context: ResolutionContext<*>,
             valueArguments: List<KtExpression>,
             name: Name,
             callElement: KtExpression?,