From e529f53b4720030617f8de7c4d93efdd85181c80 Mon Sep 17 00:00:00 2001
From: Shubham Ugare <shubhamugare@uber.com>
Date: Sun, 7 Apr 2019 12:36:09 +0530
Subject: [PATCH] Nullable switch expression support

---
 .../java/com/uber/nullaway/ErrorBuilder.java  |  1 +
 .../java/com/uber/nullaway/ErrorMessage.java  |  3 +-
 .../main/java/com/uber/nullaway/NullAway.java | 27 +++++++++-
 .../java/com/uber/nullaway/NullAwayTest.java  | 51 +++++++++++++++++++
 4 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java b/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
index b305d072b..2f2763d85 100644
--- a/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
+++ b/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
@@ -100,6 +100,7 @@ public Description createErrorDescription(
         case RETURN_NULLABLE:
         case PASS_NULLABLE:
         case ASSIGN_FIELD_NULLABLE:
+        case SWITCH_EXPRESSION_NULLABLE:
           if (config.getCastToNonNullMethod() != null) {
             builder = addCastToNonNullFix(suggestTree, builder);
           } else {
diff --git a/nullaway/src/main/java/com/uber/nullaway/ErrorMessage.java b/nullaway/src/main/java/com/uber/nullaway/ErrorMessage.java
index 6b12e8c68..147c10a44 100644
--- a/nullaway/src/main/java/com/uber/nullaway/ErrorMessage.java
+++ b/nullaway/src/main/java/com/uber/nullaway/ErrorMessage.java
@@ -51,6 +51,7 @@ public enum MessageTypes {
     NULLABLE_VARARGS_UNSUPPORTED,
     ANNOTATION_VALUE_INVALID,
     CAST_TO_NONNULL_ARG_NONNULL,
-    GET_ON_EMPTY_OPTIONAL;
+    GET_ON_EMPTY_OPTIONAL,
+    SWITCH_EXPRESSION_NULLABLE
   }
 }
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 7683d9e21..83b91baab 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -72,6 +72,7 @@
 import com.sun.source.tree.ParenthesizedTree;
 import com.sun.source.tree.ReturnTree;
 import com.sun.source.tree.StatementTree;
+import com.sun.source.tree.SwitchTree;
 import com.sun.source.tree.Tree;
 import com.sun.source.tree.TryTree;
 import com.sun.source.tree.TypeCastTree;
@@ -167,7 +168,8 @@ public class NullAway extends BugChecker
         BugChecker.LambdaExpressionTreeMatcher,
         BugChecker.IdentifierTreeMatcher,
         BugChecker.MemberReferenceTreeMatcher,
-        BugChecker.CompoundAssignmentTreeMatcher {
+        BugChecker.CompoundAssignmentTreeMatcher,
+        BugChecker.SwitchTreeMatcher {
 
   static final String INITIALIZATION_CHECK_NAME = "NullAway.Init";
 
@@ -495,6 +497,29 @@ public Description matchMethod(MethodTree tree, VisitorState state) {
     return Description.NO_MATCH;
   }
 
+  @Override
+  public Description matchSwitch(SwitchTree tree, VisitorState state) {
+    if (!matchWithinClass) {
+      return Description.NO_MATCH;
+    }
+
+    ExpressionTree switchExpression = tree.getExpression();
+    if (switchExpression instanceof ParenthesizedTree) {
+      switchExpression = ((ParenthesizedTree) switchExpression).getExpression();
+    }
+
+    if (mayBeNullExpr(state, switchExpression)) {
+      final String message = "switch expression " + switchExpression.toString() + " is @Nullable";
+      ErrorMessage errorMessage =
+          new ErrorMessage(MessageTypes.SWITCH_EXPRESSION_NULLABLE, message);
+
+      return errorBuilder.createErrorDescription(
+          errorMessage, switchExpression, buildDescription(switchExpression));
+    }
+
+    return Description.NO_MATCH;
+  }
+
   /**
    * checks that an overriding method does not override a {@code @Nullable} parameter with a
    * {@code @NonNull} parameter
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
index 996e93497..dff3d50a5 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
@@ -966,6 +966,57 @@ public void supportObjectsIsNull() {
         .doTest();
   }
 
+  @Test
+  public void supportSwitchExpression() {
+    compilationHelper
+        .addSourceLines(
+            "TestPositive.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "enum Level {",
+            " HIGH, MEDIUM, LOW }",
+            "class TestPositive {",
+            "   void foo(@Nullable Integer s) {",
+            "    // BUG: Diagnostic contains: switch expression s is @Nullable",
+            "    switch(s) {",
+            "      case 5: break;",
+            "    }",
+            "    String x = null;",
+            "    // BUG: Diagnostic contains: switch expression x is @Nullable",
+            "    switch(x) {",
+            "      default: break;",
+            "    }",
+            "    Level level = null;",
+            "    // BUG: Diagnostic contains: switch expression level is @Nullable",
+            "    switch (level) {",
+            "      default: break; }",
+            "    }",
+            "}")
+        .addSourceLines(
+            "TestNegative.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "class TestNegative {",
+            "   void foo(Integer s, short y) {",
+            "    switch(s) {",
+            "      case 5: break;",
+            "    }",
+            "    String x = \"irrelevant\";",
+            "    switch(x) {",
+            "      default: break;",
+            "    }",
+            "    switch(y) {",
+            "      default: break;",
+            "    }",
+            "    Level level = Level.HIGH;",
+            "    switch (level) {",
+            "      default: break;",
+            "    }",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void defaultPermissiveOnUnannotated() {
     compilationHelper