FIR: introduce NON_LOCAL_SUSPENSION_POINT diagnostic

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

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

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

@@ -610,5 +610,6 @@ object FirErrors {
     // Suspend errors
     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)
 
 }

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.types.coneType
 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
+import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
 
 object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
     private val SUSPEND_PROPERTIES_FQ_NAMES = setOf(
@@ -48,6 +49,10 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
                 else -> {
                 }
             }
+        } else {
+            if (!checkNonLocalReturnUsage(enclosingSuspendFunction, context)) {
+                reporter.reportOn(expression.source, FirErrors.NON_LOCAL_SUSPENSION_POINT, context)
+            }
         }
     }
 
@@ -60,4 +65,13 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
             }
         } as? FirFunction
     }
+
+    private fun checkNonLocalReturnUsage(enclosingSuspendFunction: FirFunction, context: CheckerContext): Boolean {
+        val containingFunction = context.containingDeclarations.lastIsInstanceOrNull<FirFunction>() ?: return false
+        return if (containingFunction is FirAnonymousFunction && enclosingSuspendFunction !== containingFunction) {
+            containingFunction.inlineStatus.returnAllowed
+        } else {
+            enclosingSuspendFunction === containingFunction
+        }
+    }
 }

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

@@ -261,6 +261,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_EXHAUSTIVE_WH
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_FINAL_MEMBER_IN_FINAL_CLASS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_FINAL_MEMBER_IN_OBJECT
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_LOCAL_RETURN_NOT_ALLOWED
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_LOCAL_SUSPENSION_POINT
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_MEMBER_FUNCTION_NO_BODY
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_PRIVATE_CONSTRUCTOR_IN_ENUM
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NON_PRIVATE_OR_PROTECTED_CONSTRUCTOR_IN_SEALED
@@ -1490,6 +1491,7 @@ class FirDefaultErrorMessages {
                 "Suspend property ''{0}'' should be accessed only from a coroutine or suspend function",
                 SYMBOL
             )
+            map.put(NON_LOCAL_SUSPENSION_POINT, "Suspension functions can be called only within coroutine body")
 
             // Extended checkers group
             map.put(REDUNDANT_VISIBILITY_MODIFIER, "Redundant visibility modifier")

compiler/testData/diagnostics/tests/inference/recursiveCalls/kt23531.fir.kt

@@ -14,5 +14,5 @@ suspend fun insideJob3() = doTheJob3()
 
 fun doTheJob0() = simpleAsync0 { <!ARGUMENT_TYPE_MISMATCH!>insideJob0()<!> }
 fun doTheJob1() = simpleAsync1 { <!ARGUMENT_TYPE_MISMATCH!>insideJob1()<!> }
-suspend fun doTheJob2() = simpleAsync2 { <!ARGUMENT_TYPE_MISMATCH!>insideJob2()<!> }
+suspend fun doTheJob2() = simpleAsync2 { <!ARGUMENT_TYPE_MISMATCH!><!NON_LOCAL_SUSPENSION_POINT!>insideJob2<!>()<!> }
 suspend fun doTheJob3() = simpleAsync3 { <!ARGUMENT_TYPE_MISMATCH!>insideJob3()<!> }

compiler/testData/diagnostics/tests/inference/returningLambdaInSuspendContext.fir.kt

@@ -1,39 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER -UNCHECKED_CAST -UNUSED_LAMBDA_EXPRESSION
-
-class Foo {
-    suspend operator fun <T> invoke(body: suspend (Int) -> T) = null as T
-    suspend fun <T> bar(body: suspend (Int) -> T) = null as T
-}
-
-fun <T> runBlocking(block: suspend () -> T) = null as T
-
-public inline fun <T, R> T.run(block: T.() -> R) = null as R
-
-fun main() {
-    val retry = Foo()
-
-    runBlocking {
-        retry { 1 } // OK only after fix
-    }
-    runBlocking {
-        retry { 1 } // OK
-        1
-    }
-    runBlocking {
-        retry.bar { 1 } // OK
-    }
-    runBlocking {
-        { retry { 1 } } // should be error
-    }
-    runBlocking {
-        10.run { retry { 1 } } // should be OK
-    }
-    runBlocking {
-        10.run { retry { 1 } } // should be OK
-        1
-    }
-    runBlocking {
-        { retry { 1 } } // should be error
-        1
-    }
-}

