Support @EnsuresNonNullIf

annotations/src/main/java/com/uber/nullaway/annotations/EnsuresNonNullIf.java

@@ -0,0 +1,54 @@
+package com.uber.nullaway.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * An annotation describing a nullability post-condition for an instance method. Each parameter to
+ * the annotation should be a field of the enclosing class. The method must ensure that the method
+ * returns true in case the fields are non-null. The method can also return false in case the fields
+ * are non-null (inverse logic), and in such case, you must set {@code result} to false. NullAway
+ * verifies that the property holds.
+ *
+ * <p>Here is an example:
+ *
+ * <pre>
+ * class Foo {
+ *     {@literal @}Nullable Object theField;
+ *
+ *     {@literal @}EnsuresNonNullIf("theField", result=true) // "this.theField" is also valid
+ *     boolean hasTheField() {
+ *         return theField != null;
+ *     }
+ *
+ *     void bar() {
+ *         if(!hasTheField()) {
+ *             return;
+ *         }
+ *
+ *         // No error, NullAway knows theField is non-null after call to hasTheField()
+ *         theField.toString();
+ *     }
+ * }
+ * </pre>
+ */
+@Retention(RetentionPolicy.CLASS)
+@Target({ElementType.METHOD})
+public @interface EnsuresNonNullIf {
+  /**
+   * The list of fields that are non-null after the method returns the given result.
+   *
+   * @return The list of field names
+   */
+  String[] value();
+
+  /**
+   * The return value of the method under which the postcondition holds. The default is set to true,
+   * which means the method should return true in case fields are non-null.
+   *
+   * @return true or false
+   */
+  boolean result() default true;
+}

nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java

