[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(
-        LightTreePositioningStrategies.RETURN_KEYWORD,
-        PositioningStrategies.RETURN_KEYWORD
+    val RETURN_WITH_LABEL = SourceElementPositioningStrategy(
+        LightTreePositioningStrategies.RETURN_WITH_LABEL,
+        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
+    )
 }