[java-inspections] IDEA-286507 Warn when colon-style expression switch uses yield or throw in every branch

GitOrigin-RevId: 2e79504ab44609427f51fe5652cb995ec72d17d6

java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java

@@ -59,6 +59,49 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
   @Override
   public @NotNull PsiElementVisitor buildVisitor(@NotNull ProblemsHolder holder, boolean isOnTheFly) {
     return new JavaElementVisitor() {
+
+      @Override
+      public void visitSwitchExpression(@NotNull PsiSwitchExpression expression) {
+        PsiElement switchKeyword = expression.getFirstChild();
+        if (switchKeyword == null) {
+          return;
+        }
+        PsiCodeBlock body = expression.getBody();
+        if (body == null) return;
+        boolean onlyOneYieldAfterLabel = true;
+        boolean isNotRule = true;
+        int statementAfterLabelCount = 0;
+        for (PsiStatement statement : body.getStatements()) {
+          if (statement instanceof PsiSwitchLabeledRuleStatement) {
+            isNotRule = false;
+            break;
+          }
+          if (statement instanceof PsiSwitchLabelStatement) {
+            statementAfterLabelCount = 0;
+          }
+          else {
+            statementAfterLabelCount++;
+            if (!(statement instanceof PsiYieldStatement || statement instanceof PsiThrowStatement)) {
+              onlyOneYieldAfterLabel = false;
+            }
+            else {
+              if (statementAfterLabelCount > 1) {
+                onlyOneYieldAfterLabel = false;
+              }
+            }
+          }
+        }
+        if (!isNotRule) {
+          return;
+        }
+        ProblemHighlightType warningType = ProblemHighlightType.GENERIC_ERROR_OR_WARNING;
+        if (!onlyOneYieldAfterLabel) {
+          warningType = ProblemHighlightType.INFORMATION;
+        }
+        holder.registerProblem(switchKeyword, JavaBundle.message("inspection.switch.expression.migration.inspection.switch.expression.description"),
+                               warningType, new ReplaceExpressionWithEnhancedSwitchExpressionFix());
+      }
+
       @Override
       public void visitSwitchStatement(@NotNull PsiSwitchStatement statement) {
         PsiElement switchKeyword = statement.getFirstChild();
@@ -290,6 +333,81 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
     String generate(CommentTracker ct, SwitchBranch branch);
   }
 
+  private static class ReplaceExpressionWithEnhancedSwitchExpressionFix extends PsiUpdateModCommandQuickFix{
+
+    @Override
+    protected void applyFix(@NotNull Project project, @NotNull PsiElement element, @NotNull ModPsiUpdater updater) {
+      StringBuilder builder = new StringBuilder();
+      if(!(element.getParent() instanceof PsiSwitchExpression switchExpression)) return;
+      for (@NotNull PsiElement switchExpressionChild : switchExpression.getChildren()) {
+        if (switchExpressionChild instanceof PsiCodeBlock codeBlock) {
+          for (@NotNull PsiElement codeBlockChildren : codeBlock.getChildren()) {
+            if(codeBlockChildren instanceof PsiSwitchLabelStatement switchLabelStatement) {
+              for (@NotNull PsiElement labelStatementChild : switchLabelStatement.getChildren()) {
+                if (labelStatementChild instanceof PsiJavaToken javaToken && javaToken.textMatches(":")) {
+                  builder.append("->"); //replace ':' with '->'
+                }else{
+                  builder.append(labelStatementChild.getText());
+                }
+              }
+              PsiElement nextOfSwitchLabelStatement = PsiTreeUtil.skipWhitespacesAndCommentsForward(switchLabelStatement);
+              if (!(nextOfSwitchLabelStatement instanceof PsiBlockStatement) &&
+                  findOneYieldOrThrowStatement(PsiTreeUtil.skipWhitespacesAndCommentsForward(codeBlockChildren)) == null) {
+                builder.append("{"); //wrap multiline rule into '{}'
+              }
+            }
+            else {
+              PsiStatement yieldStatement = findOneYieldOrThrowStatement(codeBlockChildren);
+              if (yieldStatement != null) {
+                for (@NotNull PsiElement yieldStatementChild : yieldStatement.getChildren()) {
+                  if (!(yieldStatementChild instanceof PsiJavaToken javaToken && javaToken.textMatches("yield"))) {
+                    builder.append(yieldStatementChild.getText()); // skip 'yield' for one-line rule
+                  }
+                }
+              }
+              else {
+                builder.append(codeBlockChildren.getText());
+                PsiElement nextOfCodeBlockChildren = PsiTreeUtil.skipWhitespacesAndCommentsForward(codeBlockChildren);
+                if (!(codeBlockChildren instanceof PsiComment || codeBlockChildren instanceof PsiWhiteSpace) &&
+                    !(PsiTreeUtil.skipWhitespacesAndCommentsBackward(codeBlockChildren) instanceof PsiJavaToken) &&
+                    findOneYieldOrThrowStatement(PsiTreeUtil.skipWhitespacesAndCommentsBackward(codeBlockChildren)) == null &&
+                    !(codeBlockChildren instanceof PsiJavaToken) &&
+                    (nextOfCodeBlockChildren instanceof PsiSwitchLabelStatement ||
+                     nextOfCodeBlockChildren instanceof PsiJavaToken javaToken && javaToken.textMatches("}"))) {
+                  builder.append("\n}"); //wrap multiline rule into '{}'
+                }
+              }
+            }
+          }
+        }
+        else {
+          builder.append(switchExpressionChild.getText());
+        }
+      }
+      PsiElementFactory factory = JavaPsiFacade.getElementFactory(element.getProject());
+      PsiExpression newSwitchExpression = factory.createExpressionFromText(builder.toString(), element);
+      switchExpression.replace(newSwitchExpression);
+    }
+
+    @Nullable
+    private static PsiStatement findOneYieldOrThrowStatement(@Nullable PsiElement switchBlockChild) {
+      if ((switchBlockChild instanceof PsiYieldStatement ||
+           switchBlockChild instanceof PsiThrowStatement) &&
+          PsiTreeUtil.skipWhitespacesAndCommentsBackward(switchBlockChild) instanceof PsiSwitchLabelStatement &&
+          (PsiTreeUtil.skipWhitespacesAndCommentsForward(switchBlockChild) instanceof PsiSwitchLabelStatement ||
+           PsiTreeUtil.skipWhitespacesAndCommentsForward(switchBlockChild) instanceof PsiJavaToken javaToken &&
+           javaToken.textMatches("}"))) {
+        return (PsiStatement)switchBlockChild;
+      }
+      return null;
+    }
+
+    @Override
+    public @NotNull String getFamilyName() {
+      return JavaBundle.message("inspection.replace.with.switch.rule.expression.fix.family.name");
+    }
+  }
+
   private static class ReplaceWithSwitchExpressionFix extends PsiUpdateModCommandQuickFix {
     private final ReplacementType myReplacementType;
 

java/java-tests/testData/inspection/switchExpressionMigration/afterOneSwitchColon.java

@@ -0,0 +1,20 @@
+// "Migrate to enhanced switch with rules" "true-preview"
+import org.jetbrains.annotations.Nullable;
+
+class X {
+
+  int test(int y) {
+    return switch (y) {
+        case 1 -> {
+            System.out.println("1");
+            yield 1;
+        }
+        case 2 -> 3;
+        case 3 -> throw new IllegalArgumentException();
+        default -> {
+            System.out.println();
+            yield 5;
+        }
+    };
+  }
+}

java/java-tests/testData/inspection/switchExpressionMigration/afterSwitchColon.java

@@ -0,0 +1,13 @@
+// "Migrate to enhanced switch with rules" "true-preview"
+import org.jetbrains.annotations.Nullable;
+
+class X {
+
+  int test(int y) {
+    return switch (y) { //some comments3
+        case 1 -> 1; //some comments2
+        case 2 -> 3; /*some comments1*/
+        default -> 5; //some comments
+    };
+  }
+}

java/java-tests/testData/inspection/switchExpressionMigration/beforeBrokenSwitchColon.java

@@ -0,0 +1,16 @@
+// "Migrate to enhanced switch with rules" "false"
+import org.jetbrains.annotations.Nullable;
+
+class X {
+
+  int test(int y) {
+    return swi<caret>tch (y) { //some comments3
+      case 1->
+        yield 1; //some comments2
+      case 2:
+        yield 3; /*some comments1*/
+      default:
+        yield 5; //some comments
+    };
+  }
+}

java/java-tests/testData/inspection/switchExpressionMigration/beforeOneSwitchColon.java

@@ -0,0 +1,20 @@
+// "Migrate to enhanced switch with rules" "true-preview"
+import org.jetbrains.annotations.Nullable;
+
+class X {
+
+  int test(int y) {
+    return swi<caret>tch (y) {
+      case 1:
+        System.out.println("1");
+        yield 1;
+      case 2:
+        yield 3;
+      case 3:
+        throw new IllegalArgumentException();
+      default:
+        System.out.println();
+        yield 5;
+    };
+  }
+}

java/java-tests/testData/inspection/switchExpressionMigration/beforeSwitchColon.java

@@ -0,0 +1,16 @@
+// "Migrate to enhanced switch with rules" "true-preview"
+import org.jetbrains.annotations.Nullable;
+
+class X {
+
+  int test(int y) {
+    return swi<caret>tch (y) { //some comments3
+      case 1:
+        yield 1; //some comments2
+      case 2:
+        yield 3; /*some comments1*/
+      default:
+        yield 5; //some comments
+    };
+  }
+}

java/openapi/resources/messages/JavaBundle.properties

@@ -703,6 +703,7 @@ inspection.replace.with.old.style.switch.statement.fix.name=Replace with old sty
 inspection.replace.with.regular.string.literal.fix=Replace with regular string literal
 inspection.replace.with.switch.expression.fix.name=Replace with 'switch' expression
 inspection.replace.with.switch.expression.fix.family.name=Migrate to enhanced switch
+inspection.replace.with.switch.rule.expression.fix.family.name=Migrate to enhanced switch with rules
 inspection.replace.with.text.block.fix=Replace with text block
 inspection.replace.with.string.template.fix=Replace with string template
 inspection.replace.with.string.concatenation.fix=Replace with string concatenation
@@ -747,6 +748,7 @@ inspection.switch.expression.backward.migration.inspection.name=Enhanced 'switch
 inspection.switch.expression.backward.statement.migration.inspection.name='switch' statement can be replaced with old style 'switch' statement
 inspection.switch.expression.migration.inspection.name=Statement can be replaced with enhanced 'switch'
 inspection.switch.expression.migration.inspection.switch.description=Switch statement can be replaced with enhanced 'switch'
+inspection.switch.expression.migration.inspection.switch.expression.description=Switch expression can be replaced with enhanced 'switch' with rules
 inspection.switch.expression.migration.warn.only.on.expression=Show warning only if conversion to expression is possible
 inspection.switch.expression.migration.option.expression.max.statements=Do not report switches having more than {0} {0, choice, 1#statement|2#statements} in a single branch
 inspection.switch.expression.migration.expression.max.statements=Maximum number of statements in one branch to convert to switch expression