Draft: Prototype conditional nullness tracking for boolean variables (Fixes #98, #1060)

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

@@ -527,6 +527,23 @@ public TransferResult<Nullness, NullnessStore> visitAssignment(
       handleEnhancedForOverKeySet(localVariableNode, rhs, input, updates);
     }
 
+    if (target instanceof LocalVariableNode localVariableNode) {
+      com.sun.tools.javac.code.Type targetType = ASTHelpers.getType(target.getTree());
+
+      // NullAway requires us to prove targetType is not null before using it!
+      if (targetType != null && targetType.getTag() == com.sun.tools.javac.code.TypeTag.BOOLEAN) {
+        NullnessStore thenStore = input.getThenStore();
+        NullnessStore elseStore = input.getElseStore();
+
+        if (!thenStore.equals(elseStore)) {
+           AccessPath booleanAp = AccessPath.fromLocal(localVariableNode);
+           if (booleanAp != null) {
+               updates.setConditional(booleanAp, thenStore, elseStore);
+           }
+        }
+      }
+    }
+
     if (target instanceof ArrayAccessNode arrayAccessNode) {
       setNonnullIfAnalyzeable(updates, arrayAccessNode.getArray());
     }
@@ -1298,6 +1315,8 @@ public interface Updates {
     void set(MethodInvocationNode node, Nullness value);
 
     void set(AccessPath ap, Nullness value);
+
+    void setConditional(AccessPath booleanAp, NullnessStore ifTrue, NullnessStore ifFalse);
   }
 
   private final class ReadableUpdates implements Updates {
@@ -1338,5 +1357,10 @@ public void set(MethodInvocationNode node, Nullness value) {
     public void set(AccessPath ap, Nullness value) {
       values.put(checkNotNull(ap), value);
     }
+
+    @Override
+    public void setConditional(AccessPath booleanAp, NullnessStore ifTrue, NullnessStore ifFalse) {
+      // Stub for conditional update extraction to be passed to the store builder
+    }
   }
 }

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

@@ -42,12 +42,20 @@
  */
 public class NullnessStore implements Store<NullnessStore> {
 
-  private static final NullnessStore EMPTY = new NullnessStore(ImmutableMap.of());
+  private static final NullnessStore EMPTY = new NullnessStore(
+      ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of());
 
   private final ImmutableMap<AccessPath, Nullness> contents;
+  private final ImmutableMap<AccessPath, ImmutableMap<AccessPath, Nullness>> conditionalIfTrue;
+  private final ImmutableMap<AccessPath, ImmutableMap<AccessPath, Nullness>> conditionalIfFalse;
 
-  private NullnessStore(Map<AccessPath, Nullness> contents) {
+  private NullnessStore(
+      Map<AccessPath, Nullness> contents,
+      Map<AccessPath, ImmutableMap<AccessPath, Nullness>> conditionalIfTrue,
+      Map<AccessPath, ImmutableMap<AccessPath, Nullness>> conditionalIfFalse) {
     this.contents = ImmutableMap.copyOf(contents);
+    this.conditionalIfTrue = ImmutableMap.copyOf(conditionalIfTrue);
+    this.conditionalIfFalse = ImmutableMap.copyOf(conditionalIfFalse);
   }
 
   /**
@@ -193,7 +201,11 @@ public NullnessStore leastUpperBound(NullnessStore other) {
         upperBoundContentsBuilder.put(ap, smallValue.leastUpperBound(largeValue));
       }
     }
-    return new NullnessStore(upperBoundContentsBuilder.build());
+    return new NullnessStore(
+        upperBoundContentsBuilder.build(), 
+        ImmutableMap.of(), 
+        ImmutableMap.of()
+    );
   }
 
   @Override
@@ -277,7 +289,9 @@ public NullnessStore filterAccessPaths(Predicate<AccessPath> pred) {
     return new NullnessStore(
         contents.entrySet().stream()
             .filter(e -> pred.test(e.getKey()))
-            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)),
+        this.conditionalIfTrue,
+        this.conditionalIfFalse);
   }
 
   /**
@@ -334,19 +348,27 @@ public Set<Element> getReceiverFields(Nullness nullness) {
   /** class for building up instances of the store. */
   public static final class Builder {
     private final ImmutableMap.Builder<AccessPath, Nullness> contents;
+    private final ImmutableMap.Builder<AccessPath, ImmutableMap<AccessPath, Nullness>> conditionalIfTrue;
+    private final ImmutableMap.Builder<AccessPath, ImmutableMap<AccessPath, Nullness>> conditionalIfFalse;
 
     Builder(NullnessStore prototype) {
       contents = ImmutableMap.builder();
+      conditionalIfTrue = ImmutableMap.builder();
+      conditionalIfFalse = ImmutableMap.builder();
+
       if (!prototype.contents.isEmpty()) {
         contents.putAll(prototype.contents);
       }
+      if (!prototype.conditionalIfTrue.isEmpty()) {
+        conditionalIfTrue.putAll(prototype.conditionalIfTrue);
+      }
+      if (!prototype.conditionalIfFalse.isEmpty()) {
+        conditionalIfFalse.putAll(prototype.conditionalIfFalse);
+      }
     }
 
     /**
-     * Sets the value for the given variable. {@code element} must come from a call to {@link
-     * LocalVariableNode#getElement()} or {@link
-     * org.checkerframework.nullaway.javacutil.TreeUtils#elementFromDeclaration} ({@link
-     * org.checkerframework.nullaway.dataflow.cfg.node.VariableDeclarationNode#getTree()}).
+     * Sets the value for the given variable.
      *
      * @param ap relevant access path
      * @param value fact for access path
@@ -357,13 +379,26 @@ public NullnessStore.Builder setInformation(AccessPath ap, Nullness value) {
       return this;
     }
 
+    public NullnessStore.Builder setConditionalInformation(
+        AccessPath booleanAp, 
+        ImmutableMap<AccessPath, Nullness> ifTrue, 
+        ImmutableMap<AccessPath, Nullness> ifFalse) {
+      conditionalIfTrue.put(booleanAp, ifTrue);
+      conditionalIfFalse.put(booleanAp, ifFalse);
+      return this;
+    }
+
     /**
      * Construct the immutable NullnessStore instance.
      *
      * @return a store constructed from everything added to the builder
      */
     public NullnessStore build() {
-      return new NullnessStore(contents.buildKeepingLast());
+      return new NullnessStore(
+          contents.buildKeepingLast(),
+          conditionalIfTrue.buildKeepingLast(),
+          conditionalIfFalse.buildKeepingLast()
+      );
     }
   }
 }

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

@@ -1295,4 +1295,27 @@ static Inner testPositive(@Nullable Outer outer) {
             """)
         .doTest();
   }
+
+  @Test
+  public void testConditionalNullnessBoolean() {
+    defaultCompilationHelper
+      .addSourceLines(
+        "Test.java",
+        "package com.uber;",
+        "import javax.annotation.Nullable;",
+        "public class Test {",
+        "  public double execute(int a) {",
+        "    Object x = (a > 0) ? new Object() : null;",
+        "    boolean triggered = (a > 6);",
+        "    if (triggered) { x = null; }",
+        "    else { triggered = false; }",
+        "    if (!triggered) {",
+        "      // BUG: Diagnostic contains: dereferenced expression",
+        "      return x.hashCode();",
+        "    }",
+        "    return 0.0;",
+        "  }",
+        "}")
+      .doTest();
+  }
 }