compiler/testData/diagnostics/tests/inference/returningLambdaInSuspendContext.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER -UNCHECKED_CAST -UNUSED_LAMBDA_EXPRESSION
 
 class Foo {

compiler/testData/diagnostics/testsWithStdLib/coroutines/inlineCrossinline/inlineSuspendOfCrossinlineOrdinary.fir.kt

@@ -32,7 +32,7 @@ suspend fun calculate() = "OK"
 fun box() {
     builder {
         test {
-            calculate()
+            <!NON_LOCAL_SUSPENSION_POINT!>calculate<!>()
         }
     }
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/inlineCrossinline/inlineSuspendOfNoinlineOrdinary.fir.kt

@@ -32,7 +32,7 @@ suspend fun calculate() = "OK"
 fun box() {
     builder {
         test {
-            calculate()
+            <!NON_LOCAL_SUSPENSION_POINT!>calculate<!>()
         }
     }
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/nonLocalSuspension.fir.kt

@@ -1,81 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-import Host.suspendFromObject
-
-suspend fun suspendHere() = 1
-suspend fun <T> another(a: T) = 1
-
-object Host {
-    suspend fun suspendFromObject() = 1
-}
-
-fun <T> builder(c: suspend () -> Unit) { }
-
-inline fun run(x: () -> Unit) {}
-
-inline fun runCross(crossinline x: () -> Unit) {}
-
-fun noinline(x: () -> Unit) {}
-
-fun foo() {
-    var result = 1
-    builder<String> {
-        suspendHere()
-        suspendFromObject()
-        
-        another("")
-        another(1)
-
-        result += suspendHere()
-        result += suspendFromObject()
-
-        run {
-            result += suspendHere()
-            result += suspendFromObject()
-            
-            run {
-                suspendHere()
-                suspendFromObject()
-            }
-        }
-
-        runCross {
-            result += suspendHere()
-            result += suspendFromObject()
-            runCross {
-                suspendHere()
-                suspendFromObject()
-            }
-        }
-
-        noinline {
-            result += suspendHere()
-            result += suspendFromObject()
-            noinline {
-                suspendHere()
-                suspendFromObject()
-            }
-        }
-
-        class A {
-            fun bar() {
-                suspendHere()
-                suspendFromObject()
-            }
-        }
-
-        object : Any() {
-            fun baz() {
-                suspendHere()
-                suspendFromObject()
-            }
-        }
-
-        builder<Int> {
-            suspendHere()
-            suspendFromObject()
-
-            another(1)
-            another("")
-        }
-    }
-}

compiler/testData/diagnostics/testsWithStdLib/coroutines/nonLocalSuspension.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 import Host.suspendFromObject
 

compiler/testData/diagnostics/testsWithStdLib/coroutines/suspesionInDefaultValue.fir.kt

@@ -6,10 +6,10 @@ suspend fun foo() = 1
 suspend fun bar(
         x: Int = 2 + foo(),
         y: suspend () -> Int = { foo() },
-        z: () -> Int = { foo() },
+        z: () -> Int = { <!NON_LOCAL_SUSPENSION_POINT!>foo<!>() },
         w: Int = myInline { foo() },
         v: Any? = object {
-            fun x() = foo()
+            fun x() = <!NON_LOCAL_SUSPENSION_POINT!>foo<!>()
             suspend fun y() = foo()
         }
 ) {}

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

@@ -3156,6 +3156,12 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.NON_LOCAL_SUSPENSION_POINT) { firDiagnostic ->
+        NonLocalSuspensionPointImpl(
+            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

@@ -2207,6 +2207,10 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         abstract val suspendCallable: KtSymbol
     }
 
+    abstract class NonLocalSuspensionPoint : KtFirDiagnostic<PsiElement>() {
+        override val diagnosticClass get() = NonLocalSuspensionPoint::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

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