From 8531e71ab39e4073ce7c13ffbddf9469444794e8 Mon Sep 17 00:00:00 2001
From: lazaro <lazaro@uber.com>
Date: Thu, 14 Oct 2021 22:08:50 +0000
Subject: [PATCH] Support methods taking compile-time constant fields in Access
 Paths.

This builds upon PR #285 by adding support to code like:

```
if (x.get(ZERO) != null && x.get(ZERO).foo() != null) {
   return x.get(ZERO).foo().bar();
}
```

Where `ZERO` is a compile-time evaluated constant. That is,
a `static final` variable or field of primitive or string type.
---
 .../uber/nullaway/dataflow/AccessPath.java    | 27 ++++++++++++++++++-
 .../java/com/uber/nullaway/NullAwayTest.java  | 27 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java
index 51f07f3ead..eaec22f3ce 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java
@@ -36,6 +36,7 @@
 import java.util.List;
 import javax.annotation.Nullable;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.VariableElement;
 import org.checkerframework.nullaway.dataflow.cfg.node.FieldAccessNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.IntegerLiteralNode;
@@ -323,7 +324,31 @@ && isBoxingMethod(ASTHelpers.getSymbol(methodInvocationTree))) {
               constantArgumentValues.add(((LiteralTree) tree).getValue().toString());
               break;
             case NULL_LITERAL:
-              // Um, probably not? Cascade to default for now.
+              // Um, probably not? Return null for now.
+              return null; // Not an AP
+            case MEMBER_SELECT: // check for Foo.CONST
+            case IDENTIFIER: // check for CONST
+              // Check for a constant field (static final)
+              Symbol symbol = ASTHelpers.getSymbol(tree);
+              if (symbol.getKind().equals(ElementKind.FIELD)) {
+                Symbol.VarSymbol varSymbol = (Symbol.VarSymbol) symbol;
+                // From docs: getConstantValue() returns the value of this variable if this is a
+                // static final field initialized to a compile-time constant. Returns null
+                // otherwise.
+                // This means that foo(FOUR) will match foo(4) iff FOUR=4 is a compile time
+                // constant :)
+                // The above will not work for static final fields of reference type, since they are
+                // initialized at class-initialization time, not compile time. Properly handling
+                // such fields would further require proving deep immutability for the object type
+                // itself.
+                Object constantValue = varSymbol.getConstantValue();
+                if (constantValue != null) {
+                  constantArgumentValues.add(constantValue.toString());
+                  break;
+                }
+              }
+              // Cascade to default, symbol is not a constant field
+              // fall through
             default:
               return null; // Not an AP
           }
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
index 3c0091faea..d8ae7fbe93 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
@@ -2399,6 +2399,33 @@ public void testConstantsInAccessPathsPositive() {
         .doTest();
   }
 
+  @Test
+  public void testVariablesInAccessPathsNegative() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "NullableContainer.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "public interface NullableContainer<K, V> {",
+            " @Nullable public V get(K k);",
+            "}")
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "public class Test {",
+            "  private static final int INT_KEY = 42;", // Guaranteed constant!
+            "  public void testEnhancedFor(NullableContainer<String, NullableContainer<Integer, Object>> c) {",
+            "    if (c.get(\"KEY_STR\") != null && c.get(\"KEY_STR\").get(INT_KEY) != null) {",
+            "      c.get(\"KEY_STR\").get(INT_KEY).toString();",
+            "      c.get(\"KEY_STR\").get(Test.INT_KEY).toString();",
+            "      c.get(\"KEY_STR\").get(42).toString();", // Extra magic!
+            "    }",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void testVariablesInAccessPathsPositive() {
     defaultCompilationHelper