FIR: introduce ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL diagnostics

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -1191,6 +1191,7 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
             parameter<Symbol>("suspendCallable")
         }
         val NON_LOCAL_SUSPENSION_POINT by error<PsiElement>(PositioningStrategy.REFERENCED_NAME_BY_QUALIFIED)
+        val ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL by error<PsiElement>(PositioningStrategy.REFERENCED_NAME_BY_QUALIFIED)
     }
 }
 

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -611,5 +611,6 @@ object FirErrors {
     val ILLEGAL_SUSPEND_FUNCTION_CALL by error1<PsiElement, FirBasedSymbol<*>>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
     val ILLEGAL_SUSPEND_PROPERTY_ACCESS by error1<PsiElement, FirBasedSymbol<*>>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
     val NON_LOCAL_SUSPENSION_POINT by error0<PsiElement>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
+    val ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL by error0<PsiElement>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
 
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirSuspendCallChecker.kt

@@ -5,20 +5,32 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.expression
 
+import org.jetbrains.kotlin.builtins.StandardNames
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.expressions.FirThisReceiverExpression
+import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.FirResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeParameterType
 import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.checkers.COROUTINE_CONTEXT_1_2_20_FQ_NAME
 import org.jetbrains.kotlin.resolve.calls.checkers.COROUTINE_CONTEXT_1_2_30_FQ_NAME
 import org.jetbrains.kotlin.resolve.calls.checkers.COROUTINE_CONTEXT_1_3_FQ_NAME
@@ -29,13 +41,16 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
         COROUTINE_CONTEXT_1_2_20_FQ_NAME, COROUTINE_CONTEXT_1_2_30_FQ_NAME, COROUTINE_CONTEXT_1_3_FQ_NAME
     )
 
+    private val RESTRICTS_SUSPENSION_CLASS_ID =
+        ClassId(StandardNames.COROUTINES_PACKAGE_FQ_NAME_RELEASE, Name.identifier("RestrictsSuspension"))
+
     @OptIn(SymbolInternals::class)
     override fun check(expression: FirQualifiedAccessExpression, context: CheckerContext, reporter: DiagnosticReporter) {
         val reference = expression.calleeReference as? FirResolvedNamedReference ?: return
         if (reference is FirResolvedCallableReference) return
         val symbol = reference.resolvedSymbol as? FirCallableSymbol ?: return
         symbol.ensureResolved(FirResolvePhase.STATUS)
-        val fir = symbol.fir as? FirMemberDeclaration ?: return
+        val fir = symbol.fir as? FirCallableDeclaration ?: return
         when (fir) {
             is FirSimpleFunction -> if (!fir.isSuspend) return
             is FirProperty -> if (symbol.callableId.asSingleFqName() !in SUSPEND_PROPERTIES_FQ_NAMES) return
@@ -53,6 +68,9 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
             if (!checkNonLocalReturnUsage(enclosingSuspendFunction, context)) {
                 reporter.reportOn(expression.source, FirErrors.NON_LOCAL_SUSPENSION_POINT, context)
             }
+            if (!checkRestrictsSuspension(expression, enclosingSuspendFunction, fir, context)) {
+                reporter.reportOn(expression.source, FirErrors.ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL, context)
+            }
         }
     }
 
@@ -74,4 +92,61 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
             enclosingSuspendFunction === containingFunction
         }
     }
