[java-intentions] IDEA-229521 "Return outside of enclosing switch statement" should have a fix to replace 'return' with 'yield'

GitOrigin-RevId: e8cc607a3f5a30e9fac7ea69671062ac14dc2673

java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightUtil.java

@@ -579,10 +579,14 @@ public final class HighlightUtil {
     return highlightInfo;
   }
 
-  static HighlightInfo checkReturnFromSwitchExpr(@NotNull PsiStatement statement) {
+  static HighlightInfo checkReturnFromSwitchExpr(@NotNull PsiReturnStatement statement) {
     if (PsiImplUtil.findEnclosingSwitchExpression(statement) != null) {
       String message = JavaErrorBundle.message("return.outside.switch.expr");
-      return HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR).range(statement).descriptionAndTooltip(message).create();
+      HighlightInfo info = HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR).range(statement).descriptionAndTooltip(message).create();
+      if (statement.getReturnValue() != null) {
+        QuickFixAction.registerQuickFixAction(info, new ReplaceWithYieldFix(statement));
+      }
+      return info;
     }
 
     return null;

java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/ReplaceWithYieldFix.java

@@ -0,0 +1,45 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.daemon.impl.quickfix;
+
+import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.codeInspection.CommonQuickFixBundle;
+import com.intellij.codeInspection.LocalQuickFixAndIntentionActionOnPsiElement;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.*;
+import com.siyeh.ig.psiutils.CommentTracker;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class ReplaceWithYieldFix extends LocalQuickFixAndIntentionActionOnPsiElement {
+  public ReplaceWithYieldFix(@NotNull PsiReturnStatement statement) {
+    super(statement);
+  }
+
+  @Override
+  public void invoke(@NotNull Project project,
+                     @NotNull PsiFile file,
+                     @Nullable Editor editor,
+                     @NotNull PsiElement startElement,
+                     @NotNull PsiElement endElement) {
+    PsiReturnStatement returnStatement = (PsiReturnStatement)startElement;
+    PsiExpression returnValue = returnStatement.getReturnValue();
+    if (returnValue == null) {
+      return;
+    }
+    TextRange range = returnStatement.getFirstChild().getTextRange();
+    // Work on document level to preserve formatting
+    file.getViewProvider().getDocument().replaceString(range.getStartOffset(), range.getEndOffset(), PsiKeyword.YIELD);
+  }
+
+  @Override
+  public @NotNull String getText() {
+    return CommonQuickFixBundle.message("fix.replace.with.x", PsiKeyword.YIELD);
+  }
+
+  @Override
+  public @NotNull String getFamilyName() {
+    return getText();
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/replaceWithYield/afterSimple.java

@@ -0,0 +1,12 @@
+// "Replace with 'yield'" "true-preview"
+class X {
+  int test(String s) {
+    return switch (s) {
+      case "foo" -> {
+        System.out.println("hello");
+        yield 123;
+      }
+      case "bar" -> 2;
+    };
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/replaceWithYield/beforeNoValue.java

@@ -0,0 +1,12 @@
+// "Replace with 'yield'" "false"
+class X {
+  int test(String s) {
+    return switch (s) {
+      case "foo" -> {
+        System.out.println("hello");
+        return<caret>;
+      }
+      case "bar" -> 2;
+    };
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/replaceWithYield/beforeSimple.java

@@ -0,0 +1,12 @@
+// "Replace with 'yield'" "true-preview"
+class X {
+  int test(String s) {
+    return switch (s) {
+      case "foo" -> {
+        System.out.println("hello");
+        return <caret>123;
+      }
+      case "bar" -> 2;
+    };
+  }
+}

java/java-tests/testSrc/com/intellij/codeInsight/daemon/impl/quickfix/ReplaceWithYieldFixTest.java

@@ -0,0 +1,12 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.daemon.impl.quickfix;
+
+import com.intellij.codeInsight.daemon.quickFix.LightQuickFixParameterizedTestCase;
+
+
+public class ReplaceWithYieldFixTest extends LightQuickFixParameterizedTestCase {
+  @Override
+  protected String getBasePath() {
+    return "/codeInsight/daemonCodeAnalyzer/quickFix/replaceWithYield";
+  }
+}