Add support for switch expressions in dataflow

SKIP-REQUIRE-JAVADOC

@@ -0,0 +1 @@
+Delete after merge.

checker/tests/nullness/java17/NullnessSwitchExpressions.java

@@ -36,8 +36,6 @@ void method2() {
           default -> throw new IllegalStateException("Invalid day: " + day);
         };
 
-    // TODO: This is a false positive.
-    // :: error: (dereference.of.nullable)
     o.toString();
   }
 
@@ -51,8 +49,6 @@ void method3() {
             String s = null;
             if (day == Day.THURSDAY) {
               s = "hello";
-              // TODO: This is a false positive.
-              // :: error: (dereference.of.nullable)
               s.toString();
             }
             yield s;

checker/tests/nullness/java17/SwitchExpressionInvariant.java

@@ -0,0 +1,27 @@
+// @below-java17-jdk-skip-test
+import java.util.List;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+public class SwitchExpressionInvariant {
+  public static boolean flag = false;
+
+  void method(
+      List<@NonNull String> nonnullStrings, List<@Nullable String> nullableStrings, int fenum) {
+
+    List<@NonNull String> list =
+        // :: error: (assignment)
+        switch (fenum) {
+            // :: error: (switch.expression)
+          case 1 -> nonnullStrings;
+          default -> nullableStrings;
+        };
+
+    List<@Nullable String> list2 =
+        switch (fenum) {
+            // :: error: (switch.expression)
+          case 1 -> nonnullStrings;
+          default -> nullableStrings;
+        };
+  }
+}

dataflow/src/main/java/org/checkerframework/dataflow/cfg/builder/CFGTranslationPhaseOne.java

@@ -46,6 +46,7 @@
 import com.sun.source.tree.SynchronizedTree;
 import com.sun.source.tree.ThrowTree;
 import com.sun.source.tree.Tree;
+import com.sun.source.tree.Tree.Kind;
 import com.sun.source.tree.TryTree;
 import com.sun.source.tree.TypeCastTree;
 import com.sun.source.tree.TypeParameterTree;
@@ -148,6 +149,7 @@
 import org.checkerframework.dataflow.cfg.node.StringConversionNode;
 import org.checkerframework.dataflow.cfg.node.StringLiteralNode;
 import org.checkerframework.dataflow.cfg.node.SuperNode;
+import org.checkerframework.dataflow.cfg.node.SwitchExpressionNode;
 import org.checkerframework.dataflow.cfg.node.SynchronizedNode;
 import org.checkerframework.dataflow.cfg.node.TernaryExpressionNode;
 import org.checkerframework.dataflow.cfg.node.ThisNode;
@@ -250,6 +252,9 @@ public class CFGTranslationPhaseOne extends TreePathScanner<Node, Void> {
   /** Nested scopes of try-catch blocks in force at the current program point. */
   private final TryStack tryStack;
 
+  /** SwitchBuild for the current switch. Used to match yield statements to enclosing switches. */
+  private SwitchBuilder switchBuilder;
+
   /**
    * Maps from AST {@link Tree}s to sets of {@link Node}s. Every Tree that produces a value will
    * have at least one corresponding Node. Trees that undergo conversions, such as boxing or
@@ -475,23 +480,29 @@ public Node scan(Tree tree, Void p) {
       return null;
     }
     // Must use String comparison to support compiling on JDK 11 and earlier.
-    if (tree.getKind().name().equals("SWITCH_EXPRESSION")) {
-      return visitSwitchExpression17(tree, p);
+    //     Features added between JDK 12 and JDK 17 inclusive.
+    switch (tree.getKind().name()) {
+        // case "BINDING_PATTERN":
+        //  return visitBindingPattern17(path.getLeaf(), p);
+      case "SWITCH_EXPRESSION":
+        return visitSwitchExpression17(tree, p);
+      case "YIELD":
+        return visitYield17(tree, p);
+      default:
+        return super.scan(tree, p);
     }
-    return super.scan(tree, p);
-
-    // TODO: Do we need to support yield trees and binding patterns to?
-    // Features added between JDK 12 and JDK 17 inclusive.
-    // switch (tree.getKind().name()) {
-    //   case "BINDING_PATTERN":
-    //     return visitBindingPattern17(path.getLeaf(), p);
-    //   case "SWITCH_EXPRESSION":
-    //     return visitSwitchExpression17(tree, p);
-    //   case "YIELD":
-    //     return visitYield17(path.getLeaf(), p);
-    //   default:
-    //     return super.scan(tree, p);
-    // }
+  }
+  /**
+   * Visit a SwitchExpressionTree
+   *
+   * @param yieldTree a YieldTree, typed as Tree to be backward-compatible
+   * @param p parameter
+   * @return the result of visiting the switch expression tree
+   */
+  public Node visitYield17(Tree yieldTree, Void p) {
+    ExpressionTree resultExpression = TreeUtils.yieldTreeGetValue(yieldTree);
+    switchBuilder.buildSwitchExpressionResult(resultExpression);
+    return null;
   }
 
   /**
@@ -503,10 +514,8 @@ public Node scan(Tree tree, Void p) {
    */
   public Node visitSwitchExpression17(Tree switchExpressionTree, Void p) {
     // TODO: Analyze switch expressions properly.
-    return new MarkerNode(
-        switchExpressionTree,
-        "switch expression tree; not analyzed #" + TreeUtils.treeUids.get(switchExpressionTree),
-        env.getTypeUtils());
+    SwitchBuilder switchBuilder = new SwitchBuilder(switchExpressionTree);
+    return switchBuilder.build();
   }
 
   /* --------------------------------------------------------- */
@@ -2128,28 +2137,63 @@ public Node visitSwitch(SwitchTree tree, Void p) {
    */
   private class SwitchBuilder {
     /** The switch tree. */
-    private final SwitchTree switchTree;
+    private final Tree switchTree;
+
+    /** The case trees of {@code switchTree} */
+    private final List<? extends CaseTree> caseTrees;
+
+    /**
+     * The Tree for the selector expression.
+     *
+     * <pre>
+     *   switch ( <em> selector expression</em> ) { ... }
+     * </pre>
+     */
+    private final ExpressionTree selectorExprTree;
+
     /** The labels for the case bodies. */
     private final Label[] caseBodyLabels;
+
     /**
      * The Node for the assignment of the switch selector expression to a synthetic local variable.
      */
     private AssignmentNode selectorExprAssignment;
 
+    /**
+     * If {@link #switchTree} is a switch expression, then this is the synthetic variable tree that
+     * all results of {@code #switchTree} are assigned. Other, this is null.
+     */
+    private @Nullable VariableTree switchExprVarTree;
+
     /**
      * Construct a SwitchBuilder.
      *
-     * @param tree a switch tree
+     * @param switchTree a switch tree
      */
-    private SwitchBuilder(SwitchTree tree) {
-      this.switchTree = tree;
+    private SwitchBuilder(Tree switchTree) {
+      this.switchTree = switchTree;
+      if (switchTree instanceof SwitchTree) {
+        SwitchTree switchStatementTree = (SwitchTree) switchTree;
+        this.caseTrees = switchStatementTree.getCases();
+        this.selectorExprTree = switchStatementTree.getExpression();
+      } else {
+        this.caseTrees = TreeUtils.switchExpressionTreeGetCases(switchTree);
+        this.selectorExprTree = TreeUtils.switchExpressionTreeGetExpression(switchTree);
+      }
       // "+ 1" for the default case.  If the switch has an explicit default case, then
       // the last element of the array is never used.
-      this.caseBodyLabels = new Label[switchTree.getCases().size() + 1];
+      this.caseBodyLabels = new Label[caseTrees.size() + 1];
     }
 
-    /** Build up the CFG for the switchTree. */
-    public void build() {
+    /**
+     * Build up the CFG for the switchTree.
+     *
+     * @return if the switch is a switch expression, then a {@link SwitchExpressionNode}; otherwise,
+     *     null
+     */
+    public @Nullable SwitchExpressionNode build() {
+      SwitchBuilder oldSwitchBuilder = switchBuilder;
+      switchBuilder = this;
       TryFinallyScopeCell oldBreakTargetL = breakTargetL;
       breakTargetL = new TryFinallyScopeCell(new Label());
       int cases = caseBodyLabels.length - 1;
@@ -2160,16 +2204,20 @@ public void build() {
 
       buildSelector();
 
-      // Build CFG for the cases.
-      extendWithNode(
-          new MarkerNode(
-              switchTree,
-              "start of switch statement #" + TreeUtils.treeUids.get(switchTree),
-              env.getTypeUtils()));
+      buildSwitchExpressionVar();
+
+      if (switchTree.getKind() == Kind.SWITCH) {
+        extendWithNode(
+            new MarkerNode(
+                switchTree,
+                "start of switch statement #" + TreeUtils.treeUids.get(switchTree),
+                env.getTypeUtils()));
+      }
 
+      // Build CFG for the cases.
       Integer defaultIndex = null;
       for (int i = 0; i < cases; ++i) {
-        CaseTree caseTree = switchTree.getCases().get(i);
+        CaseTree caseTree = caseTrees.get(i);
         if (TreeUtils.caseTreeGetExpressions(caseTree).isEmpty()) {
           defaultIndex = i;
         } else {
@@ -2180,17 +2228,35 @@ public void build() {
         // The checks of all cases must happen before the default case, therefore we build the
         // default case last.
         // Fallthrough is still handled correctly with the caseBodyLabels.
-        buildCase(switchTree.getCases().get(defaultIndex), defaultIndex);
+        buildCase(caseTrees.get(defaultIndex), defaultIndex);
       }
 
       addLabelForNextNode(breakTargetL.peekLabel());
       breakTargetL = oldBreakTargetL;
+      if (switchTree.getKind() == Kind.SWITCH) {
+        extendWithNode(
+            new MarkerNode(
+                switchTree,
+                "end of switch statement #" + TreeUtils.treeUids.get(switchTree),
+                env.getTypeUtils()));
+      }
 
-      extendWithNode(
-          new MarkerNode(
-              switchTree,
-              "end of switch statement #" + TreeUtils.treeUids.get(switchTree),
-              env.getTypeUtils()));
+      switchBuilder = oldSwitchBuilder;
+      if (switchExprVarTree != null) {
+        IdentifierTree switchExprVarUseTree = treeBuilder.buildVariableUse(switchExprVarTree);
+        handleArtificialTree(switchExprVarUseTree);
+
+        LocalVariableNode switchExprVarUseNode = new LocalVariableNode(switchExprVarUseTree);
+        switchExprVarUseNode.setInSource(false);
+        extendWithNode(switchExprVarUseNode);
+        SwitchExpressionNode switchExpressionNode =
+            new SwitchExpressionNode(
+                TreeUtils.typeOf(switchTree), switchTree, switchExprVarUseNode);
+        extendWithNode(switchExpressionNode);
+        return switchExpressionNode;
+      } else {
+        return null;
+      }
     }
 
     /**
@@ -2201,7 +2267,7 @@ public void build() {
      */
     private void buildSelector() {
       // Create a synthetic variable to which the switch selector expression will be assigned
-      TypeMirror selectorExprType = TreeUtils.typeOf(switchTree.getExpression());
+      TypeMirror selectorExprType = TreeUtils.typeOf(selectorExprTree);
       VariableTree selectorVarTree =
           treeBuilder.buildVariableDecl(selectorExprType, uniqueName("switch"), findOwner(), null);
       handleArtificialTree(selectorVarTree);
@@ -2217,17 +2283,36 @@ private void buildSelector() {
       selectorVarUseNode.setInSource(false);
       extendWithNode(selectorVarUseNode);
 
-      Node selectorExprNode = unbox(scan(switchTree.getExpression(), null));
+      Node selectorExprNode = unbox(scan(selectorExprTree, null));
 
-      AssignmentTree assign =
-          treeBuilder.buildAssignment(selectorVarUseTree, switchTree.getExpression());
+      AssignmentTree assign = treeBuilder.buildAssignment(selectorVarUseTree, selectorExprTree);
       handleArtificialTree(assign);
 
       selectorExprAssignment = new AssignmentNode(assign, selectorVarUseNode, selectorExprNode);
       selectorExprAssignment.setInSource(false);
       extendWithNode(selectorExprAssignment);
     }
 
+    /**
+     * If {@link #switchTree} is a switch expression tree, this method creates a synthetic variable
+     * whose value is the value of the switch expression.
+     */
+    private void buildSwitchExpressionVar() {
+      if (switchTree.getKind() == Kind.SWITCH) {
+        // A switch statement does not have a value, so do nothing.
+        return;
+      }
+      TypeMirror switchExprType = TreeUtils.typeOf(switchTree);
+      switchExprVarTree =
+          treeBuilder.buildVariableDecl(
+              switchExprType, uniqueName("switchExpr"), findOwner(), null);
+      handleArtificialTree(switchExprVarTree);
+
+      VariableDeclarationNode switchExprVarNode = new VariableDeclarationNode(switchExprVarTree);
+      switchExprVarNode.setInSource(false);
+      extendWithNode(switchExprVarNode);
+    }
+
     /**
      * Build the CGF for the case tree, {@code tree}.
      *
@@ -2256,11 +2341,51 @@ private void buildCase(CaseTree tree, int index) {
           scan(stmt, null);
         }
       } else {
-        scan(TreeUtils.caseTreeGetBody(tree), null);
+        Tree bodyTree = TreeUtils.caseTreeGetBody(tree);
+        if (bodyTree instanceof ExpressionTree) {
+          buildSwitchExpressionResult((ExpressionTree) bodyTree);
+        } else {
+          scan(bodyTree, null);
+        }
       }
       extendWithExtendedNode(new UnconditionalJump(nextBodyL));
       addLabelForNextNode(nextCaseL);
     }
+
+    /**
+     * Does the following for each result expression of a switch expression:
+     *
+     * <ol>
+     *   <li>Builds the CFG for the switch expression result.
+     *   <li>Creates an assignment node for the assignment of {@code resultExpression} to {@code
+     *       switchExprVarTree}.
+     *   <li>Adds an unconditional jump to {@link #breakTargetL} (the end of the switch expression.
+     * </ol>
+     *
+     * @param resultExpression the result of a switch expression; either from a yield or an
+     *     expression in a case rule.
+     */
+    void buildSwitchExpressionResult(ExpressionTree resultExpression) {
+      IdentifierTree switchExprVarUseTree = treeBuilder.buildVariableUse(switchExprVarTree);
+      handleArtificialTree(switchExprVarUseTree);
+
+      LocalVariableNode switchExprVarUseNode = new LocalVariableNode(switchExprVarUseTree);
+      switchExprVarUseNode.setInSource(false);
+      extendWithNode(switchExprVarUseNode);
+
+      Node resultExprNode = scan(resultExpression, null);
+
+      AssignmentTree assign = treeBuilder.buildAssignment(switchExprVarUseTree, resultExpression);
+      handleArtificialTree(assign);
+
+      AssignmentNode assignmentNode =
+          new AssignmentNode(assign, switchExprVarUseNode, resultExprNode);
+      assignmentNode.setInSource(false);
+      extendWithNode(assignmentNode);
+
+      assert breakTargetL != null : "no target for yield statement";
+      extendWithExtendedNode(new UnconditionalJump(breakTargetL.accessLabel()));
+    }
   }
 
   @Override

dataflow/src/main/java/org/checkerframework/dataflow/cfg/node/AbstractNodeVisitor.java

@@ -216,6 +216,11 @@ public R visitTernaryExpression(TernaryExpressionNode n, P p) {
     return visitNode(n, p);
   }
 
+  @Override
+  public R visitSwitchExpressionNode(SwitchExpressionNode n, P p) {
+    return visitNode(n, p);
+  }
+
   @Override
   public R visitAssignment(AssignmentNode n, P p) {
     return visitNode(n, p);

dataflow/src/main/java/org/checkerframework/dataflow/cfg/node/NodeVisitor.java

@@ -91,6 +91,8 @@ public interface NodeVisitor<R, P> {
 
   R visitTernaryExpression(TernaryExpressionNode n, P p);
 
+  R visitSwitchExpressionNode(SwitchExpressionNode n, P p);
+
   R visitAssignment(AssignmentNode n, P p);
 
   R visitLocalVariable(LocalVariableNode n, P p);