Add support for static fields in contracts (#1118)

This PR adds support for static fields in
`@EnsureNonnull`,`EnsureNonnullIf`,`@RequiresNonnull` annotations.
Currently the following code will throw validation errors (as well as
the annotation handlers are unable to handle static fields). However,
after this change, static fields for all three annotations are supported

```
class Foo {
  @nullable static Item staticNullableItem;

  @EnsuresNonNull("staticNullableItem")
  public static void initializeStaticField() {
    staticNullableItem = new Item();
  }

  @RequiresNonNull("staticNullableItem")
  public static void useStaticField() {
    staticNullableItem.call();
  }

  @EnsuresNonNullIf(value="staticNullableItem", result=true)
  public static boolean hasStaticNullableItem() {
    return staticNullableItem != null;
  }
}
```
Fixes #431

---------

Co-authored-by: Manu Sridharan <msridhar@gmail.com>

nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java

@@ -369,6 +369,19 @@ public static AccessPath fromFieldElement(VariableElement element) {
     return new AccessPath(null, ImmutableList.of(new FieldOrMethodCallElement(element)));
   }
 
+  /**
+   * Constructs an access path representing a static field.
+   *
+   * @param element element that must represent a static field
+   * @return an access path representing the static field
+   */
+  public static AccessPath fromStaticField(VariableElement element) {
+    Preconditions.checkArgument(
+        element.getKind().isField() && element.getModifiers().contains(Modifier.STATIC),
+        "element must be a static field but received: " + element.getKind());
+    return new AccessPath(element, ImmutableList.of(), null);
+  }
+
   private static boolean isBoxingMethod(Symbol.MethodSymbol methodSymbol) {
     if (methodSymbol.isStatic() && methodSymbol.getSimpleName().contentEquals("valueOf")) {
       Symbol.PackageSymbol enclosingPackage = ASTHelpers.enclosingPackage(methodSymbol.enclClass());

nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStore.java

@@ -32,6 +32,7 @@
 import java.util.stream.Collectors;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.Modifier;
 import org.checkerframework.nullaway.dataflow.analysis.Store;
 import org.checkerframework.nullaway.dataflow.cfg.node.FieldAccessNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.LocalVariableNode;
@@ -272,6 +273,24 @@ public Set<Element> getNonNullReceiverFields() {
     return getReceiverFields(Nullness.NONNULL);
   }
 
+  /**
+   * Return all the static fields in the store that are Non-Null.
+   *
+   * @return Set of static fields (represented as {@code Element}s) that are non-null
+   */
+  public Set<Element> getNonNullStaticFields() {
+    Set<AccessPath> nonnullAccessPaths = this.getAccessPathsWithValue(Nullness.NONNULL);
+    Set<Element> result = new LinkedHashSet<>();
+    for (AccessPath ap : nonnullAccessPaths) {
+      if (ap.getRoot() != null) {
+        if (ap.getRoot().getModifiers().contains(Modifier.STATIC)) {
+          result.add(ap.getRoot());
+        }
+      }
+    }
+    return result;
+  }
+
   /**
    * Return all the fields in the store that hold the {@code nullness} state.
    *

nullaway/src/main/java/com/uber/nullaway/handlers/AbstractFieldContractHandler.java

@@ -163,7 +163,32 @@ protected boolean validateAnnotationSyntax(
                   null));
       return false;
     } else {
+      Symbol.ClassSymbol classSymbol =
+          castToNonNull(ASTHelpers.enclosingClass(methodAnalysisContext.methodSymbol()));
       for (String fieldName : content) {
+        if (isThisDotStaticField(classSymbol, fieldName)) {
+
+          message =
+              "Cannot refer to static field "
+                  + fieldName.substring(THIS_NOTATION.length())
+                  + " using this.";
+          state.reportMatch(
+              analysis
+                  .getErrorBuilder()
+                  .createErrorDescription(
+                      new ErrorMessage(ErrorMessage.MessageTypes.ANNOTATION_VALUE_INVALID, message),
+                      tree,
+                      analysis.buildDescription(tree),
+                      state,
+                      null));
+          return false;
+        }
+        VariableElement field = getFieldOfClass(classSymbol, fieldName);
+        if (field != null) {
+          if (field.getModifiers().contains(Modifier.STATIC)) {
+            continue;
+          }
+        }
         if (fieldName.contains(".")) {
           if (!fieldName.startsWith(THIS_NOTATION)) {
             message =
@@ -188,9 +213,7 @@ protected boolean validateAnnotationSyntax(
             fieldName = fieldName.substring(fieldName.lastIndexOf(".") + 1);
           }
         }
-        Symbol.ClassSymbol classSymbol =
-            castToNonNull(ASTHelpers.enclosingClass(methodAnalysisContext.methodSymbol()));
-        VariableElement field = getInstanceFieldOfClass(classSymbol, fieldName);
+        field = getFieldOfClass(classSymbol, fieldName);
         if (field == null) {
           message =
               "For @"
@@ -223,20 +246,31 @@ protected boolean validateAnnotationSyntax(
    * @param name Name of the field.
    * @return The field with the given name, or {@code null} if the field cannot be found
    */
-  public static @Nullable VariableElement getInstanceFieldOfClass(
+  public static @Nullable VariableElement getFieldOfClass(
       Symbol.ClassSymbol classSymbol, String name) {
     Preconditions.checkNotNull(classSymbol);
     for (Element member : getEnclosedElements(classSymbol)) {
-      if (member.getKind().isField() && !member.getModifiers().contains(Modifier.STATIC)) {
+      if (member.getKind().isField()) {
         if (member.getSimpleName().toString().equals(name)) {
           return (VariableElement) member;
         }
       }
     }
     Symbol.ClassSymbol superClass = (Symbol.ClassSymbol) classSymbol.getSuperclass().tsym;
     if (superClass != null) {
-      return getInstanceFieldOfClass(superClass, name);
+      return getFieldOfClass(superClass, name);
     }
     return null;
   }
+
+  protected boolean isThisDotStaticField(Symbol.ClassSymbol classSymbol, String expression) {
+    if (expression.contains(".")) {
+      if (expression.startsWith(THIS_NOTATION)) {
+        String fieldName = expression.substring(THIS_NOTATION.length());
+        VariableElement field = getFieldOfClass(classSymbol, fieldName);
+        return field != null && field.getModifiers().contains(Modifier.STATIC);
+      }
+    }
+    return false;
+  }
 }

nullaway/src/main/java/com/uber/nullaway/handlers/contract/fieldcontract/EnsuresNonNullHandler.java

@@ -40,8 +40,10 @@
 import com.uber.nullaway.handlers.MethodAnalysisContext;
 import com.uber.nullaway.handlers.contract.ContractUtils;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.Set;
 import java.util.stream.Collectors;
+import javax.lang.model.element.Modifier;
 import javax.lang.model.element.VariableElement;
 import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
 
@@ -85,14 +87,26 @@ protected boolean validateAnnotationSemantics(
             .stream()
             .map(e -> e.getSimpleName().toString())
             .collect(Collectors.toSet());
+    Set<String> nonnullStaticFieldsAtExit =
+        analysis
+            .getNullnessAnalysis(state)
+            .getNonnullStaticFieldsAtExit(new TreePath(state.getPath(), tree), state.context)
+            .stream()
+            .map(e -> e.getSimpleName().toString())
+            .collect(Collectors.toSet());
+    Set<String> nonnullFieldsAtExit = new HashSet<String>();
+    nonnullFieldsAtExit.addAll(nonnullFieldsOfReceiverAtExit);
+    nonnullFieldsAtExit.addAll(nonnullStaticFieldsAtExit);
     Set<String> fieldNames = getAnnotationValueArray(methodSymbol, annotName, false);
     if (fieldNames == null) {
       fieldNames = Collections.emptySet();
     }
     fieldNames = ContractUtils.trimReceivers(fieldNames);
-    boolean isValidLocalPostCondition = nonnullFieldsOfReceiverAtExit.containsAll(fieldNames);
+
+    nonnullFieldsAtExit = ContractUtils.trimReceivers(nonnullFieldsAtExit);
+    boolean isValidLocalPostCondition = nonnullFieldsAtExit.containsAll(fieldNames);
     if (!isValidLocalPostCondition) {
-      fieldNames.removeAll(nonnullFieldsOfReceiverAtExit);
+      fieldNames.removeAll(nonnullFieldsAtExit);
       message =
           String.format(
               "Method is annotated with @EnsuresNonNull but fails to ensure the following fields are non-null at exit: %s",
@@ -153,14 +167,15 @@ public NullnessHint onDataflowVisitMethodInvocation(
       fieldNames = ContractUtils.trimReceivers(fieldNames);
       for (String fieldName : fieldNames) {
         VariableElement field =
-            getInstanceFieldOfClass(
-                castToNonNull(ASTHelpers.enclosingClass(methodSymbol)), fieldName);
+            getFieldOfClass(castToNonNull(ASTHelpers.enclosingClass(methodSymbol)), fieldName);
         if (field == null) {
           // Invalid annotation, will result in an error during validation. For now, skip field.
           continue;
         }
         AccessPath accessPath =
-            AccessPath.fromBaseAndElement(node.getTarget().getReceiver(), field, apContext);
+            field.getModifiers().contains(Modifier.STATIC)
+                ? AccessPath.fromStaticField(field)
+                : AccessPath.fromBaseAndElement(node.getTarget().getReceiver(), field, apContext);
         if (accessPath == null) {
           continue;
         }

nullaway/src/main/java/com/uber/nullaway/handlers/contract/fieldcontract/EnsuresNonNullIfHandler.java

@@ -51,6 +51,7 @@
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.Modifier;
 import javax.lang.model.element.VariableElement;
 import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
 import org.jspecify.annotations.Nullable;
@@ -184,6 +185,11 @@ public void onDataflowVisitReturn(
         chosenStore.getNonNullReceiverFields().stream()
             .map(e -> e.getSimpleName().toString())
             .collect(Collectors.toSet());
+    Set<String> nonNullStaticFieldsInPath =
+        chosenStore.getNonNullStaticFields().stream()
+            .map(e -> e.getSimpleName().toString())
+            .collect(Collectors.toSet());
+    nonNullFieldsInPath.addAll(nonNullStaticFieldsInPath);
     boolean allFieldsAreNonNull = nonNullFieldsInPath.containsAll(fieldNames);
 
     // Whether the return true expression evaluates to a boolean literal or not.  If null, then not
@@ -297,14 +303,16 @@ public NullnessHint onDataflowVisitMethodInvocation(
           trueIfNonNull ? thenUpdates : elseUpdates;
       for (String fieldName : fieldNames) {
         VariableElement field =
-            getInstanceFieldOfClass(
-                castToNonNull(ASTHelpers.enclosingClass(methodSymbol)), fieldName);
+            getFieldOfClass(castToNonNull(ASTHelpers.enclosingClass(methodSymbol)), fieldName);
         if (field == null) {
           // Invalid annotation, will result in an error during validation.
           continue;
         }
         AccessPath accessPath =
-            AccessPath.fromBaseAndElement(node.getTarget().getReceiver(), field, apContext);
+            field.getModifiers().contains(Modifier.STATIC)
+                ? AccessPath.fromStaticField(field)
+                : AccessPath.fromFieldElement(field);
+
         if (accessPath == null) {
           // Also likely to be an invalid annotation, will result in an error during validation.
           continue;

nullaway/src/main/java/com/uber/nullaway/handlers/contract/fieldcontract/RequiresNonNullHandler.java

@@ -43,9 +43,12 @@
 import com.uber.nullaway.handlers.MethodAnalysisContext;
 import com.uber.nullaway.handlers.contract.ContractUtils;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
+import javax.lang.model.element.Element;
+import javax.lang.model.element.Modifier;
 import javax.lang.model.element.VariableElement;
 import org.checkerframework.nullaway.dataflow.cfg.UnderlyingAST;
 import org.checkerframework.nullaway.dataflow.cfg.node.LocalVariableNode;
@@ -145,11 +148,35 @@ public void onMatchMethodInvocation(
       Symbol.ClassSymbol classSymbol = ASTHelpers.enclosingClass(methodSymbol);
       Preconditions.checkNotNull(
           classSymbol, "Could not find the enclosing class for method symbol: " + methodSymbol);
-      VariableElement field = getInstanceFieldOfClass(classSymbol, fieldName);
+      VariableElement field = getFieldOfClass(classSymbol, fieldName);
       if (field == null) {
         // we will report an error on the method declaration
         continue;
       }
+      if (field.getModifiers().contains(Modifier.STATIC)) {
+        Set<Element> nonnullStaticFields =
+            analysis
+                .getNullnessAnalysis(state)
+                .getNonnullStaticFieldsBefore(state.getPath(), state.context);
+
+        if (!nonnullStaticFields.contains(field)) {
+          String message =
+              "Expected static field "
+                  + fieldName
+                  + " to be non-null at call site due to @RequiresNonNull annotation on invoked method";
+          state.reportMatch(
+              analysis
+                  .getErrorBuilder()
+                  .createErrorDescription(
+                      new ErrorMessage(
+                          ErrorMessage.MessageTypes.PRECONDITION_NOT_SATISFIED, message),
+                      tree,
+                      analysis.buildDescription(tree),
+                      state,
+                      null));
+        }
+        continue;
+      }
       ExpressionTree methodSelectTree = tree.getMethodSelect();
       Nullness nullness =
           analysis
@@ -195,14 +222,23 @@ public NullnessStore.Builder onDataflowInitialStore(
     if (fieldNames == null) {
       return result;
     }
-    fieldNames = ContractUtils.trimReceivers(fieldNames);
+    Set<String> filteredFieldNames = new HashSet<>();
     for (String fieldName : fieldNames) {
-      VariableElement field = getInstanceFieldOfClass(ASTHelpers.getSymbol(classTree), fieldName);
+      if (!isThisDotStaticField(ASTHelpers.getSymbol(classTree), fieldName)) {
+        filteredFieldNames.add(fieldName);
+      }
+    }
+    filteredFieldNames = ContractUtils.trimReceivers(filteredFieldNames);
+    for (String fieldName : filteredFieldNames) {
+      VariableElement field = getFieldOfClass(ASTHelpers.getSymbol(classTree), fieldName);
       if (field == null) {
         // Invalid annotation, will result in an error during validation. For now, skip field.
         continue;
       }
-      AccessPath accessPath = AccessPath.fromFieldElement(field);
+      AccessPath accessPath =
+          field.getModifiers().contains(Modifier.STATIC)
+              ? AccessPath.fromStaticField(field)
+              : AccessPath.fromFieldElement(field);
       result.setInformation(accessPath, Nullness.NONNULL);
     }
     return result;

nullaway/src/test/java/com/uber/nullaway/EnsuresNonNullIfTests.java

@@ -893,4 +893,59 @@ public void ensuresNonNullMethodResultStoredInVariableInverse() {
             "Item.java", "package com.uber;", "class Item {", "  public void call() { }", "}")
         .doTest();
   }
+
+  @Test
+  public void staticFieldCorrectUse() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Foo.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import com.uber.nullaway.annotations.EnsuresNonNullIf;",
+            "class Foo {",
+            "  @Nullable static Item staticNullableItem;",
+            "  @EnsuresNonNullIf(value=\"staticNullableItem\", result=true)",
+            "  public static boolean hasStaticNullableItem() {",
+            "    return staticNullableItem != null;",
+            "  }",
+            "  public static int runOk() {",
+            "    if(!hasStaticNullableItem()) {",
+            "      return 1;",
+            "    }",
+            "    staticNullableItem.call();",
+            "    return 0;",
+            "  }",
+            "}")
+        .addSourceLines(
+            "Item.java", "package com.uber;", "class Item {", "  public void call() { }", "}")
+        .doTest();
+  }
+
+  @Test
+  public void staticFieldIncorrectUse() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Foo.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import com.uber.nullaway.annotations.EnsuresNonNullIf;",
+            "class Foo {",
+            "  @Nullable static Item nullableItem;",
+            "  @EnsuresNonNullIf(value=\"nullableItem\", result=true)",
+            "  public boolean hasNullableItem() {",
+            "    return nullableItem != null;",
+            "  }",
+            "  public int runOk() {",
+            "    if(hasNullableItem()) {",
+            "      nullableItem.call();",
+            "    }",
+            "    // BUG: Diagnostic contains: dereferenced expression nullableItem is @Nullable",
+            "    nullableItem.call();",
+            "    return 0;",
+            "  }",
+            "}")
+        .addSourceLines(
+            "Item.java", "package com.uber;", "class Item {", "  public void call() { }", "}")
+        .doTest();
+  }
 }