From 198903cfd48f10da3a2036266db638f31db57816 Mon Sep 17 00:00:00 2001
From: Liam Miller-Cushon <cushon@google.com>
Date: Tue, 30 Jul 2024 08:16:50 -0700
Subject: [PATCH] Recognize more unnecessary breaks in UnnecessaryBreakInSwitch

PiperOrigin-RevId: 657598108
---
 .../bugpatterns/UnnecessaryBreakInSwitch.java | 52 +++++++++++++------
 .../UnnecessaryBreakInSwitchTest.java         | 51 +++++++++++++++++-
 2 files changed, 87 insertions(+), 16 deletions(-)

diff --git a/core/src/main/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitch.java b/core/src/main/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitch.java
index 8deb44f59fc..0aeba268bf0 100644
--- a/core/src/main/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitch.java
+++ b/core/src/main/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitch.java
@@ -21,6 +21,7 @@
 import static com.google.errorprone.matchers.Description.NO_MATCH;
 import static com.google.errorprone.util.ASTHelpers.isRuleKind;
 
+import com.google.common.collect.ImmutableList;
 import com.google.errorprone.BugPattern;
 import com.google.errorprone.VisitorState;
 import com.google.errorprone.fixes.SuggestedFix;
@@ -29,8 +30,9 @@
 import com.sun.source.tree.BlockTree;
 import com.sun.source.tree.BreakTree;
 import com.sun.source.tree.CaseTree;
-import com.sun.source.tree.StatementTree;
+import com.sun.source.tree.IfTree;
 import com.sun.source.tree.Tree;
+import com.sun.source.util.SimpleTreeVisitor;
 
 /** A {@link BugChecker}; see the associated {@link BugPattern} annotation for details. */
 @BugPattern(
@@ -43,21 +45,41 @@ public Description matchCase(CaseTree tree, VisitorState state) {
       return NO_MATCH;
     }
     Tree body = ASTHelpers.getCaseTreeBody(tree);
-    if (!(body instanceof BlockTree)) {
+    ImmutableList<BreakTree> unnecessaryBreaks = unnecessaryBreaks(body);
+    if (unnecessaryBreaks.isEmpty()) {
       return NO_MATCH;
     }
-    BlockTree blockTree = (BlockTree) body;
-    if (blockTree.getStatements().isEmpty()) {
-      return NO_MATCH;
-    }
-    StatementTree last = getLast(blockTree.getStatements());
-    if (!(last instanceof BreakTree)) {
-      return NO_MATCH;
-    }
-    BreakTree breakTree = (BreakTree) last;
-    if (breakTree.getLabel() != null) {
-      return NO_MATCH;
-    }
-    return describeMatch(last, SuggestedFix.delete(last));
+    unnecessaryBreaks.forEach(
+        unnecessaryBreak ->
+            state.reportMatch(
+                describeMatch(unnecessaryBreak, SuggestedFix.delete(unnecessaryBreak))));
+    return NO_MATCH;
+  }
+
+  private ImmutableList<BreakTree> unnecessaryBreaks(Tree tree) {
+    ImmutableList.Builder<BreakTree> result = ImmutableList.builder();
+    new SimpleTreeVisitor<Void, Void>() {
+      @Override
+      public Void visitBreak(BreakTree node, Void unused) {
+        if (node.getLabel() == null) {
+          result.add(node);
+        }
+        return null;
+      }
+
+      @Override
+      public Void visitBlock(BlockTree node, Void unused) {
+        visit(getLast(node.getStatements(), null), null);
+        return null;
+      }
+
+      @Override
+      public Void visitIf(IfTree node, Void unused) {
+        visit(node.getThenStatement(), null);
+        visit(node.getElseStatement(), null);
+        return null;
+      }
+    }.visit(tree, null);
+    return result.build();
   }
 }
diff --git a/core/src/test/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitchTest.java b/core/src/test/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitchTest.java
index 9e3ed40d0af..49da24c7c46 100644
--- a/core/src/test/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitchTest.java
+++ b/core/src/test/java/com/google/errorprone/bugpatterns/UnnecessaryBreakInSwitchTest.java
@@ -92,7 +92,10 @@ public void negativeNotLast() {
             "  void f(int i) {",
             "    switch (i) {",
             "      default -> {",
-            "        if (true) break;",
+            "        if (true) {",
+            "          break;",
+            "        }",
+            "        System.err.println();",
             "      }",
             "    };",
             "  }",
@@ -120,4 +123,50 @@ public void negativeLabelledBreak() {
             "}")
         .doTest();
   }
+
+  @Test
+  public void negativeLoop() {
+    assumeTrue(RuntimeVersion.isAtLeast14());
+    testHelper
+        .addSourceLines(
+            "Test.java",
+            "class Test {",
+            "  void f(int i) {",
+            "    while (true) {",
+            "      switch (i) {",
+            "        default -> {",
+            "          while (true) {",
+            "            break;",
+            "          }",
+            "        }",
+            "      }",
+            "    }",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void positiveIf() {
+    assumeTrue(RuntimeVersion.isAtLeast14());
+    testHelper
+        .addSourceLines(
+            "Test.java",
+            "class Test {",
+            "  void f(int i) {",
+            "    switch (i) {",
+            "      default -> {",
+            "        if (true) {",
+            "          // BUG: Diagnostic contains: break is unnecessary",
+            "          break;",
+            "        } else {",
+            "          // BUG: Diagnostic contains: break is unnecessary",
+            "          break;",
+            "        }",
+            "      }",
+            "    };",
+            "  }",
+            "}")
+        .doTest();
+  }
 }