From 41c7c5505f0fed742912f74d58938a236ba6ac48 Mon Sep 17 00:00:00 2001
From: Manu Sridharan <msridhar@gmail.com>
Date: Wed, 18 Dec 2024 12:27:11 -0800
Subject: [PATCH] Handle case null in switch statements (#1100)

In Java 21+, when a `switch` statement contains `case null`, the
expression being switched on is allowed to be `null`.

Fixes #930
---
 .../com/uber/nullaway/jdk17/SwitchTests.java  | 23 +++++++++++++++++++
 .../main/java/com/uber/nullaway/NullAway.java |  2 +-
 .../java/com/uber/nullaway/CoreTests.java     |  2 +-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/SwitchTests.java b/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/SwitchTests.java
index f59bc3d6ac..1af17d88f8 100644
--- a/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/SwitchTests.java
+++ b/jdk-recent-unit-tests/src/test/java/com/uber/nullaway/jdk17/SwitchTests.java
@@ -244,4 +244,27 @@ public void testSwitchExprNullCaseDataflow() {
             "}")
         .doTest();
   }
+
+  @Test
+  public void switchStatementCaseNull() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "public class Test {",
+            "    public enum NullableEnum {",
+            "        VALUE1, VALUE2;",
+            "    }",
+            "    static Object handleNullableEnumCaseNullDefaultStatement(@Nullable NullableEnum nullableEnum) {",
+            "        Object result;",
+            "        switch (nullableEnum) {",
+            "            case null -> result = new Object();",
+            "            default -> result = nullableEnum.toString();",
+            "        }",
+            "        return result;",
+            "    }",
+            "}")
+        .doTest();
+  }
 }
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 6a31bb7bb6..9964550754 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -688,7 +688,7 @@ public Description matchSwitch(SwitchTree tree, VisitorState state) {
       switchSelectorExpression = ((ParenthesizedTree) switchSelectorExpression).getExpression();
     }
 
-    if (mayBeNullExpr(state, switchSelectorExpression)) {
+    if (!TreeUtils.hasNullCaseLabel(tree) && mayBeNullExpr(state, switchSelectorExpression)) {
       String message =
           "switch expression " + state.getSourceForNode(switchSelectorExpression) + " is @Nullable";
       ErrorMessage errorMessage =
diff --git a/nullaway/src/test/java/com/uber/nullaway/CoreTests.java b/nullaway/src/test/java/com/uber/nullaway/CoreTests.java
index 17cc70455c..0cf3237dd9 100644
--- a/nullaway/src/test/java/com/uber/nullaway/CoreTests.java
+++ b/nullaway/src/test/java/com/uber/nullaway/CoreTests.java
@@ -207,7 +207,7 @@ public void cfNullableArrayField() {
   }
 
   @Test
-  public void supportSwitchExpression() {
+  public void switchOnNullable() {
     defaultCompilationHelper
         .addSourceLines(
             "TestPositive.java",