Allow suspend extensions for specially annotated controllers

compiler/frontend/src/org/jetbrains/kotlin/resolve/ModifiersChecker.kt

@@ -83,7 +83,7 @@ object ModifierCheckerCore {
             NOINLINE_KEYWORD  to EnumSet.of(VALUE_PARAMETER),
             COROUTINE_KEYWORD to EnumSet.of(VALUE_PARAMETER),
             TAILREC_KEYWORD   to EnumSet.of(FUNCTION),
-            SUSPEND_KEYWORD   to EnumSet.of(MEMBER_FUNCTION),
+            SUSPEND_KEYWORD   to EnumSet.of(FUNCTION),
             EXTERNAL_KEYWORD  to EnumSet.of(FUNCTION, PROPERTY_GETTER, PROPERTY_SETTER),
             ANNOTATION_KEYWORD to EnumSet.of(ANNOTATION_CLASS),
             CROSSINLINE_KEYWORD to EnumSet.of(VALUE_PARAMETER),

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

@@ -20,12 +20,14 @@ import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
 import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
 import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 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.incremental.KotlinLookupLocation
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
 import org.jetbrains.kotlin.types.typeUtil.isUnit
@@ -33,6 +35,8 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.sure
 
 object SuspendModifierChecker : SimpleDeclarationChecker {
+    private val ALLOW_SUSPEND_EXTENSIONS_ANNOTATION_FQ_NAME = FqName("kotlin.coroutines.AllowSuspendExtensions")
+
     override fun check(
             declaration: KtDeclaration,
             descriptor: DeclarationDescriptor,
@@ -47,6 +51,25 @@ object SuspendModifierChecker : SimpleDeclarationChecker {
             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 isValidContinuation = functionDescriptor.valueParameters.lastOrNull()?.type?.isValidContinuation() ?: false
         if (!isValidContinuation) {
             report("last parameter of suspend function should have a type of Continuation<T>")

compiler/testData/diagnostics/tests/coroutines/suspendApplicability.kt

@@ -1,8 +1,15 @@
 // !DIAGNOSTICS: -UNUSED_PARAMETER -NOTHING_TO_INLINE
-<!WRONG_MODIFIER_TARGET!>suspend<!> fun notMember(x: Continuation<Int>) {
+<!INAPPLICABLE_MODIFIER!>suspend<!> fun notMember(x: Continuation<Int>) {
 
 }
 
+<!INAPPLICABLE_MODIFIER!>suspend<!> fun String.wrongExtension(x: Continuation<Int>) {
+}
+
+suspend fun Controller.correctExtension(x: Continuation<Int>) {
+}
+
+@AllowSuspendExtensions
 class Controller {
     suspend fun valid(x: Continuation<Int>) {
 
@@ -23,4 +30,8 @@ class Controller {
     <!INAPPLICABLE_MODIFIER!>suspend<!> fun starProjection(vararg x: Continuation<Any>) {
 
     }
+
+    suspend fun String.memberExtension(x: Continuation<Int>) {
+
+    }
 }

compiler/testData/diagnostics/tests/coroutines/suspendApplicability.txt

@@ -1,8 +1,10 @@
 package
 
 public suspend fun notMember(/*0*/ x: kotlin.coroutines.Continuation<kotlin.Int>): kotlin.Unit
+public suspend fun Controller.correctExtension(/*0*/ x: kotlin.coroutines.Continuation<kotlin.Int>): kotlin.Unit
+public suspend fun kotlin.String.wrongExtension(/*0*/ x: kotlin.coroutines.Continuation<kotlin.Int>): kotlin.Unit
 
-public final class Controller {
+@kotlin.coroutines.AllowSuspendExtensions() public final class Controller {
     public constructor Controller()
     public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
@@ -12,4 +14,5 @@ public final class Controller {
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
     public final suspend fun valid(/*0*/ x: kotlin.coroutines.Continuation<kotlin.Int>): kotlin.Unit
     public final suspend fun wrongParam(/*0*/ x: kotlin.collections.Collection<kotlin.Int>): kotlin.Unit
+    public final suspend fun kotlin.String.memberExtension(/*0*/ x: kotlin.coroutines.Continuation<kotlin.Int>): kotlin.Unit
 }