Preserve contract unreachability across store joins

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

@@ -81,13 +81,19 @@ private AccessPathNullnessAnalysis(VisitorState state, NullAway analysis) {
             state,
             apContext,
             analysis,
-            new CoreNullnessStoreInitializer(analysis.getGenericsChecks()));
+            new CoreNullnessStoreInitializer(analysis.getGenericsChecks()),
+            /* trackUnreachableStores= */ false);
     this.dataFlow = new DataFlow(config.assertsEnabled(), handler);
 
     if (config.checkContracts()) {
       this.contractNullnessPropagation =
           new AccessPathNullnessPropagation(
-              Nullness.NONNULL, state, apContext, analysis, new ContractNullnessStoreInitializer());
+              Nullness.NONNULL,
+              state,
+              apContext,
+              analysis,
+              new ContractNullnessStoreInitializer(),
+              /* trackUnreachableStores= */ true);
     }
   }
 
@@ -145,20 +151,16 @@ public static AccessPathNullnessAnalysis instance(VisitorState state, NullAway a
   }
 
   /**
-   * Check if any access path in the store before an expression maps to {@link Nullness#BOTTOM} in
-   * contract dataflow.
+   * Check whether the store before an expression is unreachable in contract dataflow.
    *
    * @param exprPath tree path of expression
    * @param context Javac context
-   * @return true if any access path has {@link Nullness#BOTTOM} before the expression
+   * @return true if the store before the expression is unreachable
    */
