[IfChainToSwitch] Refine analysis of breaks and yields to enable additional conversions to switch

core/src/main/java/com/google/errorprone/bugpatterns/IfChainToSwitch.java

@@ -22,6 +22,7 @@
 import static com.google.errorprone.BugPattern.SeverityLevel.WARNING;
 import static com.google.errorprone.bugpatterns.SwitchUtils.COMPILE_TIME_CONSTANT_MATCHER;
 import static com.google.errorprone.bugpatterns.SwitchUtils.getReferencedLocalVariablesInTree;
+import static com.google.errorprone.bugpatterns.SwitchUtils.hasBreakOutOfTree;
 import static com.google.errorprone.bugpatterns.SwitchUtils.isEnumValue;
 import static com.google.errorprone.bugpatterns.SwitchUtils.renderComments;
 import static com.google.errorprone.matchers.Description.NO_MATCH;
@@ -64,20 +65,25 @@
 import com.sun.source.tree.InstanceOfTree;
 import com.sun.source.tree.LiteralTree;
 import com.sun.source.tree.StatementTree;
+import com.sun.source.tree.SwitchExpressionTree;
 import com.sun.source.tree.Tree;
 import com.sun.source.tree.Tree.Kind;
 import com.sun.source.tree.VariableTree;
 import com.sun.source.tree.YieldTree;
 import com.sun.source.util.TreePath;
+import com.sun.source.util.TreePathScanner;
 import com.sun.source.util.TreeScanner;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Symbol.TypeVariableSymbol;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.TypeTag;
 import com.sun.tools.javac.code.Types;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -1547,9 +1553,16 @@ private List<SuggestedFix> deepAnalysisOfIfChain(
       VisitorState state,
       Range<Integer> ifTreeSourceRange) {
 
-    // Wrapping break/yield in a switch can potentially change its semantics.  A deeper analysis of
-    // whether semantics are preserved is not attempted here
-    if (hasBreakOrYieldInTree(ifTree)) {
+    // Yields and breaks within if-chain's blocks are allowable, provided that they do not transfer
+    // control outside of their respective block
+    boolean hasBreakOut =
+        cases.stream()
+            .filter(caseIr -> caseIr.arrowRhsOptional().isPresent())
+            .map(caseIr -> caseIr.arrowRhsOptional().get())
+            .anyMatch(arrowRhs -> hasBreakOutOfTree(arrowRhs, state));
+    boolean hasYieldOut =
+        analyzeYieldControlFlow(ifTree, state).values().stream().anyMatch(value -> ifTree == value);
+    if (hasBreakOut || hasYieldOut) {
       return new ArrayList<>();
     }
 
@@ -1635,6 +1648,35 @@ private List<SuggestedFix> deepAnalysisOfIfChain(
     return suggestedFixes;
   }
 
+  /**
+   * Returns a map from yield trees to the scope that they are yielding to. Special case: when
+   * yielding to the scope of somewhere above {@code tree}, {@code tree} itself is used as a
+   * sentinel value (no attempt is made to search up the AST from {@code tree} for the actual
+   * scope).
+   */
+  private static Map<Tree, Tree> analyzeYieldControlFlow(Tree tree, VisitorState state) {
+    ArrayDeque<Tree> yieldScope = new ArrayDeque<>();
+    yieldScope.push(tree);
+    HashMap<Tree, Tree> result = new HashMap<>();
+    // One can only yield from a switch expression
+    new TreePathScanner<Void, Void>() {
+      @Override
+      public Void visitSwitchExpression(SwitchExpressionTree switchExpressionTree, Void unused) {
+        yieldScope.push(switchExpressionTree);
+        super.visitSwitchExpression(switchExpressionTree, null);
+        yieldScope.pop();
+        return null;
+      }
+
+      @Override
+      public Void visitYield(YieldTree yieldTree, Void unused) {
+        result.put(yieldTree, yieldScope.peek());
+        return null;
+      }
+    }.scan(state.getPath(), null);
+    return result;
+  }
+
   /**
    * If a finding is available, build a {@code SuggestedFix} for it and add to the suggested fixes.
    */

core/src/main/java/com/google/errorprone/bugpatterns/SwitchUtils.java

@@ -36,6 +36,7 @@
 import com.google.errorprone.util.ErrorProneComment;
 import com.sun.source.tree.AssignmentTree;
 import com.sun.source.tree.BlockTree;
+import com.sun.source.tree.BreakTree;
 import com.sun.source.tree.CaseTree;
 import com.sun.source.tree.ContinueTree;
 import com.sun.source.tree.ExpressionTree;
@@ -322,6 +323,34 @@ public Boolean reduce(@Nullable Boolean left, @Nullable Boolean right) {
     return result != null && result;
   }
 
+  /**
+   * Determines whether any `break` statements that jump out of the {@code tree} are present within
+   * the {@code tree}.
+   */
+  static boolean hasBreakOutOfTree(Tree tree, VisitorState state) {
+    Boolean result =
+        new TreeScanner<Boolean, Void>() {
+          @Override
+          public Boolean visitBreak(BreakTree breakTree, Void unused) {
+            Tree breakTarget = skipLabel(requireNonNull(((JCTree.JCBreak) breakTree).target));
+            // If the break transfers control to somewhere above the tree, it is jumping out.
+            var pathIterator = state.getPath().iterator();
+            for (Tree at = tree; pathIterator.hasNext(); at = pathIterator.next()) {
+              if (at == breakTarget) {
+                return true;
+              }
+            }
+            return false;
+          }
+
+          @Override
+          public Boolean reduce(@Nullable Boolean left, @Nullable Boolean right) {
+            return Objects.equals(left, true) || Objects.equals(right, true);
+          }
+        }.scan(tree, null);
+    return result != null && result;
+  }
+
   /**
    * In a switch with multiple assignments, determine whether a subsequent assignment target is
    * compatible with the first assignment target.