Handle the general case of switch statement and switch expression.

Switch expressions are rewritten as switch statements in a expression block, e.g:

```
  a + switch (someEnum) {
     case VALUE1:
     case VALUE2:
       yield 2;
     case LASTVALUE:
       yield 3;
  }
```

is rewritten to:
```
  a + ^{ switch (someEnum) {
     case VALUE1:
     case VALUE2:
       return 2;
     case LASTVALUE:
       return 3;
     default:
       __builtin_unreachable();
  }}()
```

and switches with patterns or guards are rewritten into if nests that set a selector variable and the switch is converted to use that variable, e.g:

```
  switch (someEnum) {
     case String s -> s.length();
     case Integer i -> i;
     default -> 0;
  }
```
is translated as if it was written as:
```
  int selector = 0;
  if (someEnum instanceof String s) {
   selector = 1;
  } else if (someEnum instanceof Integer i) {
   selector = 2;
  }
  switch (selector) {
     case 1 -> s.length();
     case 2 -> i;
     default -> 0;
  }
```

Note that the pattern variables will be extracted out and will have the proper value in the places where they can be accessed.

PiperOrigin-RevId: 828453235

Author: rluble

translator/src/main/java/com/google/devtools/j2objc/gen/StatementGenerator.java

@@ -99,6 +99,7 @@
 import com.google.devtools.j2objc.ast.VariableDeclarationFragment;
 import com.google.devtools.j2objc.ast.VariableDeclarationStatement;
 import com.google.devtools.j2objc.ast.WhileStatement;
+import com.google.devtools.j2objc.ast.YieldStatement;
 import com.google.devtools.j2objc.util.ElementUtil;
 import com.google.devtools.j2objc.util.NameTable;
 import com.google.devtools.j2objc.util.TypeUtil;