-  public boolean hasBottomAccessPathForContractDataflow(TreePath exprPath, Context context) {
+  public boolean isUnreachableForContractDataflow(TreePath exprPath, Context context) {
     NullnessStore store =
         dataFlow.resultBeforeExpr(exprPath, context, castToNonNull(contractNullnessPropagation));
-    if (store == null) {
-      return false;
-    }
-    return !store.getAccessPathsWithValue(Nullness.BOTTOM).isEmpty();
+    return store != null && store.isUnreachable();
   }
 
   /**

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

@@ -177,6 +177,8 @@ public class AccessPathNullnessPropagation
 
   private final NullnessStoreInitializer nullnessStoreInitializer;
 
+  private final boolean trackUnreachableStores;
+
   /**
    * Updates the stored {@link VisitorState} to account for the fact that we are checking a new
    * compilation unit. Required since {@link VisitorState} objects internally contain a {@link
@@ -234,7 +236,8 @@ public AccessPathNullnessPropagation(
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       NullAway analysis,
-      NullnessStoreInitializer nullnessStoreInitializer) {
+      NullnessStoreInitializer nullnessStoreInitializer,
+      boolean trackUnreachableStores) {
     this.defaultAssumption = defaultAssumption;
     this.methodReturnsNonNull = analysis::isMethodUnannotated;
     // Overwrite the TreePath with a FailingTreePath to ensure it never gets used
@@ -244,6 +247,7 @@ public AccessPathNullnessPropagation(
     this.handler = analysis.getHandler();
     this.genericsChecks = analysis.getGenericsChecks();
     this.nullnessStoreInitializer = nullnessStoreInitializer;
+    this.trackUnreachableStores = trackUnreachableStores;
   }
 
   private static SubNodeValues values(TransferInput<Nullness, NullnessStore> input) {
@@ -1354,10 +1358,16 @@ public TransferResult<Nullness, NullnessStore> visitMarker(
   }
 
   @CheckReturnValue
-  private static ResultingStore updateStore(NullnessStore oldStore, ReadableUpdates... updates) {
+  private ResultingStore updateStore(NullnessStore oldStore, ReadableUpdates... updates) {
+    if (trackUnreachableStores && oldStore.isUnreachable()) {
+      return new ResultingStore(oldStore, NO_STORE_CHANGE);
+    }
     NullnessStore.Builder builder = oldStore.toBuilder();
     for (ReadableUpdates update : updates) {
       for (Map.Entry<AccessPath, Nullness> entry : update.values.entrySet()) {
+        if (trackUnreachableStores && entry.getValue() == BOTTOM) {
+          return new ResultingStore(NullnessStore.unreachable(), true);
+        }
         AccessPath key = entry.getKey();
         builder.setInformation(key, entry.getValue());
       }

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

@@ -45,6 +45,9 @@ public class NullnessStore implements Store<NullnessStore> {
 
   private static final NullnessStore EMPTY = new NullnessStore(ImmutableMap.of());
 
+  /** Canonical representation of a store for an unreachable control-flow path. */
+  private static final NullnessStore UNREACHABLE = new NullnessStore(ImmutableMap.of());
+
   private final ImmutableMap<AccessPath, Nullness> contents;
 
   private NullnessStore(Map<AccessPath, Nullness> contents) {
@@ -60,6 +63,25 @@ public static NullnessStore empty() {
     return EMPTY;
   }
 
+  /**
+   * Produce the canonical store for an unreachable control-flow path.
+   *
+   * @return the unreachable store
+   */
+  static NullnessStore unreachable() {
+    return UNREACHABLE;
+  }
+
+  /**
+   * Check whether this store represents an unreachable control-flow path.
+   *
+   * @return {@code true} if this is the canonical unreachable store
+   */
+  @SuppressWarnings("ReferenceEquality")
+  boolean isUnreachable() {
+    return this == UNREACHABLE;
+  }
+
   /**
    * Get the nullness for a local variable.
    *
@@ -160,6 +182,9 @@ public Nullness getNullnessOfAccessPath(AccessPath accessPath) {
   }
 
   public Builder toBuilder() {
+    if (isUnreachable()) {
+      throw new IllegalStateException("cannot create a builder from the unreachable store");
+    }
     return new Builder(this);
   }
 
@@ -174,11 +199,17 @@ public NullnessStore leastUpperBound(NullnessStore other) {
     if (this == other) {
       return this;
     }
+    if (this == UNREACHABLE) {
+      return other;
+    }
     ImmutableMap<AccessPath, Nullness> thisContents = this.contents;
     int thisContentsSize = thisContents.size();
     if (thisContentsSize == 0) {
       return this;
     }
+    if (other == UNREACHABLE) {
+      return this;
+    }
     ImmutableMap<AccessPath, Nullness> otherContents = other.contents;
     int otherContentsSize = otherContents.size();
     if (otherContentsSize == 0) {
@@ -210,24 +241,28 @@ public NullnessStore widenedUpperBound(NullnessStore vNullnessStore) {
   }
 
   @Override
+  @SuppressWarnings("ReferenceEquality")
   public boolean equals(@Nullable Object o) {
     if (this == o) {
       return true;
     }
     if (!(o instanceof NullnessStore other)) {
       return false;
     }
+    if (this == UNREACHABLE || other == UNREACHABLE) {
+      return false;
+    }
     return contents.equals(other.contents);
   }
 
   @Override
   public int hashCode() {
-    return contents.hashCode();
+    return isUnreachable() ? 31 : contents.hashCode();
   }
 
   @Override
   public String toString() {
-    return contents.toString();
+    return isUnreachable() ? "UNREACHABLE" : contents.toString();
   }
 
   @Override
@@ -257,6 +292,9 @@ public String visualize(CFGVisualizer<?, NullnessStore, ?> viz) {
    */
   public NullnessStore uprootAccessPaths(
       Map<LocalVariableNode, LocalVariableNode> localVarTranslations) {
+    if (isUnreachable()) {
+      return this;
+    }
     NullnessStore.Builder nullnessBuilder = NullnessStore.empty().toBuilder();
     for (AccessPath ap : contents.keySet()) {
       Element element = ap.getRoot();
@@ -282,6 +320,9 @@ public NullnessStore uprootAccessPaths(
    * @return NullnessStore containing only AccessPaths that pass the predicate
    */
   public NullnessStore filterAccessPaths(Predicate<AccessPath> pred) {
+    if (isUnreachable()) {
+      return this;
+    }
     return new NullnessStore(
         contents.entrySet().stream()
             .filter(e -> pred.test(e.getKey()))

nullaway/src/main/java/com/uber/nullaway/handlers/contract/ContractCheckHandler.java

@@ -173,9 +173,9 @@ public void onMatchMethod(MethodTree tree, MethodAnalysisContext methodAnalysisC
             Nullness nullness =
                 nullnessAnalysis.getNullnessForContractDataflow(expressionPath, state.context);
             if (nullness == Nullness.NULLABLE || nullness == Nullness.NULL) {
-              if (nullnessAnalysis.hasBottomAccessPathForContractDataflow(
+              if (nullnessAnalysis.isUnreachableForContractDataflow(
                   expressionPath, state.context)) {
-                // if any access path is mapped to bottom, this branch is unreachable
+                // The contract antecedent makes this branch unreachable.
                 continue;
               }
               contractViolated = true;

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

@@ -1226,6 +1226,81 @@ class Test {
         .doTest();
   }
 
+  @Test
+  public void checkContractUnreachableAfterControlFlowMerge() {
+    makeTestHelperWithArgs(
+            withJSpecifyModeArgs(
+                Arrays.asList(
+                    "-d",
+                    temporaryFolder.getRoot().getAbsolutePath(),
+                    "-XepOpt:NullAway:OnlyNullMarked=true",
+                    "-XepOpt:NullAway:CheckContracts=true")))
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.NullMarked;
+            import org.jspecify.annotations.Nullable;
+            import org.jetbrains.annotations.Contract;
+            @NullMarked
+            class Test {
+              @Contract("!null, !null -> !null")
+              static @Nullable Double conditionalOr(
+                  @Nullable Double a, @Nullable Double b) {
+                if (a == null || b == null) {
+                  return null;
+                }
+                return a + b;
+              }
+
+              @Contract("!null, !null -> !null")
+              static @Nullable Double conditionalAnd(
+                  @Nullable Double a, @Nullable Double b) {
+                if (a != null && b != null) {
+                  return a + b;
+                }
+                return null;
+              }
+
+              @Contract("!null, !null -> !null")
+              static @Nullable Double ternary(@Nullable Double a, @Nullable Double b) {
+                return (a == null || b == null) ? null : a + b;
+              }
+
+              @Contract("!null, !null -> !null")
+              static @Nullable Double twoIfs(@Nullable Double a, @Nullable Double b) {
+                if (a == null) {
+                  return null;
+                }
+                if (b == null) {
+                  return null;
+                }
+                return a + b;
+              }
+
+              @Contract("!null, _ -> !null")
+              static @Nullable Double sameVariable(
+                  @Nullable Double a, @Nullable Double unused) {
+                if (a == null || a == null) {
+                  return null;
+                }
+                return a + 1;
+              }
+
+              @Contract("!null, !null -> !null")
+              static @Nullable Double reachableMerge(
+                  @Nullable Double a, @Nullable Double b) {
+                if (a == null || b != null) {
+                  // BUG: Diagnostic contains: Method reachableMerge has @Contract
+                  return null;
+                }
+                return 0.0;
+              }
+
+            }
+            """)
+        .doTest();
+  }
+
   @Test
   public void checkDontCrashOnVoidReturn() {
     makeTestHelperWithArgs(

nullaway/src/test/java/com/uber/nullaway/dataflow/NullnessStoreTest.java

@@ -0,0 +1,28 @@
+package com.uber.nullaway.dataflow;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import org.junit.Test;
+
+public class NullnessStoreTest {
+
+  @Test
+  public void unreachableStoreLeastUpperBound() {
+    NullnessStore empty = NullnessStore.empty();
+    NullnessStore unreachable = NullnessStore.unreachable();
+
+    assertThat(unreachable.leastUpperBound(unreachable)).isSameInstanceAs(unreachable);
+    assertThat(unreachable.leastUpperBound(empty)).isSameInstanceAs(empty);
+    assertThat(empty.leastUpperBound(unreachable)).isSameInstanceAs(empty);
+  }
+
+  @Test
+  public void unreachableStoreIsDistinctFromEmptyStore() {
+    NullnessStore empty = NullnessStore.empty();
+    NullnessStore unreachable = NullnessStore.unreachable();
+
+    assertThat(unreachable.isUnreachable()).isTrue();
+    assertThat(empty.isUnreachable()).isFalse();
+    assertThat(unreachable).isNotEqualTo(empty);
+  }
+}