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
 }

core/builtins/src/kotlin/Coroutines.kt

@@ -31,3 +31,11 @@ interface Continuation<in P> {
      */
     fun resumeWithException(exception: Throwable)
 }
+
+
+/**
+ * Specifies that suspend extensions with a receiver based on corresponding controller class are allowed to be declared
+ */
+@Target(AnnotationTarget.CLASS)
+@Retention(AnnotationRetention.BINARY)
+annotation class AllowSuspendExtensions

idea/idea-completion/testData/keywords/AfterClasses.kt

@@ -34,4 +34,5 @@ class B {
 // EXIST:  external
 // EXIST:  annotation class
 // EXIST:  const
+// EXIST:  suspend
 // NOTHING_ELSE

idea/idea-completion/testData/keywords/GlobalPropertyAccessors.kt

@@ -38,4 +38,5 @@ var a : Int
 // EXIST:  external
 // EXIST:  annotation class
 // EXIST:  const
+// EXIST:  suspend
 // NOTHING_ELSE

idea/idea-completion/testData/keywords/InTopScopeAfterPackage.kt

@@ -25,4 +25,5 @@ package Test
 // EXIST:  external
 // EXIST:  annotation class
 // EXIST:  const
+// EXIST:  suspend
 // NOTHING_ELSE

idea/idea-completion/testData/keywords/TopScope.kt

@@ -24,4 +24,5 @@
 // EXIST:  external
 // EXIST:  annotation class
 // EXIST:  const
+// EXIST:  suspend
 // NOTHING_ELSE

idea/testData/checker/infos/suspendApplicability.kt

@@ -1,8 +1,15 @@
 // !DIAGNOSTICS: -UNUSED_PARAMETER -NOTHING_TO_INLINE
-<error descr="[WRONG_MODIFIER_TARGET] Modifier 'suspend' is not applicable to 'top level function'"><info descr="null">suspend</info></error> fun notMember(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
+<error descr="[INAPPLICABLE_MODIFIER] 'suspend' modifier is inapplicable. The reason is that function must be either a class member or an extension"><info descr="null">suspend</info></error> fun notMember(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
 
 }
 
+<error descr="[INAPPLICABLE_MODIFIER] 'suspend' modifier is inapplicable. The reason is that controller class must be annotated with AllowSuspendExtensions annotation"><info descr="null">suspend</info></error> fun String.wrongExtension(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
+}
+
+<info descr="null">suspend</info> fun Controller.correctExtension(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
+}
+
+@AllowSuspendExtensions
 class Controller {
     <info descr="null">suspend</info> fun valid(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
 
@@ -20,7 +27,11 @@ class Controller {
 
     }
 
-    <info descr="null">suspend</info> fun starProjection(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<*>) {
+    <error descr="[INAPPLICABLE_MODIFIER] 'suspend' modifier is inapplicable. The reason is that last parameter of suspend function should have a type of Continuation<T>"><info descr="null">suspend</info></error> fun starProjection(<info descr="null">vararg</info> <warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Any>) {
+
+    }
+
+    <info descr="null">suspend</info> fun String.memberExtension(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
 
     }
 }

libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-runtime.txt

@@ -164,6 +164,9 @@ public abstract class kotlin/collections/ShortIterator : java/util/Iterator, kot
 	public fun remove ()V
 }
 
+public abstract interface annotation class kotlin/coroutines/AllowSuspendExtensions : java/lang/annotation/Annotation {
+}
+
 public abstract interface class kotlin/coroutines/Continuation {
 	public abstract fun resume (Ljava/lang/Object;)V
 	public abstract fun resumeWithException (Ljava/lang/Throwable;)V