[FIR] Implement FirReturnAllowedChecker

Supported diagnostics:
- RETURN_NOT_ALLOWED
- RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -259,6 +259,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
         runTest("compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.kt");
     }
 
+    @TestMetadata("labeledReturnFromNotLabeledUnnamedFunction.kt")
+    public void testLabeledReturnFromNotLabeledUnnamedFunction() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/labeledReturnFromNotLabeledUnnamedFunction.kt");
+    }
+
     @TestMetadata("lambdaArgInScopeFunction.kt")
     public void testLambdaArgInScopeFunction() throws Exception {
         runTest("compiler/fir/analysis-tests/testData/resolve/lambdaArgInScopeFunction.kt");

compiler/fir/analysis-tests/testData/resolve/functionTypes.kt

@@ -1,4 +1,4 @@
-fun <T> simpleRun(f: (T) -> Unit): Unit = f(return)
+fun <T> simpleRun(f: (T) -> Unit): Unit = f(<!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!>)
 
 fun <T, R> List<T>.simpleMap(f: (T) -> R): R {
 

compiler/fir/analysis-tests/testData/resolve/labeledReturnFromNotLabeledUnnamedFunction.fir.txt

@@ -0,0 +1,10 @@
+FILE: labeledReturnFromNotLabeledUnnamedFunction.kt
+    public final fun notInline(block: R|(kotlin/Boolean) -> kotlin/Unit|): R|kotlin/String| {
+        ^notInline String()
+    }
+    public final fun test(): R|kotlin/String| {
+        ^test R|/notInline|(fun <anonymous>(b: R|kotlin/Boolean|): R|kotlin/Unit| <inline=NoInline>  {
+            ^@notInline Unit
+        }
+        )
+    }

compiler/fir/analysis-tests/testData/resolve/labeledReturnFromNotLabeledUnnamedFunction.kt

@@ -0,0 +1,9 @@
+fun notInline(block: (Boolean) -> Unit): String {
+    return ""
+}
+
+fun test(): String {
+    return notInline(fun(b: Boolean) {
+        return@notInline
+    })
+}

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -302,6 +302,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.kt");
         }
 
+        @Test
+        @TestMetadata("labeledReturnFromNotLabeledUnnamedFunction.kt")
+        public void testLabeledReturnFromNotLabeledUnnamedFunction() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/labeledReturnFromNotLabeledUnnamedFunction.kt");
+        }
+
         @Test
         @TestMetadata("lambdaArgInScopeFunction.kt")
         public void testLambdaArgInScopeFunction() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -305,6 +305,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/labelAndReceiverForInfix.kt");
         }
 
+        @Test
+        @TestMetadata("labeledReturnFromNotLabeledUnnamedFunction.kt")
+        public void testLabeledReturnFromNotLabeledUnnamedFunction() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/labeledReturnFromNotLabeledUnnamedFunction.kt");
+        }
+
         @Test
         @TestMetadata("lambdaArgInScopeFunction.kt")
         public void testLambdaArgInScopeFunction() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -34413,6 +34413,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/inline"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("inlineConstructorOfArray.kt")
