Support non-tail suspend calls in front-end

#KT-15597 In Progress

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() {
-        return isCoroutine;
+    @Override
+    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, SimpleFunctionDescriptor> ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL = Slices.createSimpleSlice();
+    WritableSlice<Call, FunctionDescriptor> 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
-                        }?.let { it to it.cast<LexicalScope>().ownerDescriptor.cast<CallableDescriptor>() }
-                ?: null to null
-
-        val enclosingSuspendFunction =
-                context.scope.parentsWithSelf.filterIsInstance<LexicalScope>().takeWhile { it != closestSuspendLambdaScope }
-                        .firstOrNull {
-                            (it.ownerDescriptor as? FunctionDescriptor)?.isSuspend == true
-                        }?.ownerDescriptor as? SimpleFunctionDescriptor
+                                    it.ownerDescriptor.safeAs<FunctionDescriptor>()?.isSuspend == true
+                        }?.cast<LexicalScope>()?.ownerDescriptor?.cast<FunctionDescriptor>()
 
         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))