@@ -230,14 +230,7 @@ public static Stream<? extends AnnotationMirror> getAllAnnotations(Symbol symbol
    */
   public static @Nullable Set<String> getAnnotationValueArray(
       Symbol.MethodSymbol methodSymbol, String annotName, boolean exactMatch) {
-    AnnotationMirror annot = null;
-    for (AnnotationMirror annotationMirror : methodSymbol.getAnnotationMirrors()) {
-      String name = AnnotationUtils.annotationName(annotationMirror);
-      if ((exactMatch && name.equals(annotName)) || (!exactMatch && name.endsWith(annotName))) {
-        annot = annotationMirror;
-        break;
-      }
-    }
+    AnnotationMirror annot = findAnnotation(methodSymbol, annotName, exactMatch);
     if (annot == null) {
       return null;
     }
@@ -255,6 +248,33 @@ public static Stream<? extends AnnotationMirror> getAllAnnotations(Symbol symbol
     return null;
   }
 
+  /**
+   * Retrieve the specific annotation of a method.
+   *
+   * @param methodSymbol A method to check for the annotation.
+   * @param annotName The qualified name or simple name of the annotation depending on the value of
+   *     {@code exactMatch}.
+   * @param exactMatch If true, the annotation name must match the full qualified name given in
+   *     {@code annotName}, otherwise, simple names will be checked.
+   * @return an {@code AnnotationMirror} representing that annotation, or null in case the
+   *     annotation with a given name {@code annotName} doesn't exist in {@code methodSymbol}.
+   */
+  public static @Nullable AnnotationMirror findAnnotation(
+      Symbol.MethodSymbol methodSymbol, String annotName, boolean exactMatch) {
+    AnnotationMirror annot = null;
+    for (AnnotationMirror annotationMirror : methodSymbol.getAnnotationMirrors()) {
+      String name = AnnotationUtils.annotationName(annotationMirror);
+      if ((exactMatch && name.equals(annotName)) || (!exactMatch && name.endsWith(annotName))) {
+        annot = annotationMirror;
+        break;
+      }
+    }
+    if (annot == null) {
+      return null;
+    }
+    return annot;
+  }
+
   /**
    * Works for method parameters defined either in source or in class files
    *

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

@@ -21,8 +21,8 @@
 import static com.uber.nullaway.NullabilityUtil.castToNonNull;
 
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
 import com.google.errorprone.VisitorState;
+import com.google.errorprone.dataflow.nullnesspropagation.NullnessAnalysis;
 import com.sun.source.tree.Tree;
 import com.sun.source.util.TreePath;
 import com.sun.tools.javac.util.Context;
@@ -157,25 +157,7 @@ public Set<Element> getNonnullFieldsOfReceiverAtExit(TreePath path, Context cont
       // be conservative and say nothing is initialized
       return Collections.emptySet();
     }
-    return getNonnullReceiverFields(nullnessResult);
-  }
-
-  private Set<Element> getNonnullReceiverFields(NullnessStore nullnessResult) {
-    Set<AccessPath> nonnullAccessPaths = nullnessResult.getAccessPathsWithValue(Nullness.NONNULL);
-    Set<Element> result = new LinkedHashSet<>();
-    for (AccessPath ap : nonnullAccessPaths) {
-      // A null root represents the receiver
-      if (ap.getRoot() == null) {
-        ImmutableList<AccessPathElement> elements = ap.getElements();
-        if (elements.size() == 1) {
-          Element elem = elements.get(0).getJavaElement();
-          if (elem.getKind().equals(ElementKind.FIELD)) {
-            result.add(elem);
-          }
-        }
-      }
-    }
-    return result;
+    return nullnessResult.getNonNullReceiverFields();
   }
 
   /**
@@ -190,7 +172,7 @@ public Set<Element> getNonnullFieldsOfReceiverBefore(TreePath path, Context cont
     if (store == null) {
       return Collections.emptySet();
     }
-    return getNonnullReceiverFields(store);
+    return store.getNonNullReceiverFields();
   }
 
   /**

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

@@ -852,7 +852,8 @@ public TransferResult<Nullness, NullnessStore> visitSuper(
   @Override
   public TransferResult<Nullness, NullnessStore> visitReturn(
       ReturnNode returnNode, TransferInput<Nullness, NullnessStore> input) {
-    handler.onDataflowVisitReturn(returnNode.getTree(), input.getThenStore(), input.getElseStore());
+    handler.onDataflowVisitReturn(
+        returnNode.getTree(), state, input.getThenStore(), input.getElseStore());
     return noStoreChanges(NULLABLE, input);
   }
 

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

@@ -19,6 +19,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.Sets.intersection;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.VisitorState;
 import com.uber.nullaway.Nullness;
@@ -30,6 +31,7 @@
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ElementKind;
 import org.checkerframework.nullaway.dataflow.analysis.Store;
 import org.checkerframework.nullaway.dataflow.cfg.node.FieldAccessNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.LocalVariableNode;
@@ -261,6 +263,39 @@ public NullnessStore filterAccessPaths(Predicate<AccessPath> pred) {
             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
   }
 
+  /**
+   * Return all the fields in the store that are Non-Null.
+   *
+   * @return Set of fields (represented as {@code Element}s) that are non-null
+   */
+  public Set<Element> getNonNullReceiverFields() {
+    return getReceiverFields(Nullness.NONNULL);
+  }
+
+  /**
+   * Return all the fields in the store that hold the {@code nullness} state.
+   *
+   * @param nullness The {@code Nullness} state
+   * @return Set of fields (represented as {@code Element}s) with the given {@code nullness}.
+   */
+  public Set<Element> getReceiverFields(Nullness nullness) {
+    Set<AccessPath> nonnullAccessPaths = this.getAccessPathsWithValue(nullness);
+    Set<Element> result = new LinkedHashSet<>();
+    for (AccessPath ap : nonnullAccessPaths) {
+      // A null root represents the receiver
+      if (ap.getRoot() == null) {
+        ImmutableList<AccessPathElement> elements = ap.getElements();
+        if (elements.size() == 1) {
+          Element elem = elements.get(0).getJavaElement();
+          if (elem.getKind().equals(ElementKind.FIELD)) {
+            result.add(elem);
+          }
+        }
+      }
+    }
+    return result;
+  }
+
   /** class for building up instances of the store. */
   public static final class Builder {
     private final Map<AccessPath, Nullness> contents;

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

@@ -174,7 +174,7 @@ public NullnessHint onDataflowVisitFieldAccess(
 
   @Override
   public void onDataflowVisitReturn(
-      ReturnTree tree, NullnessStore thenStore, NullnessStore elseStore) {
+      ReturnTree tree, VisitorState state, NullnessStore thenStore, NullnessStore elseStore) {
     // NoOp
   }
 

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

@@ -220,9 +220,9 @@ public NullnessHint onDataflowVisitFieldAccess(
 
   @Override
   public void onDataflowVisitReturn(
-      ReturnTree tree, NullnessStore thenStore, NullnessStore elseStore) {
+      ReturnTree tree, VisitorState state, NullnessStore thenStore, NullnessStore elseStore) {
     for (Handler h : handlers) {
-      h.onDataflowVisitReturn(tree, thenStore, elseStore);
+      h.onDataflowVisitReturn(tree, state, thenStore, elseStore);
     }
   }
 

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

@@ -269,12 +269,14 @@ NullnessHint onDataflowVisitFieldAccess(
    * Called when the Dataflow analysis visits a return statement.
    *
    * @param tree The AST node for the return statement being matched.
+   * @param state The current visitor state
    * @param thenStore The NullnessStore for the true case of the expression inside the return
    *     statement.
    * @param elseStore The NullnessStore for the false case of the expression inside the return
    *     statement.
    */
-  void onDataflowVisitReturn(ReturnTree tree, NullnessStore thenStore, NullnessStore elseStore);
+  void onDataflowVisitReturn(
+      ReturnTree tree, VisitorState state, NullnessStore thenStore, NullnessStore elseStore);
 
   /**
    * Called when the Dataflow analysis visits the result expression inside the body of lambda.

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

@@ -27,6 +27,7 @@
 import com.uber.nullaway.handlers.contract.ContractCheckHandler;
 import com.uber.nullaway.handlers.contract.ContractHandler;
 import com.uber.nullaway.handlers.contract.fieldcontract.EnsuresNonNullHandler;
+import com.uber.nullaway.handlers.contract.fieldcontract.EnsuresNonNullIfHandler;
 import com.uber.nullaway.handlers.contract.fieldcontract.RequiresNonNullHandler;
 import com.uber.nullaway.handlers.temporary.FluentFutureHandler;
 
@@ -69,6 +70,7 @@ public static Handler buildDefault(Config config) {
     handlerListBuilder.add(new GrpcHandler());
     handlerListBuilder.add(new RequiresNonNullHandler());
     handlerListBuilder.add(new EnsuresNonNullHandler());
+    handlerListBuilder.add(new EnsuresNonNullIfHandler());
     handlerListBuilder.add(new SynchronousCallbackHandler());
     if (config.serializationIsActive() && config.getSerializationConfig().fieldInitInfoEnabled) {
       handlerListBuilder.add(

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

@@ -595,7 +595,7 @@ public NullnessStore.Builder onDataflowInitialStore(
 
   @Override
   public void onDataflowVisitReturn(
-      ReturnTree tree, NullnessStore thenStore, NullnessStore elseStore) {
+      ReturnTree tree, VisitorState state, NullnessStore thenStore, NullnessStore elseStore) {
     Tree filterTree = returnToEnclosingMethodOrLambda.get(tree);
     if (filterTree != null) {
       assert (filterTree instanceof MethodTree || filterTree instanceof LambdaExpressionTree);

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

@@ -40,7 +40,6 @@
 import com.uber.nullaway.handlers.MethodAnalysisContext;
 import com.uber.nullaway.handlers.contract.ContractUtils;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.lang.model.element.VariableElement;
@@ -124,46 +123,8 @@ protected void validateOverridingRules(
       VisitorState state,
       MethodTree tree,
       Symbol.MethodSymbol overriddenMethod) {
-    Set<String> overriddenFieldNames = getAnnotationValueArray(overriddenMethod, annotName, false);
-    if (overriddenFieldNames == null) {
-      return;
-    }
-    if (overridingFieldNames == null) {
-      overridingFieldNames = Collections.emptySet();
-    }
-    if (overridingFieldNames.containsAll(overriddenFieldNames)) {
-      return;
-    }
-    overriddenFieldNames.removeAll(overridingFieldNames);
-
-    StringBuilder errorMessage = new StringBuilder();
-    errorMessage
-        .append(
-            "postcondition inheritance is violated, this method must guarantee that all fields written in the @EnsuresNonNull annotation of overridden method ")
-        .append(castToNonNull(ASTHelpers.enclosingClass(overriddenMethod)).getSimpleName())
-        .append(".")
-        .append(overriddenMethod.getSimpleName())
-        .append(" are @NonNull at exit point as well. Fields [");
-    Iterator<String> iterator = overriddenFieldNames.iterator();
-    while (iterator.hasNext()) {
-      errorMessage.append(iterator.next());
-      if (iterator.hasNext()) {
-        errorMessage.append(", ");
-      }
-    }
-    errorMessage.append(
-        "] must explicitly appear as parameters at this method @EnsuresNonNull annotation");
-    state.reportMatch(
-        analysis
-            .getErrorBuilder()
-            .createErrorDescription(
-                new ErrorMessage(
-                    ErrorMessage.MessageTypes.WRONG_OVERRIDE_POSTCONDITION,
-                    errorMessage.toString()),
-                tree,
-                analysis.buildDescription(tree),
-                state,
-                null));
+    FieldContractUtils.ensureStrictPostConditionInheritance(
+        annotName, overridingFieldNames, analysis, state, tree, overriddenMethod);
   }
 
   /**