Add additional checks for RestrictsSuspension-marked functions

Namely, check that when one calls a restricted function
the reciever used for that calls is obtained exactly from the enclosing
suspend function

 #KT-24859 Fixed

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

@@ -8,13 +8,10 @@ package org.jetbrains.kotlin.resolve.calls.checkers
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.config.restrictsSuspensionFqName
 import org.jetbrains.kotlin.config.isBuiltInCoroutineContext
+import org.jetbrains.kotlin.config.restrictsSuspensionFqName
 import org.jetbrains.kotlin.coroutines.hasSuspendFunctionType
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.descriptors.PropertyDescriptor
-import org.jetbrains.kotlin.descriptors.PropertyGetterDescriptor
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.name.Name
@@ -23,7 +20,6 @@ import org.jetbrains.kotlin.psi.KtThisExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
-import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.resolve.scopes.HierarchicalScope
@@ -131,36 +127,59 @@ fun checkCoroutinesFeature(languageVersionSettings: LanguageVersionSettings, dia
 }
 
 private fun checkRestrictsSuspension(
-    enclosingCallableDescriptor: CallableDescriptor,
+    enclosingSuspendCallableDescriptor: CallableDescriptor,
     resolvedCall: ResolvedCall<*>,
     reportOn: PsiElement,
     context: CallCheckerContext
 ) {
-    val enclosingSuspendReceiverValue = enclosingCallableDescriptor.extensionReceiverParameter?.value ?: return
+    fun ReceiverValue.isRestrictsSuspensionReceiver() = (listOf(type) + type.supertypes()).any {
-
-    fun ReceiverValue.isRestrictsSuspensionReceiver() = (type.supertypes() + type).any {
         it.constructor.declarationDescriptor?.annotations?.hasAnnotation(context.languageVersionSettings.restrictsSuspensionFqName()) == true
     }
 
     infix fun ReceiverValue.sameInstance(other: ReceiverValue?): Boolean {
         if (other == null) return false
-        if (this === other) return true
+        // Implicit receiver should be reference equal
+        if (this.original === other.original) return true
 
         val referenceExpression = ((other as? ExpressionReceiver)?.expression as? KtThisExpression)?.instanceReference
         val referenceTarget = referenceExpression?.let {
             context.trace.get(BindingContext.REFERENCE_TARGET, referenceExpression)
         }
 
-        return this === (referenceTarget as? CallableDescriptor)?.extensionReceiverParameter?.value
+        val referenceReceiverValue = when (referenceTarget) {
+            is CallableDescriptor -> referenceTarget.extensionReceiverParameter?.value
+            is ClassDescriptor -> referenceTarget.thisAsReceiverParameter.value
+            else -> null
+        }
+
+        return this === referenceReceiverValue
     }
 
-    if (!enclosingSuspendReceiverValue.isRestrictsSuspensionReceiver()) return
+    fun reportError() {
+        context.trace.report(Errors.ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL.on(reportOn))
+    }
+
+    val enclosingSuspendExtensionReceiverValue = enclosingSuspendCallableDescriptor.extensionReceiverParameter?.value
+    val enclosingSuspendDispatchReceiverValue = enclosingSuspendCallableDescriptor.dispatchReceiverParameter?.value
+
+    val receivers = listOfNotNull(resolvedCall.dispatchReceiver, resolvedCall.extensionReceiver)
+    for (receiverValue in receivers) {
+        if (!receiverValue.isRestrictsSuspensionReceiver()) continue
+        if (enclosingSuspendExtensionReceiverValue?.sameInstance(receiverValue) == true) continue
+        if (enclosingSuspendDispatchReceiverValue?.sameInstance(receiverValue) == true) continue
+
+        reportError()
+        return
+    }
+
+    if (enclosingSuspendExtensionReceiverValue?.isRestrictsSuspensionReceiver() != true) return
 
     // member of suspend receiver
-    if (enclosingSuspendReceiverValue sameInstance resolvedCall.dispatchReceiver?.original) return
+    if (enclosingSuspendExtensionReceiverValue sameInstance resolvedCall.dispatchReceiver) return
 
-    if (enclosingSuspendReceiverValue sameInstance resolvedCall.extensionReceiver?.original &&
+    if (enclosingSuspendExtensionReceiverValue sameInstance resolvedCall.extensionReceiver &&
-        resolvedCall.candidateDescriptor.extensionReceiverParameter!!.value.isRestrictsSuspensionReceiver()) return
+        resolvedCall.candidateDescriptor.extensionReceiverParameter!!.value.isRestrictsSuspensionReceiver()
+    ) return
 
-    context.trace.report(Errors.ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL.on(reportOn))
+    reportError()
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/memberExtension.kt

@@ -0,0 +1,62 @@
+// COMMON_COROUTINES_TEST
+// SKIP_TXT
+@COROUTINES_PACKAGE.RestrictsSuspension
+class RestrictedController {
+    suspend fun member() {
+        ext()
+        member()
+        memberExt()
+    }
+
+    suspend fun RestrictedController.memberExt() {
+        ext()
+        member()
+        memberExt()
+    }
+}
+
+suspend fun RestrictedController.ext() {
+    ext()
+    member()
+    memberExt()
+}
+
+fun generate(<!UNUSED_PARAMETER!>c<!>: suspend RestrictedController.() -> Unit) {}
+
+fun runBlocking(<!UNUSED_PARAMETER!>x<!>: suspend () -> Unit) {}
+
+fun test() {
+    generate a@{
+        ext()
+        member()
+        memberExt()
+
+        this@a.ext()
+        this@a.member()
+        this@a.memberExt()
+
+        generate b@{
+            ext()
+            member()
+            memberExt()
+
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>ext<!>()
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>member<!>()
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>memberExt<!>()
+
+            this@b.ext()
+            this@b.member()
+            this@b.memberExt()
+        }
+
+        runBlocking {
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>ext<!>()
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>member<!>()
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>memberExt<!>()
+
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>ext<!>()
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>member<!>()
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>memberExt<!>()
+        }
+    }
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield.kt

@@ -0,0 +1,91 @@
+// COMMON_COROUTINES_TEST
+// SKIP_TXT
+@COROUTINES_PACKAGE.RestrictsSuspension
+class RestrictedController<T> {
+    suspend fun yield(<!UNUSED_PARAMETER!>x<!>: T) {}
+
+    suspend fun anotherYield(x: T) {
+        yield(x)
+        this.yield(x)
+
+        yield2(x)
+        this.yield2(x)
+
+        with(this) {
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(x)
+            this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(x)
+
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(x)
+            this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(x)
+        }
+    }
+}
+
+fun <T> buildSequence(<!UNUSED_PARAMETER!>c<!>: suspend RestrictedController<T>.() -> Unit) {}
+suspend fun <T> RestrictedController<T>.yield2(<!UNUSED_PARAMETER!>x<!>: T) {}
+
+fun test() {
+    buildSequence<Int> a@{
+        buildSequence<Int> b@{
+            yield(1)
+            yield2(1)
+            this@b.yield(1)
+            this@b.yield2(1)
+
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(2) // Should be error
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(2) // Should be error
+
+            with(this) {
+                <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(3)
+                this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(3)
+
+                <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(3)
+                this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(3)
+            }
+        }
+    }
+
+    buildSequence<Int> {
+        buildSequence<String> {
+            yield("a")
+            yield2("a")
+            this.yield("b")
+            this.yield2("b")
+
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(1) // Should be error
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(1) // Should be error
+
+            with(this) {
+                <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>("")
+                this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>("")
+
+                <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>("")
+                this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>("")
+            }
+        }
+    }
+
+    buildSequence<Int> a@{
+        yield(1)
+        yield2(1)
+        buildSequence {
+            yield("")
+            yield2("")
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>(1)
+            this@a.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>(1)
+
+            with(this) {
+                <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>("")
+                this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>("")
+
+                <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>("")
+                this@with.<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield2<!>("")
+            }
+        }
+    }
+
+    buildSequence<String> {
+        yield("")
+        RestrictedController<String>().<!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>("1")
+    }
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/wrongEnclosingFunction.kt

@@ -0,0 +1,18 @@
+// COMMON_COROUTINES_TEST
+// SKIP_TXT
+@COROUTINES_PACKAGE.RestrictsSuspension
+class RestrictedController {
+    suspend fun yield() {}
+}
+
+fun generate(<!UNUSED_PARAMETER!>c<!>: suspend RestrictedController.() -> Unit) {}
+
+fun runBlocking(<!UNUSED_PARAMETER!>x<!>: suspend () -> Unit) {}
+
+fun test() {
+    generate {
+        runBlocking {
+            <!ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL!>yield<!>()
+        }
+    }
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestWithStdLibGenerated.java

@@ -1835,6 +1835,18 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
             }
 
+            @TestMetadata("memberExtension.kt")
+            public void testMemberExtension_1_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/memberExtension.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("memberExtension.kt")
+            public void testMemberExtension_1_3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/memberExtension.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
+            }
+
             @TestMetadata("notRelatedFun.kt")
             public void testNotRelatedFun_1_2() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/notRelatedFun.kt");
@@ -1847,6 +1859,18 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
             }
 
+            @TestMetadata("outerYield.kt")
+            public void testOuterYield_1_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("outerYield.kt")
+            public void testOuterYield_1_3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
+            }
+
             @TestMetadata("sameInstance.kt")
             public void testSameInstance_1_2() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/sameInstance.kt");
@@ -1870,6 +1894,18 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/simpleForbidden.kt");
                 doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
             }