+            public void testInlineConstructorOfArray() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/inline/inlineConstructorOfArray.kt");
+            }
+
             @Test
             @TestMetadata("inlineOnlySuppressesNothingToInline.kt")
             public void testInlineOnlySuppressesNothingToInline() throws Exception {

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

@@ -54,7 +54,7 @@ enum class PositioningStrategy(private val strategy: String? = null) {
     NAME_OF_NAMED_ARGUMENT,
     VALUE_ARGUMENTS,
     SUPERTYPES_LIST,
-    RETURN_KEYWORD,
+    RETURN_WITH_LABEL,
 
     ;
 

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

@@ -59,7 +59,6 @@ object DIAGNOSTICS_LIST : DiagnosticList() {
         val BREAK_OR_CONTINUE_OUTSIDE_A_LOOP by error<FirSourceElement, PsiElement>()
         val NOT_A_LOOP_LABEL by error<FirSourceElement, PsiElement>()
         val VARIABLE_EXPECTED by error<FirSourceElement, PsiElement>()
-        val RETURN_NOT_ALLOWED by error<FirSourceElement, PsiElement>()
         val DELEGATION_IN_INTERFACE by error<FirSourceElement, PsiElement>()
         val NESTED_CLASS_NOT_ALLOWED by error<FirSourceElement, KtNamedDeclaration>(PositioningStrategy.DECLARATION_NAME) {
             parameter<String>("declaration")
@@ -596,6 +595,11 @@ object DIAGNOSTICS_LIST : DiagnosticList() {
         val VARIABLE_NEVER_READ by warning<FirSourceElement, KtNamedDeclaration>(PositioningStrategy.DECLARATION_NAME)
         val USELESS_CALL_ON_NOT_NULL by warning<FirSourceElement, PsiElement>(PositioningStrategy.SELECTOR_BY_QUALIFIED)
     }
+
+    val RETURNS by object : DiagnosticGroup("Returns") {
+        val RETURN_NOT_ALLOWED by error<FirSourceElement, KtReturnExpression>(PositioningStrategy.RETURN_WITH_LABEL)
+        val RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY by error<FirSourceElement, KtReturnExpression>(PositioningStrategy.RETURN_WITH_LABEL)
+    }
 }
 
 private val exposedVisibilityDiagnosticInit: DiagnosticBuilder.() -> Unit = {

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

@@ -46,6 +46,7 @@ import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtPropertyAccessor
 import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtTypeAlias
 import org.jetbrains.kotlin.psi.KtTypeParameter
@@ -76,7 +77,6 @@ object FirErrors {
     val BREAK_OR_CONTINUE_OUTSIDE_A_LOOP by error0<FirSourceElement, PsiElement>()
     val NOT_A_LOOP_LABEL by error0<FirSourceElement, PsiElement>()
     val VARIABLE_EXPECTED by error0<FirSourceElement, PsiElement>()
-    val RETURN_NOT_ALLOWED by error0<FirSourceElement, PsiElement>()
     val DELEGATION_IN_INTERFACE by error0<FirSourceElement, PsiElement>()
     val NESTED_CLASS_NOT_ALLOWED by error1<FirSourceElement, KtNamedDeclaration, String>(SourceElementPositioningStrategies.DECLARATION_NAME)
 
@@ -359,4 +359,8 @@ object FirErrors {
     val VARIABLE_NEVER_READ by warning0<FirSourceElement, KtNamedDeclaration>(SourceElementPositioningStrategies.DECLARATION_NAME)
     val USELESS_CALL_ON_NOT_NULL by warning0<FirSourceElement, PsiElement>(SourceElementPositioningStrategies.SELECTOR_BY_QUALIFIED)
 
+    // Returns
+    val RETURN_NOT_ALLOWED by error0<FirSourceElement, KtReturnExpression>(SourceElementPositioningStrategies.RETURN_WITH_LABEL)
+    val RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY by error0<FirSourceElement, KtReturnExpression>(SourceElementPositioningStrategies.RETURN_WITH_LABEL)
+
 }

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.checkers.expression
+
+import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+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.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+
+object FirReturnAllowedChecker : FirReturnExpressionChecker() {
+    override fun check(expression: FirReturnExpression, context: CheckerContext, reporter: DiagnosticReporter) {
+        val source = expression.source
+        if (source?.kind == FirFakeSourceElementKind.ImplicitReturn) return
+
+        val targetSymbol = expression.target.labeledElement.symbol
+
+        if (!isReturnAllowed(targetSymbol, context)) {
+            reporter.reportOn(source, FirErrors.RETURN_NOT_ALLOWED, 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)
+        }
+    }
+
+    private fun isReturnAllowed(targetSymbol: FirFunctionSymbol<*>, context: CheckerContext): Boolean {
+        for (containingDeclaration in context.containingDeclarations.asReversed()) {
+            when (containingDeclaration) {
+                // return from member of local class or anonymous object
+                is FirClass<*> -> return false
+                is FirFunction<*> -> {
+                    when {
+                        containingDeclaration.symbol == targetSymbol -> return true
+                        containingDeclaration is FirAnonymousFunction -> {
+                            if (!containingDeclaration.inlineStatus.returnAllowed) return false
+                        }
+                        else -> return false
+                    }
+                }
+                is FirProperty -> if (!containingDeclaration.isLocal) return false
+                is FirValueParameter -> return true
+            }
+        }
+        return true
+    }
+}

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

@@ -184,6 +184,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.REDUNDANT_VISIBIL
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.REIFIED_TYPE_IN_CATCH_CLAUSE
 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.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_ON_OVERRIDE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.SEALED_CLASS_CONSTRUCTOR_CALL
@@ -274,7 +275,6 @@ class FirDefaultErrorMessages : DefaultErrorMessages.Extension {
             map.put(BREAK_OR_CONTINUE_OUTSIDE_A_LOOP, "'break' and 'continue' are only allowed inside a loop")
             map.put(NOT_A_LOOP_LABEL, "The label does not denote a loop") // *
             map.put(VARIABLE_EXPECTED, "Variable expected")
-            map.put(RETURN_NOT_ALLOWED, "'return' is not allowed here")
             map.put(DELEGATION_IN_INTERFACE, "Interfaces cannot use delegation")
             map.put(NESTED_CLASS_NOT_ALLOWED, "{0} is not allowed here", TO_STRING)
 
@@ -742,6 +742,10 @@ class FirDefaultErrorMessages : DefaultErrorMessages.Extension {
             // Type alias
             map.put(TOPLEVEL_TYPEALIASES_ONLY, "Nested and local type aliases are not supported")
 
+            // Returns
+            map.put(RETURN_NOT_ALLOWED, "'return' is not allowed here")
+            map.put(RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY, "Returns are not allowed for functions with expression body. Use block body in '{...}'")
+
             // Extended checkers group
             map.put(REDUNDANT_VISIBILITY_MODIFIER, "Redundant visibility modifier")
             map.put(REDUNDANT_MODALITY_MODIFIER, "Redundant modality modifier")

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

@@ -497,13 +497,17 @@ object LightTreePositioningStrategies {
         }
     }
 
-    val RETURN_KEYWORD = object : LightTreePositioningStrategy() {
+    val RETURN_WITH_LABEL = object : LightTreePositioningStrategy() {
         override fun mark(
             node: LighterASTNode,
             startOffset: Int,
             endOffset: Int,
             tree: FlyweightCapableTreeStructure<LighterASTNode>
         ): List<TextRange> {
+            val labeledExpression = tree.findChildByType(node, KtNodeTypes.LABEL_QUALIFIER)
+            if (labeledExpression != null) {
+                return markRange(node, labeledExpression, startOffset, endOffset, tree, node)
+            }
             return markElement(tree.returnKeyword(node) ?: node, startOffset, endOffset, tree)
         }
     }

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

@@ -168,8 +168,8 @@ object SourceElementPositioningStrategies {
         PositioningStrategies.ARRAY_ACCESS
     )
 
-    val RETURN_KEYWORD = SourceElementPositioningStrategy(
+    val RETURN_WITH_LABEL = SourceElementPositioningStrategy(
-        LightTreePositioningStrategies.RETURN_KEYWORD,
+        LightTreePositioningStrategies.RETURN_WITH_LABEL,
-        PositioningStrategies.RETURN_KEYWORD
+        PositioningStrategies.RETURN_WITH_LABEL
     )
 }

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/checkers/CommonExpressionCheckers.kt

@@ -42,4 +42,8 @@ object CommonExpressionCheckers : ExpressionCheckers() {
     override val whenExpressionCheckers: Set<FirWhenExpressionChecker> = setOf(
         FirExhaustiveWhenChecker
     )
+
+    override val returnExpressionCheckers: Set<FirReturnExpressionChecker> = setOf(
+        FirReturnAllowedChecker
+    )
 }

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/PositioningStrategies.kt

@@ -421,12 +421,6 @@ object PositioningStrategies {
         }
     }
 
-    val RETURN_KEYWORD: PositioningStrategy<KtReturnExpression> = object : PositioningStrategy<KtReturnExpression>() {
-        override fun mark(element: KtReturnExpression): List<TextRange> {
-            return markElement(element.returnKeyword)
-        }
-    }
-
     private class ModifierSetBasedPositioningStrategy(private val modifierSet: TokenSet) : PositioningStrategy<KtModifierListOwner>() {
         override fun mark(element: KtModifierListOwner): List<TextRange> {
             val modifierList = element.modifierList

compiler/testData/diagnostics/tests/FunctionReturnTypes.fir.kt

@@ -8,12 +8,12 @@ fun unitEmpty() : Unit {}
 fun unitEmptyReturn() : Unit {return}
 fun unitIntReturn() : Unit {return 1}
 fun unitUnitReturn() : Unit {return Unit}
-fun test1() : Any = {return}
+fun test1() : Any = {<!RETURN_NOT_ALLOWED!>return<!>}
 fun test2() : Any = a@ {return@a 1}
 fun test3() : Any { return }
-fun test4(): ()-> Unit = { return@test4 }
+fun test4(): ()-> Unit = { <!RETURN_NOT_ALLOWED!>return@test4<!> }
 fun test5(): Any = l@{ return@l }
-fun test6(): Any = {return 1}
+fun test6(): Any = {<!RETURN_NOT_ALLOWED!>return<!> 1}
 
 fun bbb() {
     return 1
@@ -137,7 +137,7 @@ fun blockNoReturnIfUnitInOneBranch(): Int {
 fun nonBlockReturnIfEmptyIf(): Int = if (1 < 2) {} else {}
 fun nonBlockNoReturnIfUnitInOneBranch(): Int = if (1 < 2) {} else 2
 
-val a = return 1
+val a = <!RETURN_NOT_ALLOWED!>return<!> 1
 
 class A() {
 }

compiler/testData/diagnostics/tests/Return.fir.kt

@@ -6,12 +6,12 @@ class A {
             if (1 < 2)
                 return@inner
             else
-            return@outer
+            <!RETURN_NOT_ALLOWED!>return@outer<!>
         }
         if (1 < 2)
-            return@A
+            <!RETURN_NOT_ALLOWED!>return@A<!>
         else if (2 < 3)
-            return@inner
+            <!RETURN_NOT_ALLOWED!>return@inner<!>
         return@outer
     }
 }

compiler/testData/diagnostics/tests/ReturnInFunctionWithoutBody.fir.kt

@@ -1,3 +1,3 @@
-fun foo1(): () -> String = return { "some long expression "}
+fun foo1(): () -> String = <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> { "some long expression "}
-fun foo2(): () -> String = return@label { "some long expression "}
+fun foo2(): () -> String = <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY, RETURN_NOT_ALLOWED!>return@label<!> { "some long expression "}
-fun foo3(): () -> String = return<!SYNTAX!>@<!> { "some long expression "}
+fun foo3(): () -> String = <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!><!SYNTAX!>@<!> { "some long expression "}

compiler/testData/diagnostics/tests/controlFlowAnalysis/kt10823.fir.kt

@@ -1,6 +1,6 @@
 fun find2(): Any? {
     fun visit(element: Any) {
-        return@find2 element
+        <!RETURN_NOT_ALLOWED!>return@find2<!> element
     }
     return null
 }
@@ -10,7 +10,7 @@ fun find2(): Any? {
 fun find(): Any? {
     object : Any() {
         fun visit(element: Any) {
-            return@find element
+            <!RETURN_NOT_ALLOWED!>return@find<!> element
         }
     }
     return null
@@ -18,7 +18,7 @@ fun find(): Any? {
 
 fun find4(): Any? {
     inline fun visit(element: Any) {
-        return@find4 element
+        <!RETURN_NOT_ALLOWED!>return@find4<!> element
     }
     return null
 }
@@ -26,7 +26,7 @@ fun find4(): Any? {
 fun find3(): Any? {
     object : Any() {
         inline fun visit(element: Any) {
-            return@find3 element
+            <!RETURN_NOT_ALLOWED!>return@find3<!> element
         }
     }
     return null

compiler/testData/diagnostics/tests/controlStructures/kt799.fir.kt

@@ -12,11 +12,11 @@ fun test() {
     doSmth(if (true) 3 else return, <!TOO_MANY_ARGUMENTS!>1<!>)
 }
 
-val a : Nothing = return 1
+val a : Nothing = <!RETURN_NOT_ALLOWED!>return<!> 1
 
-val b = return 1
+val b = <!RETURN_NOT_ALLOWED!>return<!> 1
 
-val c = doSmth(if (true) 3 else return)
+val c = doSmth(if (true) 3 else <!RETURN_NOT_ALLOWED!>return<!>)
 
 
 fun f(mi: Int = if (true) 0 else return) {}

compiler/testData/diagnostics/tests/controlStructures/localReturnInsidePropertyAccessor.fir.kt

@@ -10,7 +10,7 @@ fun f() = object : ClassData {
 fun g() = object : ClassData {
     init {
         if (true) {
-            return 0
+            <!RETURN_NOT_ALLOWED!>return<!> 0
         }
     }
 

compiler/testData/diagnostics/tests/controlStructures/notAFunctionLabel_after.fir.kt

@@ -31,7 +31,7 @@ fun testAnnotatedLambdaLabel() =
 
 fun testLambdaMultipleLabels1() =
     lambda1@ lambda2@ {
-        return@lambda1
+        <!RETURN_NOT_ALLOWED!>return@lambda1<!>
     }
 
 fun testLambdaMultipleLabels2() =
@@ -46,7 +46,7 @@ fun testAnonymousFunctionLabel() =
 
 fun testLoopLabelInReturn(xs: List<Int>) {
     L@ for (x in xs) {
-        if (x > 0) return@L
+        if (x > 0) <!RETURN_NOT_ALLOWED!>return@L<!>
     }
 }
 

compiler/testData/diagnostics/tests/controlStructures/notAFunctionLabel_before.fir.kt

@@ -31,7 +31,7 @@ fun testAnnotatedLambdaLabel() =
 
 fun testLambdaMultipleLabels1() =
     lambda1@ lambda2@ {
-        return@lambda1
+        <!RETURN_NOT_ALLOWED!>return@lambda1<!>
     }
 
 fun testLambdaMultipleLabels2() =
@@ -46,7 +46,7 @@ fun testAnonymousFunctionLabel() =
 
 fun testLoopLabelInReturn(xs: List<Int>) {
     L@ for (x in xs) {
-        if (x > 0) return@L
+        if (x > 0) <!RETURN_NOT_ALLOWED!>return@L<!>
     }
 }
 

compiler/testData/diagnostics/tests/declarationChecks/ScalaLikeNamedFun.fir.kt

@@ -1,34 +1,34 @@
 // !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
 // KT-5068 Add special error for scala-like syntax 'fun foo(): Int = { 1 }'
 
-fun test1(): Int = { return 1 }
+fun test1(): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
 fun test2(): Int = { 1 }
-val test3: () -> Int = fun (): Int = { return 1 }
+val test3: () -> Int = fun (): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
 val test4: () -> Int = fun (): Int = { 1 }
 fun test5(): Int { return { 1 } }
 fun test6(): Int = fun (): Int = 1
 
 fun outer() {
-    fun test1(): Int = { return 1 }
+    fun test1(): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
     fun test2(): Int = { 1 }
-    val test3: () -> Int = fun (): Int = { return 1 }
+    val test3: () -> Int = fun (): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
     val test4: () -> Int = fun (): Int = { 1 }
     fun test5(): Int { return { 1 } }
     fun test6(): Int = fun (): Int = 1
 }
 
 class Outer {
-    fun test1(): Int = { return 1 }
+    fun test1(): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
     fun test2(): Int = { 1 }
-    val test3: () -> Int = fun (): Int = { return 1 }
+    val test3: () -> Int = fun (): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
     val test4: () -> Int = fun (): Int = { 1 }
     fun test5(): Int { return { 1 } }
     fun test6(): Int = fun (): Int = 1
 
     class Nested {
-        fun test1(): Int = { return 1 }
+        fun test1(): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
         fun test2(): Int = { 1 }
-        val test3: () -> Int = fun (): Int = { return 1 }
+        val test3: () -> Int = fun (): Int = { <!RETURN_NOT_ALLOWED!>return<!> 1 }
         val test4: () -> Int = fun (): Int = { 1 }
         fun test5(): Int { return { 1 } }
         fun test6(): Int = fun (): Int = 1

compiler/testData/diagnostics/tests/functionAsExpression/AutoLabels.fir.kt

@@ -7,5 +7,5 @@ fun autolabel(l: List<Int>) = l.map (fun (i: Int): Int {
 })
 
 fun unresolvedMapLabel(l: List<Int>) = l.map (l@ fun(i: Int): Int {
-    return@map 4
+    <!RETURN_NOT_ALLOWED!>return@map<!> 4
 })

compiler/testData/diagnostics/tests/functionAsExpression/ForbiddenNonLocalReturn.fir.kt

@@ -1,22 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_VARIABLE
-
-fun test() {
-    fun bar() {
-        val bas = fun() {
-            return@bar
-        }
-    }
-
-    val bar = fun() {
-        return@test
-    }
-}
-
-fun foo() {
-    val bal = bag@ fun () {
-        val bar = fun() {
-            return@bag
-        }
-        return@bag
-    }
-}

compiler/testData/diagnostics/tests/functionAsExpression/ForbiddenNonLocalReturn.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_VARIABLE
 
 fun test() {

compiler/testData/diagnostics/tests/functionLiterals/return/AutoLabelsNonLocal.fir.kt

@@ -1,11 +0,0 @@
-fun f() {
-    foo {
-        bar {
-            return@foo 1
-        }
-        return@foo 1
-    }
-}
-
-fun foo(a: Any) {}
-fun bar(a: Any) {}

compiler/testData/diagnostics/tests/functionLiterals/return/AutoLabelsNonLocal.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun f() {
     foo {
         bar {

compiler/testData/diagnostics/tests/functionLiterals/return/ForbiddenNonLocalReturnNoType.fir.kt

@@ -1,19 +0,0 @@
-// !WITH_NEW_INFERENCE
-fun test() {
-    run {return}
-    run {}
-}
-
-fun test2() {
-    run {return@test2}
-    run {}
-}
-
-fun test3() {
-    fun test4() {
-        run {return@test3}
-        run {}
-    }
-}
-
-fun <T> run(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/ForbiddenNonLocalReturnNoType.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !WITH_NEW_INFERENCE
 fun test() {
     run {<!RETURN_NOT_ALLOWED!>return<!>}

compiler/testData/diagnostics/tests/functionLiterals/return/LocalAndNonLocalReturnInLambda.fir.kt

@@ -1,11 +0,0 @@
-// !CHECK_TYPE
-
-fun test2(a: Int) {
-    val x = run f@{
-      if (a > 0) return
-      return@f 1
-    }
-    checkSubtype<Int>(x)
-}
-
-fun <T> run(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalAndNonLocalReturnInLambda.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !CHECK_TYPE
 
 fun test2(a: Int) {

compiler/testData/diagnostics/tests/inline/nonLocalReturns/labeledReturn.fir.kt

@@ -1,31 +0,0 @@
-fun inlineCall(): String {
-    inlineFun {
-        if (true) {
-            return@inlineCall ""
-        }
-        1
-    }
-
-    return "x"
-}
-
-inline fun inlineFun(s: ()->Int) {
-    s()
-}
-
-
-fun noInlineCall(): String {
-    noInline {
-        if (true) {
-            return@noInlineCall ""
-        }
-        1
-    }
-
-    return "x"
-}
-
-
-fun noInline(s: ()->Int) {
-    s()
-}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/labeledReturn.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun inlineCall(): String {
     inlineFun {
         if (true) {

compiler/testData/diagnostics/tests/inline/nonLocalReturns/noInlineAnnotation.fir.kt

@@ -1,10 +0,0 @@
-// !DIAGNOSTICS: -NOTHING_TO_INLINE
-fun main() {
-    test {
-        return
-    }
-}
-
-inline fun test(noinline lambda: () -> Unit) {
-    lambda()
-}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/noInlineAnnotation.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -NOTHING_TO_INLINE
 fun main() {
     test {

compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambda.fir.kt

@@ -1,43 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE -NOTHING_TO_INLINE
-
-fun <R> fun1(p: () -> R) {
-    inlineFun {
-        p()
-    }
-}
-
-fun <R> fun1ValueArgument(p: () -> R) {
-    inlineFun ({
-                   p()
-               })
-}
-
-fun <R> fun3(p: () -> R) {
-    inlineFun {
-        return;
-    }
-}
-
-fun <R> fun3ValueArgument(p: () -> R) {
-    inlineFun ({
-                   return;
-               })
-}
-
-
-fun <R> fun4(p: () -> R) {
-    inlineFun lambda@ {
-        return@lambda p();
-    }
-}
-
-fun <R> fun4ValueArgument(p: () -> R) {
-    inlineFun (lambda@ {
-        return@lambda p();
-    })
-}
-
-
-inline fun <R> inlineFun(crossinline p: () -> R) {
-    p()
-}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambda.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE -NOTHING_TO_INLINE
 
 fun <R> fun1(p: () -> R) {

compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambdaBinaryExpr.fir.kt

@@ -1,25 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE -NOTHING_TO_INLINE
-
-class Z {
-    inline infix fun <R> inlineFun(crossinline p: () -> R) {
-        p()
-    }
-}
-
-fun <R> fun1(p: () -> R) {
-    Z() inlineFun {
-        p()
-    }
-}
-
-fun <R> fun3(p: () -> R) {
-    Z() inlineFun {
-        return;
-    }
-}
-
-fun <R> fun4(p: () -> R) {
-    Z() inlineFun lambda@ {
-        return@lambda p();
-    }
-}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambdaBinaryExpr.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE -NOTHING_TO_INLINE
 
 class Z {

compiler/testData/diagnostics/tests/inline/nonLocalReturns/propertyAccessorsAndConstructor.fir.kt

@@ -1,24 +0,0 @@
-inline fun <R> doCall(p: () -> R) {
-    p()
-}
-
-inline fun <R> doCallInt(p: () -> R): R {
-    return p()
-}
-
-class A {
-    var result: Int = doCallInt { return this };
-
-    var field: Int
-        get() {
-            doCall { return 1 }
-            return 2
-        }
-        set(v: Int) {
-            doCall {
-                result = v / 2
-                return
-            }
-            result = v
-        }
-}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/propertyAccessorsAndConstructor.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 inline fun <R> doCall(p: () -> R) {
     p()
 }

compiler/testData/diagnostics/tests/labels/kt361.fir.kt

@@ -1,11 +0,0 @@
-fun nonlocals(b : Boolean) {
-    a@{
-        fun foo() {
-            if (b) {
-                return@a 1 // The label must be resolved, but an error should be reported for a non-local return
-            }
-        }
-
-        return@a 5
-    }
-}

compiler/testData/diagnostics/tests/labels/kt361.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun nonlocals(b : Boolean) {
     a@{
         fun foo() {

compiler/testData/diagnostics/tests/regressions/kt411.fir.kt

@@ -23,7 +23,7 @@ fun t2() : String {
         if (true) {
             return@l 1
         }
-        return "s"
+        <!RETURN_NOT_ALLOWED!>return<!> "s"
     }
     return "s"
 }
@@ -32,10 +32,10 @@ fun t3() : String {
     invoker(
     l@{
         if (true) {
-            return@t3 "1"
+            <!RETURN_NOT_ALLOWED!>return@t3<!> "1"
         }
         else {
-            return 2
+            <!RETURN_NOT_ALLOWED!>return<!> 2
         }
         return@l 0
     }
@@ -47,7 +47,7 @@ fun t3() : String {
     )
     invoker(
     {
-        return "0"
+        <!RETURN_NOT_ALLOWED!>return<!> "0"
     }
     )
     return "2"

compiler/testData/diagnostics/tests/secondaryConstructors/return.fir.kt

@@ -1,8 +1,8 @@
 // !WITH_NEW_INFERENCE
 class A {
     init {
-        return
+        <!RETURN_NOT_ALLOWED!>return<!>
-        return 1
+        <!RETURN_NOT_ALLOWED!>return<!> 1
     }
     constructor() {
         if (1 == 1) {

compiler/testData/diagnostics/testsWithStdLib/coroutines/suspendFunctionType/inline.fir.kt

@@ -16,15 +16,15 @@ suspend fun baz() {
     }
 
     foo2 {
-        return@baz
+        <!RETURN_NOT_ALLOWED!>return@baz<!>
     }
 
     foo3 {
-        return@baz
+        <!RETURN_NOT_ALLOWED!>return@baz<!>
     }
 
     foo4 {
-        return@baz
+        <!RETURN_NOT_ALLOWED!>return@baz<!>
     }
 
     bar1 {
@@ -32,14 +32,14 @@ suspend fun baz() {
     }
 
     bar2 {
-        return@baz
+        <!RETURN_NOT_ALLOWED!>return@baz<!>
     }
 
     bar3 {
-        return@baz
+        <!RETURN_NOT_ALLOWED!>return@baz<!>
     }
 
     bar4 {
-        return@baz
+        <!RETURN_NOT_ALLOWED!>return@baz<!>
     }
 }

compiler/testData/diagnostics/testsWithStdLib/factoryPattern/returnFromInlineLambda.fir.kt

@@ -1,51 +0,0 @@
-// !LANGUAGE: +NewInference +OverloadResolutionByLambdaReturnType
-// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE -UNUSED_EXPRESSION -EXPERIMENTAL_API_USAGE -EXPERIMENTAL_UNSIGNED_LITERALS
-// ISSUE: KT-11265
-
-// FILE: OverloadResolutionByLambdaReturnType.kt
-
-package kotlin
-
-annotation class OverloadResolutionByLambdaReturnType
-
-// FILE: main.kt
-
-import kotlin.OverloadResolutionByLambdaReturnType
-
-inline fun <T, R> Array<out T>.myFlatMap(transform: (T) -> Iterable<R>): List<R> {
-    TODO()
-}
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("seqFlatMap")
-inline fun <T, R> Array<out T>.myFlatMap(transform: (T) -> Sequence<R>): List<R> {
-    TODO()
-}
-
-fun String.toList(): List<String> = null!!
-
-fun test_1(a: Array<String>, b: Boolean) {
-    a.myFlatMap { it.toList().ifEmpty { return } }
-    a.myFlatMap {
-        if (b) return
-        it.toList()
-    }
-}
-
-fun <T, R> Array<out T>.noInlineFlatMap(transform: (T) -> Iterable<R>): List<R> {
-    TODO()
-}
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("noInlineSeqFlatMap")
-fun <T, R> Array<out T>.noInlineFlatMap(transform: (T) -> Sequence<R>): List<R> {
-    TODO()
-}
-
-fun test_2(a: Array<String>, b: Boolean) {
-    a.noInlineFlatMap { it.toList().ifEmpty { return } }
-    a.noInlineFlatMap {
-        if (b) return
-        it.toList()
-    }
-}

compiler/testData/diagnostics/testsWithStdLib/factoryPattern/returnFromInlineLambda.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +NewInference +OverloadResolutionByLambdaReturnType
 // !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE -UNUSED_EXPRESSION -EXPERIMENTAL_API_USAGE -EXPERIMENTAL_UNSIGNED_LITERALS
 // ISSUE: KT-11265

compiler/testData/diagnostics/testsWithStdLib/inline/inlineConstructorOfArray.kt

@@ -0,0 +1,20 @@
+// FIR_IDENTICAL
+fun testArray(b: Boolean) {
+    Array(5) { i ->
+        if (b) return
+        i
+    }
+    throw AssertionError()
+}
+
+fun testMyArray(b: Boolean) {
+    MyArray(5) { i ->
+        if (b) <!RETURN_NOT_ALLOWED!>return<!>
+        i
+    }
+    throw AssertionError()
+}
+
+class MyArray<T> {
+    constructor(size: Int, init: (Int) -> T)
+}

compiler/testData/diagnostics/testsWithStdLib/inline/inlineConstructorOfArray.txt

@@ -0,0 +1,11 @@
+package
+
+public fun testArray(/*0*/ b: kotlin.Boolean): kotlin.Unit
+public fun testMyArray(/*0*/ b: kotlin.Boolean): kotlin.Unit
+
+public final class MyArray</*0*/ T> {
+    public constructor MyArray</*0*/ T>(/*0*/ size: kotlin.Int, /*1*/ init: (kotlin.Int) -> T)
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -34509,6 +34509,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/inline"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("inlineConstructorOfArray.kt")
+            public void testInlineConstructorOfArray() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/inline/inlineConstructorOfArray.kt");
+            }
+
             @Test
             @TestMetadata("inlineOnlySuppressesNothingToInline.kt")
             public void testInlineOnlySuppressesNothingToInline() throws Exception {

compiler/tests-spec/testData/diagnostics/linked/expressions/when-expression/p-6/pos/5.2.fir.kt

@@ -225,8 +225,8 @@ fun case_23(value_1: Nothing) {
 
 // TESTCASE NUMBER: 24
 fun case_24(value_1: Nothing?) = when (value_1) {
-    throw Exception(), return "" -> ""
+    throw Exception(), <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> "" -> ""
-    null, return return return "", throw throw throw Exception() -> ""
+    null, <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> "", throw throw throw Exception() -> ""
     else -> ""
 }
 
@@ -237,8 +237,8 @@ fun case_24(value_1: Nothing?) = when (value_1) {
  */
 fun case_25(value_1: Boolean) = when (value_1) {
     true -> {}
-    throw Exception(), return -> {}
+    throw Exception(), <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> -> {}
-    false, return return return, throw throw throw Exception() -> {}
+    false, <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!> <!RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY!>return<!>, throw throw throw Exception() -> {}
 }
 
 /*

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

@@ -111,12 +111,6 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
-    add(FirErrors.RETURN_NOT_ALLOWED) { firDiagnostic ->
-        ReturnNotAllowedImpl(
-            firDiagnostic as FirPsiDiagnostic<*>,
-            token,
-        )
-    }
     add(FirErrors.DELEGATION_IN_INTERFACE) { firDiagnostic ->
         DelegationInInterfaceImpl(
             firDiagnostic as FirPsiDiagnostic<*>,
@@ -1633,4 +1627,16 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.RETURN_NOT_ALLOWED) { firDiagnostic ->
+        ReturnNotAllowedImpl(
+            firDiagnostic as FirPsiDiagnostic<*>,
+            token,
+        )
+    }
+    add(FirErrors.RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY) { firDiagnostic ->
+        ReturnInFunctionWithExpressionBodyImpl(
+            firDiagnostic as FirPsiDiagnostic<*>,
+            token,
+        )
+    }
 }

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

@@ -40,6 +40,7 @@ import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtPropertyAccessor
 import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtTypeAlias
 import org.jetbrains.kotlin.psi.KtTypeParameter
@@ -97,10 +98,6 @@ sealed class KtFirDiagnostic<PSI: PsiElement> : KtDiagnosticWithPsi<PSI> {
         override val diagnosticClass get() = VariableExpected::class
     }
 
-    abstract class ReturnNotAllowed : KtFirDiagnostic<PsiElement>() {
-        override val diagnosticClass get() = ReturnNotAllowed::class
-    }
-
     abstract class DelegationInInterface : KtFirDiagnostic<PsiElement>() {
         override val diagnosticClass get() = DelegationInInterface::class
     }
@@ -1149,4 +1146,12 @@ sealed class KtFirDiagnostic<PSI: PsiElement> : KtDiagnosticWithPsi<PSI> {
         override val diagnosticClass get() = UselessCallOnNotNull::class
     }
 
+    abstract class ReturnNotAllowed : KtFirDiagnostic<KtReturnExpression>() {
+        override val diagnosticClass get() = ReturnNotAllowed::class
+    }
+
+    abstract class ReturnInFunctionWithExpressionBody : KtFirDiagnostic<KtReturnExpression>() {
+        override val diagnosticClass get() = ReturnInFunctionWithExpressionBody::class
+    }
+
 }

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

@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtPropertyAccessor
 import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtTypeAlias
 import org.jetbrains.kotlin.psi.KtTypeParameter
@@ -128,13 +129,6 @@ internal class VariableExpectedImpl(
     override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
 }
 
-internal class ReturnNotAllowedImpl(
-    firDiagnostic: FirPsiDiagnostic<*>,
-    override val token: ValidityToken,
-) : KtFirDiagnostic.ReturnNotAllowed(), KtAbstractFirDiagnostic<PsiElement> {
-    override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
-}
-
 internal class DelegationInInterfaceImpl(
     firDiagnostic: FirPsiDiagnostic<*>,
     override val token: ValidityToken,
@@ -1864,3 +1858,17 @@ internal class UselessCallOnNotNullImpl(
     override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
 }
 
+internal class ReturnNotAllowedImpl(
+    firDiagnostic: FirPsiDiagnostic<*>,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.ReturnNotAllowed(), KtAbstractFirDiagnostic<KtReturnExpression> {
+    override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
+}
+
+internal class ReturnInFunctionWithExpressionBodyImpl(
+    firDiagnostic: FirPsiDiagnostic<*>,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.ReturnInFunctionWithExpressionBody(), KtAbstractFirDiagnostic<KtReturnExpression> {
+    override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
+}
+

idea/testData/checker/FunctionReturnTypes.fir.kt

@@ -8,7 +8,7 @@ fun unitEmpty() : Unit {}
 fun unitEmptyReturn() : Unit {return}
 fun unitIntReturn() : Unit {return 1}
 fun unitUnitReturn() : Unit {return Unit}
-fun test1() : Any = { return }
+fun test1() : Any = { <error descr="[RETURN_NOT_ALLOWED] 'return' is not allowed here">return</error> }
 fun test2() : Any = a@ {return@a 1}
 fun test3() : Any { return }
 
@@ -133,7 +133,7 @@ fun blockNoReturnIfUnitInOneBranch(): Int {
 fun nonBlockReturnIfEmptyIf(): Int = if (1 < 2) {} else {}
 fun nonBlockNoReturnIfUnitInOneBranch(): Int = if (1 < 2) {} else 2
 
-val a = return 1
+val a = <error descr="[RETURN_NOT_ALLOWED] 'return' is not allowed here">return</error> 1
 
 class A() {
 }