@@ -785,20 +786,8 @@ public boolean visit(SwitchCase node) {
   @Override
   @SuppressWarnings("UngroupedOverloads")
   public boolean visit(SwitchExpression node) {
-    Expression expr = node.getExpression();
-    buffer.append("switch (");
-    expr.accept(this);
-    buffer.append(") ");
-    buffer.append("{\n");
-    buffer.indent();
-    for (Statement stmt : node.getStatements()) {
-      buffer.printIndent();
-      stmt.accept(this);
-    }
-    buffer.unindent();
-    buffer.printIndent();
-    buffer.append("}");
-    return false;
+    // Switch expressions are rewritten to switch statements by SwitchConstructRewriter.
+    throw new AssertionError("switch expression not converted");
   }
 
   @Override
@@ -997,6 +986,12 @@ public boolean visit(WhileStatement node) {
     return false;
   }
 
+  @Override
+  public boolean visit(YieldStatement node) {
+    // Yield statements are rewritten to return statements in SwitchConstructRewriter.
+    throw new AssertionError("yield statement not converted");
+  }
+
   @Override
   public boolean visit(Initializer node) {
     // All Initializer nodes should have been converted during initialization

translator/src/main/java/com/google/devtools/j2objc/javac/TreeConverter.java

@@ -14,6 +14,8 @@
 
 package com.google.devtools.j2objc.javac;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import com.google.common.collect.Lists;
 import com.google.devtools.j2objc.Options;
 import com.google.devtools.j2objc.ast.AbstractTypeDeclaration;
@@ -94,6 +96,7 @@
 import com.google.devtools.j2objc.ast.SuperMethodInvocation;
 import com.google.devtools.j2objc.ast.SuperMethodReference;
 import com.google.devtools.j2objc.ast.SwitchCase;
+import com.google.devtools.j2objc.ast.SwitchConstruct;
 import com.google.devtools.j2objc.ast.SwitchExpression;
 import com.google.devtools.j2objc.ast.SwitchStatement;
 import com.google.devtools.j2objc.ast.SynchronizedStatement;
@@ -206,7 +209,6 @@
 import com.sun.tools.javac.util.Position;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 import javax.lang.model.element.AnnotationMirror;
@@ -318,8 +320,6 @@ private TreeNode convertInner(Tree javacNode, TreePath parent) {
         return convertBlock((BlockTree) javacNode, parent);
       case "BREAK":
         return convertBreakStatement((BreakTree) javacNode);
-      case "CASE":
-        return convertCase((CaseTree) javacNode, parent);
       case "CATCH":
         return convertCatch((CatchTree) javacNode, parent);
       case "CLASS":
@@ -675,19 +675,6 @@ private TreeNode convertBreakStatement(BreakTree node) {
     return newNode;
   }
 
-  private TreeNode convertCase(CaseTree node, TreePath parent) {
-    // Case statements are converted in convertSwitch().
-    // TODO(b/456297342): Implement handling for patterns in case statements.
-    SwitchCase newNode = new SwitchCase();
-    List<? extends ExpressionTree> expressions = node.getExpressions();
-    for (ExpressionTree expressionTree : expressions) {
-      TreePath path = getTreePath(parent, node);
-      newNode.addExpression((Expression) convert(expressionTree, path));
-    }
-    newNode.setIsDefault(expressions.isEmpty());
-    return newNode;
-  }
-
   private TreeNode convertCatch(CatchTree node, TreePath parent) {
     TreePath path = getTreePath(parent, node);
     return new CatchClause()
@@ -1371,12 +1358,6 @@ private TreeNode convertRecord(ClassTree node, TreePath parent) {
 
   private TreeNode convertReturn(ReturnTree node, TreePath parent) {
     Expression expr = (Expression) convert(node.getExpression(), getTreePath(parent, node));
-    if (expr != null && expr.getKind() == TreeNode.Kind.SWITCH_EXPRESSION) {
-      // This is returning just a switch expression and the construction transforms yields to
-      // returns in this case so it is safe to just return the switch.
-      // TODO(b/456285695): Move this to a pass.
-      return new ExpressionStatement(expr);
-    }
     return new ReturnStatement(expr);
   }
 
@@ -1386,139 +1367,59 @@ private TreeNode convertStringLiteral(LiteralTree node, TreePath parent) {
 
   private TreeNode convertSwitch(SwitchTree node, TreePath parent) {
     TreePath path = getTreePath(parent, node);
-    SwitchStatement newNode =
+    SwitchStatement switchStatement =
         new SwitchStatement().setExpression(convertWithoutParens(node.getExpression(), path));
-    for (CaseTree switchCase : node.getCases()) {
-      TreePath switchCasePath = getTreePath(path, switchCase);
-      if (switchCase.getCaseKind() == CaseKind.RULE) {
-        SwitchCase switchCaseStmt = convertCaseRule(switchCase, switchCasePath);
-        newNode.addStatement(switchCaseStmt);
-        newNode.addStatement(new BreakStatement());
-      } else {
-        newNode.addStatement((SwitchCase) convert(switchCase, path));
-        for (StatementTree s : switchCase.getStatements()) {
-          newNode.addStatement((Statement) convert(s, switchCasePath));
-        }
-      }
+    for (CaseTree caseTree : node.getCases()) {
+      TreePath switchCasePath = getTreePath(parent, caseTree);
+      SwitchCase switchCase = convertSwitchCaseLabel(caseTree, switchCasePath);
+      switchStatement.addStatement(switchCase);
+      convertSwitchCaseBody(caseTree, switchCasePath, switchStatement, switchCase);
     }
-    return newNode;
+
+    return switchStatement;
   }
 
   private SwitchExpression convertSwitchExpression(SwitchExpressionTree node, TreePath parent) {
     TreePath path = getTreePath(parent, node);
-    SwitchExpression newNode =
-        new SwitchExpression().setExpression(convertWithoutParens(node.getExpression(), path));
-    newNode.setTypeMirror(getTypeMirror(path));
-
-    Tree parentTree = parent.getLeaf();
-    // TODO(b/456285695): Move the optimization of the two cases to a pass. And implement the
-    // general case, possibly using block expressions.
-    boolean exprReturned = parentTree.getKind() == Tree.Kind.RETURN;
-    boolean exprSaved = parentTree.getKind() == Tree.Kind.VARIABLE;
-    VariableElement yieldSymbol = null;
-    if (parentTree instanceof JCVariableDecl varDecl) {
-      yieldSymbol = varDecl.sym;
-    }
-
-    List<? extends CaseTree> cases = node.getCases();
-    for (CaseTree caseTree : cases) {
-      TreePath switchCasePath = getTreePath(path, caseTree);
-
-      SwitchCase switchCase = convertCaseRule(caseTree, switchCasePath);
-
-      // Convert any cases that have multiple expressions into a list of
-      // SwitchCases for all but the last expression, remove them from
-      // the SwitchExpressionCase and insert their SwitchCase equivalent
-      // ahead of the current case.
-      List<Expression> caseExprs = switchCase.getExpressions();
-      if (caseExprs.size() > 1) {
-        Iterator<Expression> caseExprIter = caseExprs.iterator();
-        Expression caseExpr = caseExprIter.next();
-        do {
-          caseExprIter.remove();
-          SwitchCase newCase = new SwitchCase();
-          newCase.addExpression(caseExpr.copy());
-          newNode.addStatement(newCase);
-          caseExpr = caseExprIter.next();
-        } while (caseExprIter.hasNext());
-      }
 
-      TreeNode body = switchCase.getBody();
-      if (body.getKind() == TreeNode.Kind.YIELD_STATEMENT) {
-        Expression yield = ((YieldStatement) body).getExpression().copy();
-        if (exprReturned) {
-          switchCase.setBody(new ReturnStatement(yield));
-        } else if (exprSaved) {
-          Assignment assignment = new Assignment(new SimpleName(yieldSymbol), yield);
-          switchCase.setBody(new ExpressionStatement(assignment));
-        }
-      } else if (body.getKind() == TreeNode.Kind.BLOCK) {
-        List<Statement> blockStmts = ((Block) body).getStatements();
-        if (!blockStmts.isEmpty()) {
-          int lastIndex = blockStmts.size() - 1;
-          Statement lastStmt = blockStmts.get(lastIndex);
-          if (lastStmt.getKind() == TreeNode.Kind.YIELD_STATEMENT) {
-            Expression yield = ((YieldStatement) lastStmt).getExpression().copy();
-            if (exprReturned) {
-              switchCase.setBody(new ReturnStatement(yield));
-            } else if (exprSaved) {
-              Assignment assignment = new Assignment(new SimpleName(yieldSymbol), yield);
-              switchCase.setBody(new ExpressionStatement(assignment));
-            }
-          }
-        }
-      }
-      newNode.addStatement(switchCase);
-      if (exprSaved) {
-        newNode.addStatement(new BreakStatement());
-      }
-    }
+    SwitchExpression switchExpression =
+        new SwitchExpression()
+            .setExpression(convertWithoutParens(node.getExpression(), path))
+            .setTypeMirror(getTypeMirror(path));
 
-    if (!newNode.hasDefaultCase()) {
-      addUnreachableDirective(newNode);
+    for (CaseTree caseTree : node.getCases()) {
+      TreePath switchCasePath = getTreePath(parent, caseTree);
+      SwitchCase switchCase = convertSwitchCaseLabel(caseTree, switchCasePath);
+      switchExpression.addStatement(switchCase);
+      convertSwitchCaseBody(caseTree, switchCasePath, switchExpression, switchCase);
     }
 
-    return newNode;
-  }
-
-  // If a switch expression doesn't have a default case,  add a function call to tell clang that all
-  // cases are handled. This can be asserted because switch expressions are checked to be
-  // exhaustive.
-  private void addUnreachableDirective(SwitchExpression switchExpression) {
-    // If the switch expression has no default case, add one
-    SwitchCase defaultCase = new SwitchCase().setIsDefault(true);
-    switchExpression.addStatement(defaultCase);
-
-    TypeMirror voidType = newUnit.getEnv().typeUtil().getVoid();
-    FunctionElement element = new FunctionElement("__builtin_unreachable", voidType, null);
-    FunctionInvocation unreachableInvocation = new FunctionInvocation(element, voidType);
-
-    if (switchExpression.getStatements().isEmpty()
-        || ((SwitchCase) switchExpression.getStatements().get(0)).getBody() != null) {
-      // If the switch expression has rules, add it as a rule.
-      defaultCase.setBody(new ExpressionStatement(unreachableInvocation));
-    } else {
-      // Otherwise, add it as a statement.
-      switchExpression.addStatement(new ExpressionStatement(unreachableInvocation));
+    if (!switchExpression.hasDefaultCase()) {
+      // Switch expressions are exhaustive, so if there is no default case, add one to signal the
+      // compiler that the default case is unreachable.
+      addUnreachableDirective(switchExpression);
     }
+
+    return switchExpression;
   }
 
-  private SwitchCase convertCaseRule(CaseTree caseTree, TreePath parent) {
+  private SwitchCase convertSwitchCaseLabel(CaseTree caseTree, TreePath parent) {
+    TreePath switchCasePath = getTreePath(parent, caseTree);
     SwitchCase switchCase = new SwitchCase();
     List<? extends CaseLabelTree> caseExpressionsList = caseTree.getLabels();
     for (CaseLabelTree caseLabelTree : caseExpressionsList) {
       switch (caseLabelTree.getKind()) {
         case CONSTANT_CASE_LABEL -> {
           ConstantCaseLabelTree constantCaseLabelTree = (ConstantCaseLabelTree) caseLabelTree;
           switchCase.addExpression(
-              (Expression) convert(constantCaseLabelTree.getConstantExpression(), parent));
+              (Expression) convert(constantCaseLabelTree.getConstantExpression(), switchCasePath));
         }
         case PATTERN_CASE_LABEL -> {
           PatternCaseLabelTree patternCaseLabelTree = (PatternCaseLabelTree) caseLabelTree;
           if (patternCaseLabelTree.getPattern() instanceof BindingPatternTree bindingPatternTree) {
             VariableTree varTree = (VariableTree) bindingPatternTree.getVariable();
             VariableElement var =
-                ((VariableDeclaration) convertVariableDeclaration(varTree, parent))
+                ((VariableDeclaration) convertVariableDeclaration(varTree, switchCasePath))
                     .getVariableElement();
             Pattern.BindingPattern pattern = new Pattern.BindingPattern(var);
             switchCase.setPattern(pattern);
@@ -1528,25 +1429,72 @@ private SwitchCase convertCaseRule(CaseTree caseTree, TreePath parent) {
         default -> throw new AssertionError("unknown case label type: " + caseLabelTree.getKind());
       }
     }
-    Tree javacBody = caseTree.getBody();
-    if (javacBody != null) {
-      TreeNode body = convert(javacBody, parent);
-      if (body instanceof Expression) {
-        body = new YieldStatement((Expression) body);
+    switchCase.setGuard((Expression) convert(caseTree.getGuard(), switchCasePath));
+
+    return switchCase;
+  }
+
+  private void convertSwitchCaseBody(
+      CaseTree caseTree, TreePath parent, SwitchConstruct switchConstruct, SwitchCase switchCase) {
+    TreePath switchCasePath = getTreePath(parent, caseTree);
+    if (caseTree.getCaseKind() == CaseKind.RULE) {
+      if (caseTree.getBody() != null) {
+        boolean isImplicitYield =
+            switchConstruct instanceof SwitchExpression switchExpression
+                && switchExpression.getTypeMirror().getKind() != TypeKind.VOID;
+
+        TreeNode body = convert(caseTree.getBody(), parent);
+        if (body instanceof Expression expression) {
+          if (isImplicitYield) {
+            // Switch expression that are not void have an implicit yield.
+            body = new YieldStatement(expression);
+          } else {
+            body = new ExpressionStatement(expression);
+          }
+        }
+        if (!isImplicitYield) {
+          // Add a break statement to the switch statement with rules to avoid fallback.
+          Block block = body instanceof Block b ? b : new Block().addStatement((Statement) body);
+          block.addStatement(new BreakStatement());
+          body = block;
+        }
+        switchCase.setBody((Statement) body);
+
+      } else {
+        Block body = new Block();
+        for (StatementTree s : caseTree.getStatements()) {
+          body.addStatement((Statement) convert(s, switchCasePath));
+        }
+        switchCase.setBody(body);
       }
-      switchCase.setBody((Statement) body);
     } else {
-      Block body = new Block();
-      List<? extends StatementTree> statementTrees = caseTree.getStatements();
-      for (StatementTree statementTree : statementTrees) {
-        body.addStatement((Statement) convert(statementTree, parent));
+      checkArgument(caseTree.getBody() == null);
+      for (StatementTree s : caseTree.getStatements()) {
+        switchConstruct.addStatement((Statement) convert(s, switchCasePath));
       }
-      switchCase.setBody(body);
     }
+  }
+
+  // If a switch expression doesn't have a default case,  add a function call to tell clang that all
+  // cases are handled. This can be asserted because switch expressions are checked to be
+  // exhaustive.
+  private void addUnreachableDirective(SwitchExpression switchExpression) {
+    // If the switch expression has no default case, add one
+    SwitchCase defaultCase = new SwitchCase().setIsDefault(true);
+    switchExpression.addStatement(defaultCase);
 
-    switchCase.setGuard((Expression) convert(caseTree.getGuard(), parent));
+    TypeMirror voidType = newUnit.getEnv().typeUtil().getVoid();
+    FunctionElement element = new FunctionElement("__builtin_unreachable", voidType, null);
+    FunctionInvocation unreachableInvocation = new FunctionInvocation(element, voidType);
 
-    return switchCase;
+    if (switchExpression.getStatements().isEmpty()
+        || ((SwitchCase) switchExpression.getStatements().get(0)).getBody() != null) {
+      // If the switch expression has rules, add it as a rule.
+      defaultCase.setBody(new ExpressionStatement(unreachableInvocation));
+    } else {
+      // Otherwise, add it as a statement.
+      switchExpression.addStatement(new ExpressionStatement(unreachableInvocation));
+    }
   }
 
   private TreeNode convertSynchronized(SynchronizedTree node, TreePath parent) {

translator/src/main/java/com/google/devtools/j2objc/pipeline/TranslationProcessor.java

@@ -64,7 +64,7 @@
 import com.google.devtools.j2objc.translate.SerializationStripper;
 import com.google.devtools.j2objc.translate.StaticVarRewriter;
 import com.google.devtools.j2objc.translate.SuperMethodInvocationRewriter;
-import com.google.devtools.j2objc.translate.SwitchCaseRewriter;
+import com.google.devtools.j2objc.translate.SwitchConstructRewriter;
 import com.google.devtools.j2objc.translate.SwitchRewriter;
 import com.google.devtools.j2objc.translate.UnsequencedExpressionRewriter;
 import com.google.devtools.j2objc.translate.VarargsRewriter;
@@ -211,6 +211,10 @@ public static void applyMutations(
     new AbstractMethodRewriter(unit, deadCodeMap).run();
     ticker.tick("AbstractMethodRewriter");
 
+    // Before: InstanceOfPatternRewriter - to handle instanceof patterns from the switch rewrite.
+    new SwitchConstructRewriter(unit).run();
+    ticker.tick("SwitchConstructRewriter");
+
     // Before: VariableRenamer - VariableRenamer renames variables in a scope-aware manner.
     new InstanceOfPatternRewriter(unit).run();
     ticker.tick("InstanceOfPatternRewriter");
@@ -342,11 +346,6 @@ public static void applyMutations(
     new CastResolver(unit).run();
     ticker.tick("CastResolver");
 
-    // After: SwitchRewriter to handle Java 21 patterns and guards in switches,
-    // and CastResolver to avoid duplicate cast checks.
-    new SwitchCaseRewriter(unit).run();
-    ticker.tick("SwitchCaseRewriter");
-
     // After: InnerClassExtractor, Functionizer - Expects all types to be
     //   top-level and functionizing to have occurred.
     new PrivateDeclarationResolver(unit).run();