DO NOT LAND TESTING BENCHMARKING

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

@@ -230,7 +230,14 @@ public static Stream<? extends AnnotationMirror> getAllAnnotations(Symbol symbol
    */
   public static @Nullable Set<String> getAnnotationValueArray(
       Symbol.MethodSymbol methodSymbol, String annotName, boolean exactMatch) {
-    AnnotationMirror annot = findAnnotation(methodSymbol, annotName, 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;
     }
@@ -248,33 +255,6 @@ 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,7 +157,25 @@ public Set<Element> getNonnullFieldsOfReceiverAtExit(TreePath path, Context cont
       // be conservative and say nothing is initialized
       return Collections.emptySet();
     }
-    return nullnessResult.getNonNullReceiverFields();
+    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;
   }
 
   /**
@@ -172,7 +190,7 @@ public Set<Element> getNonnullFieldsOfReceiverBefore(TreePath path, Context cont
     if (store == null) {
       return Collections.emptySet();
     }
-    return store.getNonNullReceiverFields();
+    return getNonnullReceiverFields(store);
   }
 
   /**

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

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

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

@@ -19,7 +19,6 @@
 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;
@@ -31,7 +30,6 @@
 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;
@@ -263,39 +261,6 @@ 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, VisitorState state, NullnessStore thenStore, NullnessStore elseStore) {
+      ReturnTree tree, 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, VisitorState state, NullnessStore thenStore, NullnessStore elseStore) {
+      ReturnTree tree, NullnessStore thenStore, NullnessStore elseStore) {
     for (Handler h : handlers) {
-      h.onDataflowVisitReturn(tree, state, thenStore, elseStore);
+      h.onDataflowVisitReturn(tree, thenStore, elseStore);
     }
   }
 

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

@@ -269,14 +269,12 @@ 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, VisitorState state, NullnessStore thenStore, NullnessStore elseStore);
+  void onDataflowVisitReturn(ReturnTree tree, 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,7 +27,6 @@
 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;
 
@@ -70,7 +69,6 @@ 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, VisitorState state, NullnessStore thenStore, NullnessStore elseStore) {
+      ReturnTree tree, 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,6 +40,7 @@
 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;
@@ -123,8 +124,46 @@
       VisitorState state,
       MethodTree tree,
       Symbol.MethodSymbol overriddenMethod) {
-    FieldContractUtils.ensureStrictPostConditionInheritance(
-        annotName, overridingFieldNames, analysis, state, tree, 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));
   }
 
   /**