FIR: introduce RETURN_FOR_BUILT_IN_SUSPEND diagnostic

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

@@ -1194,6 +1194,7 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
         val ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL by error<PsiElement>(PositioningStrategy.REFERENCED_NAME_BY_QUALIFIED)
         val NON_MODIFIER_FORM_FOR_BUILT_IN_SUSPEND by error<PsiElement>(PositioningStrategy.REFERENCED_NAME_BY_QUALIFIED)
         val MODIFIER_FORM_FOR_NON_BUILT_IN_SUSPEND by error<PsiElement>(PositioningStrategy.REFERENCED_NAME_BY_QUALIFIED)
+        val RETURN_FOR_BUILT_IN_SUSPEND by error<KtReturnExpression>()
     }
 }
 

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

@@ -614,5 +614,6 @@ object FirErrors {
     val ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL by error0<PsiElement>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
     val NON_MODIFIER_FORM_FOR_BUILT_IN_SUSPEND by error0<PsiElement>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
     val MODIFIER_FORM_FOR_NON_BUILT_IN_SUSPEND by error0<PsiElement>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
+    val RETURN_FOR_BUILT_IN_SUSPEND by error0<KtReturnExpression>()
 
 }

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

@@ -6,14 +6,18 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.expression
 
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import org.jetbrains.kotlin.fir.FirRealSourceElementKind
 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.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 
 object FirReturnAllowedChecker : FirReturnExpressionChecker() {
     override fun check(expression: FirReturnExpression, context: CheckerContext, reporter: DiagnosticReporter) {
@@ -26,6 +30,25 @@ object FirReturnAllowedChecker : FirReturnExpressionChecker() {
             reporter.reportOn(source, FirErrors.RETURN_NOT_ALLOWED, context)
         }
 
+        if (targetSymbol is FirAnonymousFunctionSymbol) {
+            val label = targetSymbol.label
+            if (label?.source?.kind !is FirRealSourceElementKind) {
+                val functionCall = context.qualifiedAccessOrAnnotationCalls.asReversed().find {
+                    it is FirFunctionCall &&
+                            ((it.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirNamedFunctionSymbol)?.callableId ==
+                            FirSuspendCallChecker.KOTLIN_SUSPEND_BUILT_IN_FUNCTION_CALLABLE_ID
+                }
+                if (functionCall is FirFunctionCall &&
+                    functionCall.arguments.any {
+                        it is FirLambdaArgumentExpression &&
+                                (it.expression as? FirAnonymousFunctionExpression)?.anonymousFunction?.symbol == targetSymbol
+                    }
+                ) {
+                    reporter.reportOn(source, FirErrors.RETURN_FOR_BUILT_IN_SUSPEND, context)
+                }
+            }
+        }
+
         val containingDeclaration = context.containingDeclarations.last()
         if (containingDeclaration is FirFunction && containingDeclaration.body is FirSingleExpressionBlock) {
             reporter.reportOn(source, FirErrors.RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY, context)

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

@@ -29,8 +29,10 @@ 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.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.StandardClassIds
 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
@@ -47,6 +49,8 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
 
     private val BUILTIN_SUSPEND_NAME = KOTLIN_SUSPEND_BUILT_IN_FUNCTION_FQ_NAME.shortName()
 
+    internal val KOTLIN_SUSPEND_BUILT_IN_FUNCTION_CALLABLE_ID = CallableId(StandardClassIds.BASE_KOTLIN_PACKAGE, BUILTIN_SUSPEND_NAME)
+
     @OptIn(SymbolInternals::class)
     override fun check(expression: FirQualifiedAccessExpression, context: CheckerContext, reporter: DiagnosticReporter) {
         val reference = expression.calleeReference as? FirResolvedNamedReference ?: return
@@ -89,7 +93,7 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
         context: CheckerContext,
         reporter: DiagnosticReporter
     ) {
-        if (symbol.callableId.asSingleFqName() == KOTLIN_SUSPEND_BUILT_IN_FUNCTION_FQ_NAME) {
+        if (symbol.callableId == KOTLIN_SUSPEND_BUILT_IN_FUNCTION_CALLABLE_ID) {
             if (reference.name != BUILTIN_SUSPEND_NAME ||
                 expression.explicitReceiver != null ||
                 !expression.hasFormOfSuspendModifierForLambda()

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

@@ -339,6 +339,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.REPEATED_MODIFIER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RESERVED_MEMBER_INSIDE_INLINE_CLASS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RESOLUTION_TO_CLASSIFIER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RESULT_TYPE_MISMATCH
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RETURN_FOR_BUILT_IN_SUSPEND
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RETURN_NOT_ALLOWED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RETURN_TYPE_MISMATCH
@@ -1507,6 +1508,7 @@ class FirDefaultErrorMessages {
                 MODIFIER_FORM_FOR_NON_BUILT_IN_SUSPEND,
                 "Calls having a form of ''suspend {}'' are deprecated because ''suspend'' in the context will have a meaning of a modifier. Add empty argument list to the call: ''suspend() { ... }''"
             )
+            map.put(RETURN_FOR_BUILT_IN_SUSPEND, "Using implicit label for this lambda is prohibited")
 
             // Extended checkers group
             map.put(REDUNDANT_VISIBILITY_MODIFIER, "Redundant visibility modifier")

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -143,7 +143,10 @@ class ExpressionsConverter(
             symbol = FirAnonymousFunctionSymbol()
             isLambda = true
             label = context.firLabels.pop() ?: context.calleeNamesForLambda.lastOrNull()?.let {
-                buildLabel { name = it.asString() }
+                buildLabel {
+                    source = expressionSource.fakeElement(FirFakeSourceElementKind.GeneratedLambdaLabel)
+                    name = it.asString()
+                }
             }
             target = FirFunctionTarget(labelName = label?.name, isLambda = true)
             context.firFunctionTargets += target
@@ -332,7 +335,10 @@ class ExpressionsConverter(
         var firExpression: FirElement? = null
         labeledExpression.forEachChildren {
             when (it.tokenType) {
-                LABEL_QUALIFIER -> context.firLabels += buildLabel { name = it.toString().replace("@", "") }
+                LABEL_QUALIFIER -> context.firLabels += buildLabel {
+                    source = labeledExpression.toFirLightSourceElement(tree)
+                    name = it.toString().replace("@", "")
+                }
                 BLOCK -> firExpression = declarationsConverter.convertBlock(it)
                 PROPERTY -> firExpression = declarationsConverter.convertPropertyDeclaration(it)
                 else -> if (it.isExpression()) firExpression = getAsFirExpression(it)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/impl/FirFunctionSymbol.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.symbols.impl
 
+import org.jetbrains.kotlin.fir.FirLabel
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.symbols.AccessorSymbol
@@ -63,7 +64,9 @@ sealed class FirFunctionWithoutNameSymbol<F : FirFunction>(
     stubName: Name
 ) : FirFunctionSymbol<F>(CallableId(FqName("special"), stubName))
 
-class FirAnonymousFunctionSymbol : FirFunctionWithoutNameSymbol<FirAnonymousFunction>(Name.identifier("anonymous"))
+class FirAnonymousFunctionSymbol : FirFunctionWithoutNameSymbol<FirAnonymousFunction>(Name.identifier("anonymous")) {
+    val label: FirLabel? get() = fir.label
+}
 
 class FirPropertyAccessorSymbol : FirFunctionWithoutNameSymbol<FirPropertyAccessor>(Name.identifier("accessor"))
 

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

@@ -1,12 +1,12 @@
 // SKIP_TXT
 fun bar() {
     suspend {
-        return@suspend
+        <!RETURN_FOR_BUILT_IN_SUSPEND!>return@suspend<!>
     }
 
     suspend {
         run {
-            return@suspend
+            <!RETURN_FOR_BUILT_IN_SUSPEND!>return@suspend<!>
         }
     }
 
@@ -21,7 +21,7 @@ fun bar() {
     val x = suspend@{
         suspend {
             // Might be resolved to outer lambda, but doesn't make sense because suspend-lambdas here is noinline
-            return@suspend
+            <!RETURN_FOR_BUILT_IN_SUSPEND!>return@suspend<!>
         }
     }
 }

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

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

@@ -2223,6 +2223,10 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         override val diagnosticClass get() = ModifierFormForNonBuiltInSuspend::class
     }
 
+    abstract class ReturnForBuiltInSuspend : KtFirDiagnostic<KtReturnExpression>() {
+        override val diagnosticClass get() = ReturnForBuiltInSuspend::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

@@ -3589,6 +3589,13 @@ internal class ModifierFormForNonBuiltInSuspendImpl(
     override val firDiagnostic: FirPsiDiagnostic by weakRef(firDiagnostic)
 }
 
+internal class ReturnForBuiltInSuspendImpl(
+    firDiagnostic: FirPsiDiagnostic,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.ReturnForBuiltInSuspend(), KtAbstractFirDiagnostic<KtReturnExpression> {
+    override val firDiagnostic: FirPsiDiagnostic by weakRef(firDiagnostic)
+}
+
 internal class ConflictingJvmDeclarationsImpl(
     firDiagnostic: FirPsiDiagnostic,
     override val token: ValidityToken,