Support new suspend convention in frontend

#KT-14924 In Progress

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

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.ExplicitSmartCasts;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.ImplicitSmartCasts;
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant;
+import org.jetbrains.kotlin.resolve.coroutine.CoroutineReceiverValue;
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
 import org.jetbrains.kotlin.resolve.scopes.receivers.Qualifier;
@@ -135,6 +136,9 @@ 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, CoroutineReceiverValue> COROUTINE_RECEIVER_FOR_SUSPENSION_POINT = Slices.createSimpleSlice();
+    WritableSlice<Call, SimpleFunctionDescriptor> ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_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/TargetPlatform.kt

@@ -64,7 +64,6 @@ private val DEFAULT_DECLARATION_CHECKERS = listOf(
         UnderscoreChecker,
         InlineParameterChecker,
         InfixModifierChecker(),
-        SuspendModifierChecker,
         CoroutineModifierChecker,
         SinceKotlinAnnotationValueChecker,
         ReifiedTypeParameterAnnotationChecker()

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/coroutineCallChecker.kt

@@ -19,22 +19,47 @@ package org.jetbrains.kotlin.resolve.calls.checkers
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.coroutine.CoroutineReceiverValue
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.resolve.scopes.utils.getImplicitReceiversHierarchy
+import org.jetbrains.kotlin.resolve.scopes.utils.parentsWithSelf
 
 object CoroutineSuspendCallChecker : CallChecker {
     override fun check(resolvedCall: ResolvedCall<*>, reportOn: PsiElement, context: CallCheckerContext) {
-        val descriptor = resolvedCall.candidateDescriptor as? FunctionDescriptor ?: return
-        if (!descriptor.isSuspend || descriptor.initialSignatureDescriptor == null) return
+        val descriptor = resolvedCall.candidateDescriptor as? SimpleFunctionDescriptor ?: return
+        if (!descriptor.isSuspend) return
 
-        val dispatchReceiverOwner = (resolvedCall.dispatchReceiver as? CoroutineReceiverValue)?.declarationDescriptor ?: return
-        val callElement = resolvedCall.call.callElement as KtExpression
+        val closestCoroutineReceiver =
+                context.scope.getImplicitReceiversHierarchy().firstOrNull { it.value is CoroutineReceiverValue }?.value as CoroutineReceiverValue?
 
-        if (!InlineUtil.checkNonLocalReturnUsage(dispatchReceiverOwner, callElement, context.resolutionContext)) {
-            context.trace.report(Errors.NON_LOCAL_SUSPENSION_POINT.on(reportOn))
+        val enclosingSuspendFunction =
+                context.scope.parentsWithSelf.filterIsInstance<LexicalScope>().takeWhile {
+                    closestCoroutineReceiver == null || it.implicitReceiver?.value != closestCoroutineReceiver
+                }.firstOrNull {
+                    (it.ownerDescriptor as? FunctionDescriptor)?.isSuspend == true
+                }?.ownerDescriptor as? SimpleFunctionDescriptor
+
+        when {
+            enclosingSuspendFunction != null -> {
+                // TODO: check if tail call
+                context.trace.record(BindingContext.ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL, resolvedCall.call, enclosingSuspendFunction)
+            }
+            closestCoroutineReceiver != null -> {
+
+                val callElement = resolvedCall.call.callElement as KtExpression
+
+                if (!InlineUtil.checkNonLocalReturnUsage(closestCoroutineReceiver.declarationDescriptor, callElement, context.resolutionContext)) {
+                    context.trace.report(Errors.NON_LOCAL_SUSPENSION_POINT.on(reportOn))
+                }
+
+                context.trace.record(BindingContext.COROUTINE_RECEIVER_FOR_SUSPENSION_POINT, resolvedCall.call, closestCoroutineReceiver)
+            }
         }
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/checkers/SuspendModifierChecker.kt

@@ -1,85 +0,0 @@
-/*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.jetbrains.kotlin.resolve.checkers
-
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.coroutines.isValidContinuation
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.diagnostics.DiagnosticSink
-import org.jetbrains.kotlin.diagnostics.Errors
-import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.types.typeUtil.isUnit
-import org.jetbrains.kotlin.utils.sure
-
-object SuspendModifierChecker : SimpleDeclarationChecker {
-    private val ALLOW_SUSPEND_EXTENSIONS_ANNOTATION_FQ_NAME =
-            KotlinBuiltIns.COROUTINES_PACKAGE_FQ_NAME.child(Name.identifier("AllowSuspendExtensions"))
-
-    override fun check(
-            declaration: KtDeclaration,
-            descriptor: DeclarationDescriptor,
-            diagnosticHolder: DiagnosticSink,
-            bindingContext: BindingContext
-    ) {
-        val functionDescriptor = descriptor as? FunctionDescriptor ?: return
-        if (!functionDescriptor.isSuspend) return
-
-        val suspendModifierElement = declaration.modifierList?.getModifier(KtTokens.SUSPEND_KEYWORD).sure { declaration.text }
-        fun report(message: String) {
-            diagnosticHolder.report(Errors.INAPPLICABLE_MODIFIER.on(suspendModifierElement, KtTokens.SUSPEND_KEYWORD, message))
-        }
-
-        if (functionDescriptor.dispatchReceiverParameter == null) {
-            if (functionDescriptor.extensionReceiverParameter == null) {
-                report("function must be either a class member or an extension")
-                return
-            }
-
-            val classDescriptor =
-                    functionDescriptor.extensionReceiverParameter!!.type.constructor.declarationDescriptor as? ClassDescriptor
-            if (classDescriptor == null) {
-                report("function must be an extension to class")
-                return
-            }
-
-            if (!classDescriptor.annotations.hasAnnotation(ALLOW_SUSPEND_EXTENSIONS_ANNOTATION_FQ_NAME)) {
-                report("controller class must be annotated with AllowSuspendExtensions annotation")
-                return
-            }
-        }
-
-        val continuationParameterType = functionDescriptor.valueParameters.lastOrNull()?.type
-        val isValidContinuation = continuationParameterType?.isValidContinuation() ?: false
-        if (!isValidContinuation) {
-            report("last parameter of suspend function should have a type of Continuation<T>")
-            return
-        }
-
-        if (continuationParameterType?.arguments?.firstOrNull()?.isStarProjection == true) {
-            report("Continuation<*> is prohibited as a last parameter of suspend function")
-        }
-
-        if (functionDescriptor.returnType?.isUnit() != true) {
-            report("return type of suspension function must be a kotlin.Unit, but ${functionDescriptor.returnType} was found")
-        }
-    }
-}