Report RETURN_NOT_ALLOWED and RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY only on the return with label (KT-13340)

#KT-13340 Fixed

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/Errors.java

@@ -774,8 +774,8 @@ public interface Errors {
 
     DiagnosticFactory0<KtSimpleNameExpression> EXPRESSION_EXPECTED_PACKAGE_FOUND = DiagnosticFactory0.create(ERROR);
 
-    DiagnosticFactory0<KtReturnExpression> RETURN_NOT_ALLOWED = DiagnosticFactory0.create(ERROR);
-    DiagnosticFactory0<KtReturnExpression> RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY = DiagnosticFactory0.create(ERROR);
+    DiagnosticFactory0<KtReturnExpression> RETURN_NOT_ALLOWED = DiagnosticFactory0.create(ERROR, PositioningStrategies.RETURN_WITH_LABEL);
+    DiagnosticFactory0<KtReturnExpression> RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY = DiagnosticFactory0.create(ERROR, PositioningStrategies.RETURN_WITH_LABEL);
     DiagnosticFactory0<KtDeclarationWithBody>
             NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY = DiagnosticFactory0.create(ERROR, DECLARATION_WITH_BODY);
 

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

@@ -477,4 +477,14 @@ object PositioningStrategies {
         }
     }
 
+    @JvmField val RETURN_WITH_LABEL: PositioningStrategy<KtReturnExpression> = object: PositioningStrategy<KtReturnExpression>() {
+        override fun mark(element: KtReturnExpression): List<TextRange> {
+            val labeledExpression = element.labeledExpression
+            if (labeledExpression != null) {
+                return markRange(element, labeledExpression)
+            }
+
+            return markElement(element.returnKeyword)
+        }
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/KtReturnExpression.java

@@ -17,8 +17,11 @@
 package org.jetbrains.kotlin.psi;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.KtNodeTypes;
+import org.jetbrains.kotlin.lexer.KtTokens;
 
 public class KtReturnExpression extends KtExpressionWithLabel implements KtStatementExpression {
     public KtReturnExpression(@NotNull ASTNode node) {
@@ -34,4 +37,15 @@ public class KtReturnExpression extends KtExpressionWithLabel implements KtState
     public KtExpression getReturnedExpression() {
         return findChildByClass(KtExpression.class);
     }
+
+    @NotNull
+    public PsiElement getReturnKeyword() {
+        //noinspection ConstantConditions
+        return findChildByType(KtTokens.RETURN_KEYWORD);
+    }
+
+    @Nullable
+    public PsiElement getLabeledExpression() {
+        return findChildByType(KtNodeTypes.LABEL_QUALIFIER);
+    }
 }