Support non-tail suspend calls in front-end

#KT-15597 In Progress

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

@@ -2229,9 +2229,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
     @NotNull
     private Type getReturnTypeForNonLocalReturn(DeclarationDescriptor elementDescriptor) {
-        return (elementDescriptor instanceof AnonymousFunctionDescriptor
+        return elementDescriptor instanceof FunctionDescriptor && ((FunctionDescriptor) elementDescriptor).isSuspend()
-            && ((AnonymousFunctionDescriptor) elementDescriptor).isCoroutine())
-            || (elementDescriptor instanceof FunctionDescriptor && ((FunctionDescriptor) elementDescriptor).isSuspend())
         ? getBoxedReturnTypeForSuspend((FunctionDescriptor) elementDescriptor)
         : typeMapper.mapReturnType((CallableDescriptor) elementDescriptor);
     }

compiler/frontend/src/org/jetbrains/kotlin/cfg/ControlFlowInformationProvider.kt

@@ -827,7 +827,7 @@ class ControlFlowInformationProvider private constructor(
             val isUsedAsExpression = instruction.owner.getUsages(instruction.outputValue).isNotEmpty()
 
             if (!isUsedAsExpression || !instruction.isTailCall(enclosingSuspendFunction) || isInsideTry(element)) {
-                trace.report(Errors.SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE.on(element))
+                trace.record(BindingContext.CONTAINS_NON_TAIL_SUSPEND_CALLS, currentFunction.original)
             }
         }
     }

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

@@ -21,6 +21,6 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 
-val CallableDescriptor.isSuspendLambda get() = this is AnonymousFunctionDescriptor && this.isCoroutine
+val CallableDescriptor.isSuspendLambda get() = this is AnonymousFunctionDescriptor && this.isSuspend
 
 val ValueParameterDescriptor.hasSuspendFunctionType get() = returnType?.isSuspendFunctionType == true

compiler/frontend/src/org/jetbrains/kotlin/descriptors/impl/AnonymousFunctionDescriptor.java