+
+    private fun checkRestrictsSuspension(
+        expression: FirQualifiedAccessExpression,
+        enclosingSuspendFunction: FirFunction,
+        calledDeclaration: FirCallableDeclaration,
+        context: CheckerContext
+    ): Boolean {
+        val session = context.session
+        val enclosingSuspendFunctionDispatchReceiverOwner =
+            (enclosingSuspendFunction.dispatchReceiverType as? ConeClassLikeType)?.lookupTag?.toFirRegularClass(session)
+        val enclosingSuspendFunctionExtensionReceiverOwner = enclosingSuspendFunction.takeIf { it.receiverTypeRef != null }
+        val dispatchReceiverExpression = expression.dispatchReceiver.takeIf { it !is FirNoReceiverExpression }
+        val extensionReceiverExpression = expression.extensionReceiver.takeIf { it !is FirNoReceiverExpression }
+        for (receiverExpression in listOfNotNull(dispatchReceiverExpression, extensionReceiverExpression)) {
+            if (!receiverExpression.typeRef.coneType.isRestrictSuspensionReceiver(session)) continue
+            if (sameInstanceOfReceiver(receiverExpression, enclosingSuspendFunctionDispatchReceiverOwner)) continue
+            if (sameInstanceOfReceiver(receiverExpression, enclosingSuspendFunctionExtensionReceiverOwner)) continue
+
+            return false
+        }
+        if (enclosingSuspendFunctionExtensionReceiverOwner?.receiverTypeRef?.coneType?.isRestrictSuspensionReceiver(session) != true) {
+            return true
+        }
+        if (sameInstanceOfReceiver(dispatchReceiverExpression, enclosingSuspendFunctionExtensionReceiverOwner)) {
+            return true
+        }
+        if (sameInstanceOfReceiver(extensionReceiverExpression, enclosingSuspendFunctionExtensionReceiverOwner)) {
+            if (calledDeclaration.receiverTypeRef?.coneType?.isRestrictSuspensionReceiver(session) == true) {
+                return true
+            }
+        }
+        return false
+    }
+
+    private fun ConeKotlinType.isRestrictSuspensionReceiver(session: FirSession): Boolean {
+        when (this) {
+            is ConeClassLikeType -> {
+                val regularClass = fullyExpandedType(session).lookupTag.toFirRegularClass(session) ?: return false
+                if (regularClass.hasAnnotation(RESTRICTS_SUSPENSION_CLASS_ID)) {
+                    return true
+                }
+                return regularClass.superTypeRefs.any { it.coneType.isRestrictSuspensionReceiver(session) }
+            }
+            is ConeTypeParameterType -> {
+                return lookupTag.typeParameterSymbol.resolvedBounds.any { it.coneType.isRestrictSuspensionReceiver(session) }
+            }
+            else -> return false
+        }
+    }
+
+    private fun sameInstanceOfReceiver(useSiteReceiverExpression: FirExpression?, declarationSiteReceiverOwner: FirDeclaration?): Boolean {
+        if (declarationSiteReceiverOwner == null || useSiteReceiverExpression == null) return false
+        if (useSiteReceiverExpression is FirThisReceiverExpression) {
+            return useSiteReceiverExpression.calleeReference.boundSymbol == declarationSiteReceiverOwner.symbol
+        }
+        return false
+    }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDefaultErrorMessages.kt

@@ -190,6 +190,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.GETTER_VISIBILITY
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.HAS_NEXT_FUNCTION_AMBIGUITY
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ILLEGAL_CONST_EXPRESSION
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ILLEGAL_KOTLIN_VERSION_STRING_VALUE
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ILLEGAL_SELECTOR
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ILLEGAL_SUSPEND_FUNCTION_CALL
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ILLEGAL_SUSPEND_PROPERTY_ACCESS
@@ -1492,6 +1493,10 @@ class FirDefaultErrorMessages {
                 SYMBOL
             )
             map.put(NON_LOCAL_SUSPENSION_POINT, "Suspension functions can be called only within coroutine body")
+            map.put(
+                ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL,
+                "Restricted suspending functions can only invoke member or extension suspending functions on their restricted coroutine scope"
+            )
 
             // Extended checkers group
             map.put(REDUNDANT_VISIBILITY_MODIFIER, "Redundant visibility modifier")

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/extensions.fir.kt