+
+            @TestMetadata("wrongEnclosingFunction.kt")
+            public void testWrongEnclosingFunction_1_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/wrongEnclosingFunction.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("wrongEnclosingFunction.kt")
+            public void testWrongEnclosingFunction_1_3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/wrongEnclosingFunction.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
+            }
         }
 
         @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/suspendFunctionType")

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsTestWithStdLibUsingJavacGenerated.java

@@ -1835,6 +1835,18 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
             }
 
+            @TestMetadata("memberExtension.kt")
+            public void testMemberExtension_1_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/memberExtension.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("memberExtension.kt")
+            public void testMemberExtension_1_3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/memberExtension.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
+            }
+
             @TestMetadata("notRelatedFun.kt")
             public void testNotRelatedFun_1_2() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/notRelatedFun.kt");
@@ -1847,6 +1859,18 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
             }
 
+            @TestMetadata("outerYield.kt")
+            public void testOuterYield_1_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("outerYield.kt")
+            public void testOuterYield_1_3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
+            }
+
             @TestMetadata("sameInstance.kt")
             public void testSameInstance_1_2() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/sameInstance.kt");
@@ -1870,6 +1894,18 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/simpleForbidden.kt");
                 doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
             }
+
+            @TestMetadata("wrongEnclosingFunction.kt")
+            public void testWrongEnclosingFunction_1_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/wrongEnclosingFunction.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("wrongEnclosingFunction.kt")
+            public void testWrongEnclosingFunction_1_3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/wrongEnclosingFunction.kt");
+                doTestWithCoroutinesPackageReplacement(fileName, "kotlin.coroutines");
+            }
         }
 
         @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/coroutines/suspendFunctionType")