@@ -23,20 +23,21 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.name.Name;
 
 public class AnonymousFunctionDescriptor extends SimpleFunctionDescriptorImpl {
-    private final boolean isCoroutine;
+    private final boolean isSuspend;
 
     public AnonymousFunctionDescriptor(
             @NotNull DeclarationDescriptor containingDeclaration,
             @NotNull Annotations annotations,
             @NotNull Kind kind,
             @NotNull SourceElement source,
-            boolean isCoroutine
+            boolean isSuspend
     ) {
         super(containingDeclaration, null, annotations, Name.special("<anonymous>"), kind, source);
-        this.isCoroutine = isCoroutine;
+        this.isSuspend = isSuspend;
     }
 
-    public boolean isCoroutine() {
+    @Override
-        return isCoroutine;
+    public boolean isSuspend() {
+        return isSuspend;
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/Errors.java

@@ -891,7 +891,6 @@ public interface Errors {
     DiagnosticFactory0<PsiElement> NON_LOCAL_SUSPENSION_POINT = DiagnosticFactory0.create(ERROR);
     DiagnosticFactory0<PsiElement> ILLEGAL_SUSPEND_FUNCTION_CALL = DiagnosticFactory0.create(ERROR);
     DiagnosticFactory0<PsiElement> ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL = DiagnosticFactory0.create(ERROR);
-    DiagnosticFactory0<PsiElement> SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE = DiagnosticFactory0.create(ERROR);
 
 
     // Error sets

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/rendering/DefaultErrorMessages.java

@@ -891,7 +891,6 @@ public class DefaultErrorMessages {
         MAP.put(NON_LOCAL_SUSPENSION_POINT, "Suspension functions can be called only within coroutine body");
         MAP.put(ILLEGAL_SUSPEND_FUNCTION_CALL, "Suspend functions are only allowed to be called from a coroutine or another suspend function");
         MAP.put(ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL, "Restricted suspending functions can only invoke member or extension suspending functions on their restricted coroutine scope");
-        MAP.put(SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE, "Only tail calls are allowed to another suspension function, and these calls must be used as return values");
 
         MAP.setImmutable();
 

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

@@ -135,8 +135,8 @@ public interface BindingContext {
     WritableSlice<KtExpression, ResolvedCall<FunctionDescriptor>> LOOP_RANGE_NEXT_RESOLVED_CALL = Slices.createSimpleSlice();
     WritableSlice<KtExpression, ResolvedCall<FunctionDescriptor>> RETURN_HANDLE_RESULT_RESOLVED_CALL = Slices.createSimpleSlice();
 
-    WritableSlice<Call, CallableDescriptor> ENCLOSING_SUSPEND_LAMBDA_FOR_SUSPENSION_POINT = Slices.createSimpleSlice();
+    WritableSlice<Call, FunctionDescriptor> ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL = Slices.createSimpleSlice();
-    WritableSlice<Call, SimpleFunctionDescriptor> ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL = Slices.createSimpleSlice();
+    WritableSlice<FunctionDescriptor, Boolean> CONTAINS_NON_TAIL_SUSPEND_CALLS = Slices.createSimpleSetSlice();
 
     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/checkers/coroutineCallChecker.kt

@@ -20,10 +20,8 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.coroutines.hasSuspendFunctionType
-import org.jetbrains.kotlin.coroutines.isSuspendLambda
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.diagnostics.Errors
@@ -47,40 +45,26 @@ object CoroutineSuspendCallChecker : CallChecker {
         val descriptor = resolvedCall.candidateDescriptor as? SimpleFunctionDescriptor ?: return
         if (!descriptor.isSuspend) return
 
-        val (closestSuspendLambdaScope, closestSuspensionLambdaDescriptor) =
+        val enclosingSuspendFunction =
                 context.scope
                         .parentsWithSelf.firstOrNull {
                                 it is LexicalScope && it.kind == LexicalScopeKind.FUNCTION_INNER_SCOPE &&
-                                    it.ownerDescriptor.safeAs<CallableDescriptor>()?.isSuspendLambda == true
+                                    it.ownerDescriptor.safeAs<FunctionDescriptor>()?.isSuspend == true
-                        }?.let { it to it.cast<LexicalScope>().ownerDescriptor.cast<CallableDescriptor>() }
+                        }?.cast<LexicalScope>()?.ownerDescriptor?.cast<FunctionDescriptor>()
-                ?: null to null
-
-        val enclosingSuspendFunction =
-                context.scope.parentsWithSelf.filterIsInstance<LexicalScope>().takeWhile { it != closestSuspendLambdaScope }
-                        .firstOrNull {
-                            (it.ownerDescriptor as? FunctionDescriptor)?.isSuspend == true
-                        }?.ownerDescriptor as? SimpleFunctionDescriptor
 
         when {
             enclosingSuspendFunction != null -> {
-                // Tail calls checks happen during control flow analysis
-                // Here we only record enclosing function mapping (for backends purposes)
-                context.trace.record(BindingContext.ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL, resolvedCall.call, enclosingSuspendFunction)
-
-                checkRestrictsSuspension(enclosingSuspendFunction, resolvedCall, reportOn, context)
-            }
-            closestSuspensionLambdaDescriptor != null -> {
                 val callElement = resolvedCall.call.callElement as KtExpression
 
-                if (!InlineUtil.checkNonLocalReturnUsage(closestSuspensionLambdaDescriptor, callElement, context.resolutionContext)) {
+                if (!InlineUtil.checkNonLocalReturnUsage(enclosingSuspendFunction, callElement, context.resolutionContext)) {
                     context.trace.report(Errors.NON_LOCAL_SUSPENSION_POINT.on(reportOn))
                 }
 
                 context.trace.record(
-                        BindingContext.ENCLOSING_SUSPEND_LAMBDA_FOR_SUSPENSION_POINT, resolvedCall.call, closestSuspensionLambdaDescriptor
+                        BindingContext.ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL, resolvedCall.call, enclosingSuspendFunction
                 )
 
-                checkRestrictsSuspension(closestSuspensionLambdaDescriptor, resolvedCall, reportOn, context)
+                checkRestrictsSuspension(enclosingSuspendFunction, resolvedCall, reportOn, context)
             }
             else -> {
                 context.trace.report(Errors.ILLEGAL_SUSPEND_FUNCTION_CALL.on(reportOn))

compiler/testData/diagnostics/tests/coroutines/tailCalls/forbidden.kt

@@ -5,15 +5,15 @@ import kotlin.coroutines.intrinsics.*
 fun nonSuspend() {}
 
 suspend fun foo() {
-    <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { x: Continuation<Int> -> }<!>
+    suspendCoroutineOrReturn { x: Continuation<Int> -> }
 
     nonSuspend()
 }
 
 suspend fun unitSuspend() {
-    <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { x: Continuation<Int> -> }<!>
+    suspendCoroutineOrReturn { x: Continuation<Int> -> }
 }
 
 suspend fun baz(): Int = run {
-    <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { x: Continuation<Int> -> }<!>
+    suspendCoroutineOrReturn { x: Continuation<Int> -> }
 }

compiler/testData/diagnostics/tests/coroutines/tailCalls/localFunctions.kt

@@ -3,7 +3,7 @@ suspend fun unit() {}
 
 suspend fun foo() {
     suspend fun bar() {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE, SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>baz()<!>
+        baz()
         return unit()
     }
 

compiler/testData/diagnostics/tests/coroutines/tailCalls/nothingTypedSuspendFunction.kt

@@ -2,7 +2,7 @@
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 
-suspend fun suspendLogAndThrow(exception: Throwable): Nothing = <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { c ->
+suspend fun suspendLogAndThrow(exception: Throwable): Nothing = suspendCoroutineOrReturn { c ->
     c.resumeWithException(exception)
     SUSPENDED_MARKER
-}<!>
+}

compiler/testData/diagnostics/tests/coroutines/tailCalls/recursive.kt

@@ -1,5 +1,5 @@
 suspend fun unit1() {
-    <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>unit1()<!>
+    unit1()
 }
 
 suspend fun unit2() {
@@ -13,6 +13,6 @@ suspend fun int1(): Int {
 suspend fun int2(): Int = int2()
 
 suspend fun int3(): Int {
-    <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>int3()<!>
+    int3()
     return int3()
 }

compiler/testData/diagnostics/tests/coroutines/tailCalls/tryCatch.kt

@@ -8,15 +8,15 @@ suspend fun baz(): Int = 1
 
 suspend fun tryCatch(): Int {
     return try {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { x: Continuation<Int> -> }<!>
+        suspendCoroutineOrReturn { x: Continuation<Int> -> }
     } catch (e: Exception) {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>baz()<!> // another suspend function
+        baz() // another suspend function
     }
 }
 
 suspend fun tryFinally(): Int {
     return try {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { x: Continuation<Int> -> }<!>
+        suspendCoroutineOrReturn { x: Continuation<Int> -> }
     } finally {
         nonSuspend()
     }
@@ -26,16 +26,16 @@ suspend fun returnInFinally(): Int {
     try {
     } finally {
         // Probably this is too restrictive, but it does not matter much
-        return <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE, SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>baz()<!>
+        return baz()
     }
 }
 
 suspend fun tryCatchFinally(): Int {
     return try {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>suspendCoroutineOrReturn { x: Continuation<Int> -> }<!>
+        suspendCoroutineOrReturn { x: Continuation<Int> -> }
     } catch (e: Exception) {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>baz()<!> // another suspend function
+        baz() // another suspend function
     } finally {
-        <!SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE, SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE!>baz()<!>
+        baz()
     }
 }

js/js.frontend/src/org/jetbrains/kotlin/js/descriptorUtils/descriptorUtils.kt

@@ -54,4 +54,4 @@ fun KotlinType.getJetTypeFqName(printTypeArguments: Boolean): String {
 fun ClassDescriptor.hasPrimaryConstructor(): Boolean = unsubstitutedPrimaryConstructor != null
 
 val DeclarationDescriptor.isCoroutineLambda: Boolean
-    get() = this is AnonymousFunctionDescriptor && isCoroutine
+    get() = this is AnonymousFunctionDescriptor && isSuspend

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/TranslationUtils.java

@@ -339,7 +339,7 @@ public final class TranslationUtils {
     @NotNull
     public static JsExpression translateContinuationArgument(@NotNull TranslationContext context, @NotNull ResolvedCall<?> resolvedCall) {
         CallableDescriptor continuationDescriptor =
-                context.bindingContext().get(BindingContext.ENCLOSING_SUSPEND_LAMBDA_FOR_SUSPENSION_POINT, resolvedCall.getCall());
+                context.bindingContext().get(BindingContext.ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL, resolvedCall.getCall());
 
         if (continuationDescriptor == null) {
             continuationDescriptor = getEnclosingContinuationParameter(context);