@@ -1,128 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER -SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE
-interface SuperInterface
-
-@kotlin.coroutines.RestrictsSuspension
-open class RestrictedController : SuperInterface
-
-class SubClass : RestrictedController()
-
-suspend fun Any?.extAny() {}
-suspend fun SuperInterface.extSuper() {}
-suspend fun RestrictedController.ext() {}
-suspend fun SubClass.extSub() {}
-
-class A {
-    suspend fun Any?.memExtAny() {}
-    suspend fun SuperInterface.memExtSuper() {}
-    suspend fun RestrictedController.memExt() {}
-    suspend fun SubClass.memExtSub() {}
-}
-
-
-fun generate1(f: suspend SuperInterface.() -> Unit) {}
-fun generate2(f: suspend RestrictedController.() -> Unit) {}
-fun generate3(f: suspend SubClass.() -> Unit) {}
-
-fun A.test() {
-    generate1 {
-        extAny()
-        memExtAny()
-        extSuper()
-        memExtSuper()
-        with(A()) {
-            extAny()
-            memExtAny()
-            extSuper()
-            memExtSuper()
-        }
-    }
-    generate2 {
-        extAny()
-        memExtAny()
-        extSuper()
-        memExtSuper()
-
-        ext()
-        memExt()
-        with(A()) {
-            extAny()
-            memExtAny()
-            extSuper()
-            memExtSuper()
-            ext()
-            memExt()
-        }
-    }
-    generate3 {
-        extAny()
-        memExtAny()
-        extSuper()
-        memExtSuper()
-
-        ext()
-        memExt()
-        extSub()
-        memExtSub()
-        with(A()) {
-            extAny()
-            memExtAny()
-            extSuper()
-            memExtSuper()
-            ext()
-            memExt()
-            extSub()
-            memExtSub()
-        }
-    }
-
-    suspend fun SuperInterface.fun1() {
-        extAny()
-        memExtAny()
-        extSuper()
-        memExtSuper()
-        with(A()) {
-            extAny()
-            memExtAny()
-            extSuper()
-            memExtSuper()
-        }
-    }
-    suspend fun RestrictedController.fun2() {
-        extAny()
-        memExtAny()
-        extSuper()
-        memExtSuper()
-
-        ext()
-        memExt()
-        with(A()) {
-            extAny()
-            memExtAny()
-            extSuper()
-            memExtSuper()
-            ext()
-            memExt()
-        }
-    }
-    suspend fun SubClass.fun3() {
-        extAny()
-        memExtAny()
-        extSuper()
-        memExtSuper()
-
-        ext()
-        memExt()
-        extSub()
-        memExtSub()
-        with(A()) {
-            extAny()
-            memExtAny()
-            extSuper()
-            memExtSuper()
-            ext()
-            memExt()
-            extSub()
-            memExtSub()
-        }
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER -SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE
 interface SuperInterface
 

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

@@ -1,61 +0,0 @@
-// SKIP_TXT
-@kotlin.coroutines.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(c: suspend RestrictedController.() -> Unit) {}
-
-fun runBlocking(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.ext()
-            this@a.member()
-            this@a.memberExt()
-
-            this@b.ext()
-            this@b.member()
-            this@b.memberExt()
-        }
-
-        runBlocking {
-            ext()
-            member()
-            memberExt()
-
-            this@a.ext()
-            this@a.member()
-            this@a.memberExt()
-        }
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // SKIP_TXT
 @kotlin.coroutines.RestrictsSuspension
 class RestrictedController {

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/notRelatedFun.fir.kt

@@ -1,71 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER -SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE
-interface SuperInterface
-
-@kotlin.coroutines.RestrictsSuspension
-open class RestrictedController : SuperInterface
-
-class SubClass : RestrictedController()
-
-suspend fun topLevel() {}
-
-class A {
-    suspend fun member() {}
-}
-
-fun generate1(f: suspend SuperInterface.() -> Unit) {}
-fun generate2(f: suspend RestrictedController.() -> Unit) {}
-fun generate3(f: suspend SubClass.() -> Unit) {}
-
-fun A.test() {
-    generate1 {
-        topLevel()
-        member()
-        with(A()) {
-            topLevel()
-            member()
-        }
-    }
-    generate2 {
-        topLevel()
-        member()
-        with(A()) {
-            topLevel()
-            member()
-        }
-    }
-    generate3 {
-        topLevel()
-        member()
-        with(A()) {
-            topLevel()
-            member()
-        }
-    }
-
-    suspend fun SuperInterface.fun1() {
-        topLevel()
-        member()
-        with(A()) {
-            topLevel()
-            member()
-        }
-    }
-    suspend fun RestrictedController.fun2() {
-        topLevel()
-        member()
-        with(A()) {
-            topLevel()
-            member()
-        }
-    }
-    suspend fun SubClass.fun3() {
-        topLevel()
-        member()
-        with(A()) {
-            topLevel()
-            member()
-        }
-    }
-
-}
-

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER -SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE
 interface SuperInterface
 

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield_1_3.fir.kt

@@ -1,99 +0,0 @@
-// !LANGUAGE: +ReleaseCoroutines +ExperimentalBuilderInference
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// SKIP_TXT
-
-@file:OptIn(ExperimentalTypeInference::class)
-
-import kotlin.experimental.ExperimentalTypeInference
-
-@kotlin.coroutines.RestrictsSuspension
-class RestrictedController<T> {
-    suspend fun yield(x: T) {}
-
-    suspend fun anotherYield(x: T) {
-        yield(x)
-        this.yield(x)
-
-        yield2(x)
-        this.yield2(x)
-
-        with(this) {
-            yield(x)
-            this@with.yield(x)
-
-            yield2(x)
-            this@with.yield2(x)
-        }
-    }
-}
-
-fun <T> buildSequence(@BuilderInference c: suspend RestrictedController<T>.() -> Unit) {}
-
-@BuilderInference
-suspend fun <T> RestrictedController<T>.yield2(x: T) {}
-
-fun test() {
-    buildSequence<Int> a@{
-        buildSequence<Int> b@{
-            yield(1)
-            yield2(1)
-            this@b.yield(1)
-            this@b.yield2(1)
-
-            this@a.yield(2) // Should be error
-            this@a.yield2(2) // Should be error
-
-            with(this) {
-                yield(3)
-                this@with.yield(3)
-
-                yield2(3)
-                this@with.yield2(3)
-            }
-        }
-    }
-
-    buildSequence<Int> {
-        buildSequence<String> {
-            yield("a")
-            yield2("a")
-            this.yield("b")
-            this.yield2("b")
-
-            yield(1) // Should be error
-            yield2(1) // Should be error
-
-            with(this) {
-                yield("")
-                this@with.yield("")
-
-                yield2("")
-                this@with.yield2("")
-            }
-        }
-    }
-
-    buildSequence<Int> a@{
-        yield(1)
-        yield2(1)
-        buildSequence {
-            yield("")
-            yield2("")
-            this@a.yield(1)
-            this@a.yield2(1)
-
-            with(this) {
-                yield("")
-                this@with.yield("")
-
-                yield2("")
-                this@with.yield2("")
-            }
-        }
-    }
-
-    buildSequence<String> {
-        yield("")
-        RestrictedController<String>().yield("1")
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ReleaseCoroutines +ExperimentalBuilderInference
 // !USE_EXPERIMENTAL: kotlin.RequiresOptIn
 // SKIP_TXT

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/sameInstance.fir.kt

@@ -1,52 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER -SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE
-@kotlin.coroutines.RestrictsSuspension
-class RestrictedController {
-    suspend fun member() {}
-}
-
-suspend fun RestrictedController.extension() {}
-
-fun generate(f: suspend RestrictedController.() -> Unit) {}
-
-fun test() {
-    generate() l@ {
-        member()
-        extension()
-
-        this.member()
-        this.extension()
-
-        val foo = this
-        foo.member()
-        foo.extension()
-
-        this@l.member()
-        this@l.extension()
-
-        with(1) {
-            this@l.member()
-            this@l.extension()
-        }
-    }
-}
-
-suspend fun RestrictedController.l() {
-    member()
-    extension()
-
-    this.member()
-    this.extension()
-
-    val foo = this
-    foo.member()
-    foo.extension()
-
-    this@l.member()
-    this@l.extension()
-
-    with(1) {
-        this@l.member()
-        this@l.extension()
-    }
-
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER -SUSPENSION_CALL_MUST_BE_USED_AS_RETURN_VALUE
 @kotlin.coroutines.RestrictsSuspension
 class RestrictedController {

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/simpleForbidden.fir.kt

@@ -1,14 +0,0 @@
-@kotlin.coroutines.RestrictsSuspension
-class RestrictedController
-
-suspend fun Any?.extFun() {}
-suspend fun suspendFun() {}
-
-fun generate(c: suspend RestrictedController.() -> Unit) {}
-
-fun test() {
-    generate {
-        extFun()
-        suspendFun()
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 @kotlin.coroutines.RestrictsSuspension
 class RestrictedController
 

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

@@ -1,17 +0,0 @@
-// SKIP_TXT
-@kotlin.coroutines.RestrictsSuspension
-class RestrictedController {
-    suspend fun yield() {}
-}
-
-fun generate(c: suspend RestrictedController.() -> Unit) {}
-
-fun runBlocking(x: suspend () -> Unit) {}
-
-fun test() {
-    generate {
-        runBlocking {
-            yield()
-        }
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // SKIP_TXT
 @kotlin.coroutines.RestrictsSuspension
 class RestrictedController {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -3162,6 +3162,12 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL) { firDiagnostic ->
+        IllegalRestrictedSuspendingFunctionCallImpl(
+            firDiagnostic as FirPsiDiagnostic,
+            token,
+        )
+    }
     add(FirJvmErrors.CONFLICTING_JVM_DECLARATIONS) { firDiagnostic ->
         ConflictingJvmDeclarationsImpl(
             firDiagnostic as FirPsiDiagnostic,

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -2211,6 +2211,10 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         override val diagnosticClass get() = NonLocalSuspensionPoint::class
     }
 
+    abstract class IllegalRestrictedSuspendingFunctionCall : KtFirDiagnostic<PsiElement>() {
+        override val diagnosticClass get() = IllegalRestrictedSuspendingFunctionCall::class
+    }
+
     abstract class ConflictingJvmDeclarations : KtFirDiagnostic<PsiElement>() {
         override val diagnosticClass get() = ConflictingJvmDeclarations::class
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -3568,6 +3568,13 @@ internal class NonLocalSuspensionPointImpl(
     override val firDiagnostic: FirPsiDiagnostic by weakRef(firDiagnostic)
 }
 
+internal class IllegalRestrictedSuspendingFunctionCallImpl(
+    firDiagnostic: FirPsiDiagnostic,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.IllegalRestrictedSuspendingFunctionCall(), KtAbstractFirDiagnostic<PsiElement> {
+    override val firDiagnostic: FirPsiDiagnostic by weakRef(firDiagnostic)
+}
+
 internal class ConflictingJvmDeclarationsImpl(
     firDiagnostic: FirPsiDiagnostic,
     override val token: ValidityToken,