Handle JDK 21 case operands in type refinement (#928)

Fixes #927. We also extract some common logic in
`AccessPathNullnessPropagation` around handling equality comparisons to
avoid duplication.

2 files changed, 80 insertions, 38 deletions

diff --git a/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/NullAwaySwitchTests.java b/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/NullAwaySwitchTests.java
index ca349f5..ab5ef6a 100644
--- a/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/NullAwaySwitchTests.java
+++ b/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/NullAwaySwitchTests.java
@@ -206,4 +206,42 @@ public class NullAwaySwitchTests {
             "}")
         .doTest();
   }
+
+  @Test
+  public void testSwitchExprNullCaseDataflow() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "SwitchNullCase.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "class SwitchNullCase {",
+            "  public enum NullableEnum {",
+            "    A,",
+            "    B,",
+            "  }",
+            "  static Object testPositive1(@Nullable NullableEnum nullableEnum) {",
+            "    return switch (nullableEnum) {",
+            "      case A -> new Object();",
+            "      // BUG: Diagnostic contains: dereferenced expression nullableEnum is @Nullable",
+            "      case null -> nullableEnum.hashCode();",
+            "      default -> nullableEnum.toString();",
+            "    };",
+            "  }",
+            "  static Object testPositive2(@Nullable NullableEnum nullableEnum) {",
+            "    return switch (nullableEnum) {",
+            "      case A -> new Object();",
+            "      // BUG: Diagnostic contains: dereferenced expression nullableEnum is @Nullable",
+            "      case null, default -> nullableEnum.toString();",
+            "    };",
+            "  }",
+            "  static Object testNegative(@Nullable NullableEnum nullableEnum) {",
+            "    return switch (nullableEnum) {",
+            "      case A, B -> nullableEnum.hashCode();",
+            "      case null -> new Object();",
+            "      default -> nullableEnum.toString();",
+            "    };",
+            "  }",
+            "}")
+        .doTest();
+  }
 }
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
index 2996450..a1efbe5 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
@@ -405,54 +405,43 @@ public class AccessPathNullnessPropagation
   @Override
   public TransferResult<Nullness, NullnessStore> visitEqualTo(
       EqualToNode equalToNode, TransferInput<Nullness, NullnessStore> input) {
-    ReadableUpdates thenUpdates = new ReadableUpdates();
-    ReadableUpdates elseUpdates = new ReadableUpdates();
-    handleEqualityComparison(
-        true,
-        equalToNode.getLeftOperand(),
-        equalToNode.getRightOperand(),
-        values(input),
-        thenUpdates,
-        elseUpdates);
-    ResultingStore thenStore = updateStore(input.getThenStore(), thenUpdates);
-    ResultingStore elseStore = updateStore(input.getElseStore(), elseUpdates);
-    return conditionalResult(
-        thenStore.store, elseStore.store, thenStore.storeChanged || elseStore.storeChanged);
+    return handleEqualityComparison(
+        input, equalToNode.getLeftOperand(), equalToNode.getRightOperand(), true);
   }
 
   @Override
   public TransferResult<Nullness, NullnessStore> visitNotEqual(
       NotEqualNode notEqualNode, TransferInput<Nullness, NullnessStore> input) {
-    ReadableUpdates thenUpdates = new ReadableUpdates();
-    ReadableUpdates elseUpdates = new ReadableUpdates();
-    handleEqualityComparison(
-        false,
-        notEqualNode.getLeftOperand(),
-        notEqualNode.getRightOperand(),
-        values(input),
-        thenUpdates,
-        elseUpdates);
-    ResultingStore thenStore = updateStore(input.getThenStore(), thenUpdates);
-    ResultingStore elseStore = updateStore(input.getElseStore(), elseUpdates);
-    return conditionalResult(
-        thenStore.store, elseStore.store, thenStore.storeChanged || elseStore.storeChanged);
+    return handleEqualityComparison(
+        input, notEqualNode.getLeftOperand(), notEqualNode.getRightOperand(), false);
   }
 
-  private void handleEqualityComparison(
-      boolean equalTo,
-      Node leftNode,
-      Node rightNode,
-      SubNodeValues inputs,
-      Updates thenUpdates,
-      Updates elseUpdates) {
-    Nullness leftVal = inputs.valueOfSubNode(leftNode);
-    Nullness rightVal = inputs.valueOfSubNode(rightNode);
+  /**
+   * Handle nullability refinements from an equality comparison.
+   *
+   * @param input transfer input for the operation
+   * @param leftOperand left operand of the comparison
+   * @param rightOperand right operand of the comparison
+   * @param equalTo if {@code true}, the comparison is an equality comparison, otherwise it is a
+   *     dis-equality ({@code !=}) comparison
+   * @return a TransferResult reflecting any updates from the comparison
+   */
+  private TransferResult<Nullness, NullnessStore> handleEqualityComparison(
+      TransferInput<Nullness, NullnessStore> input,
+      Node leftOperand,
+      Node rightOperand,
+      boolean equalTo) {
+    ReadableUpdates thenUpdates = new ReadableUpdates();
+    ReadableUpdates elseUpdates = new ReadableUpdates();
+    SubNodeValues inputs = values(input);
+    Nullness leftVal = inputs.valueOfSubNode(leftOperand);
+    Nullness rightVal = inputs.valueOfSubNode(rightOperand);
     Nullness equalBranchValue = leftVal.greatestLowerBound(rightVal);
     Updates equalBranchUpdates = equalTo ? thenUpdates : elseUpdates;
     Updates notEqualBranchUpdates = equalTo ? elseUpdates : thenUpdates;
 
-    Node realLeftNode = unwrapAssignExpr(leftNode);
-    Node realRightNode = unwrapAssignExpr(rightNode);
+    Node realLeftNode = unwrapAssignExpr(leftOperand);
+    Node realRightNode = unwrapAssignExpr(rightOperand);
 
     AccessPath leftAP = AccessPath.getAccessPathForNode(realLeftNode, state, apContext);
     if (leftAP != null) {
@@ -467,6 +456,10 @@ public class AccessPathNullnessPropagation
       notEqualBranchUpdates.set(
           rightAP, rightVal.greatestLowerBound(leftVal.deducedValueWhenNotEqual()));
     }
+    ResultingStore thenStore = updateStore(input.getThenStore(), thenUpdates);
+    ResultingStore elseStore = updateStore(input.getElseStore(), elseUpdates);
+    return conditionalResult(
+        thenStore.store, elseStore.store, thenStore.storeChanged || elseStore.storeChanged);
   }
 
   @Override
@@ -932,7 +925,18 @@ public class AccessPathNullnessPropagation
   @Override
   public TransferResult<Nullness, NullnessStore> visitCase(
       CaseNode caseNode, TransferInput<Nullness, NullnessStore> input) {
-    return noStoreChanges(NULLABLE, input);
+    List<Node> caseOperands = caseNode.getCaseOperands();
+    if (caseOperands.isEmpty()) {
+      return noStoreChanges(NULLABLE, input);
+    } else {
+      // `null` can only appear on its own as a case operand, or together with the default case
+      // (i.e., `case null, default:`).  So, it is safe to only look at the first case operand, and
+      // update the stores based on that.  We treat the case operation as an equality comparison
+      // between the switch expression and the case operand.
+      Node switchOperand = caseNode.getSwitchOperand().getExpression();
+      Node caseOperand = caseOperands.get(0);
+      return handleEqualityComparison(input, switchOperand, caseOperand, true);
+    }
   }
 
   @Override