JSpecify: Handle @Nonnull elements in @Nullable content arrays

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

@@ -46,6 +46,7 @@
 import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.VariableElement;
+import org.checkerframework.nullaway.dataflow.cfg.node.ArrayAccessNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.ClassNameNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.FieldAccessNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.IntegerLiteralNode;
@@ -206,7 +207,7 @@
   @Nullable
   public static AccessPath fromBaseAndElement(
       Node base, Element element, AccessPathContext apContext) {
-    return fromNodeElementAndContext(base, new AccessPathElement(element), apContext);
+    return fromNodeElementAndContext(base, new FieldOrMethodCallElement(element), apContext);
   }
 
   @Nullable
@@ -239,7 +240,7 @@
   public static AccessPath fromBaseMethodAndConstantArgs(
       Node base, Element method, List<String> constantArguments, AccessPathContext apContext) {
     return fromNodeElementAndContext(
-        base, new AccessPathElement(method, constantArguments), apContext);
+        base, new FieldOrMethodCallElement(method, constantArguments), apContext);
   }
 
   /**
@@ -334,6 +335,26 @@
       return fromFieldAccess((FieldAccessNode) node, apContext);
     } else if (node instanceof MethodInvocationNode) {
       return fromMethodCall((MethodInvocationNode) node, state, apContext);
+    } else if (node instanceof ArrayAccessNode) {
+      return fromArrayAccess((ArrayAccessNode) node, apContext);
+    } else {
+      return null;
+    }
+  }
+
+  @Nullable
+  private static AccessPath fromArrayAccess(ArrayAccessNode node, AccessPathContext apContext) {
+    return fromNodeAndContext(node, apContext);
+  }
+
+  @Nullable
+  private static Element getElementFromArrayNode(Node arrayNode) {
+    if (arrayNode instanceof LocalVariableNode) {
+      return ((LocalVariableNode) arrayNode).getElement();
+    } else if (arrayNode instanceof FieldAccessNode) {
+      return ((FieldAccessNode) arrayNode).getElement();
+    } else if (arrayNode instanceof MethodInvocationNode) {
+      return ASTHelpers.getSymbol(((MethodInvocationNode) arrayNode).getTree());
     } else {
       return null;
     }
@@ -350,7 +371,7 @@
     Preconditions.checkArgument(
         element.getKind().isField(),
         "element must be of type: FIELD but received: " + element.getKind());
-    return new AccessPath(null, ImmutableList.of(new AccessPathElement(element)));
+    return new AccessPath(null, ImmutableList.of(new FieldOrMethodCallElement(element)));
   }
 
   private static boolean isBoxingMethod(Symbol.MethodSymbol methodSymbol) {
@@ -384,11 +405,31 @@
         result = new AccessPath(fieldAccess.getElement(), ImmutableList.copyOf(elements), mapKey);
       } else {
         // instance field access
-        elements.push(new AccessPathElement(fieldAccess.getElement()));
+        elements.push(new FieldOrMethodCallElement(fieldAccess.getElement()));
         result =
             buildAccessPathRecursive(
                 stripCasts(fieldAccess.getReceiver()), elements, apContext, mapKey);
       }
+    } else if (node instanceof ArrayAccessNode) {
+      ArrayAccessNode arrayAccess = (ArrayAccessNode) node;
+      Node arrayNode = arrayAccess.getArray();
+      Node indexNode = arrayAccess.getIndex();
+      Element arrayElement = getElementFromArrayNode(arrayNode);
+      if (arrayElement == null) {
+        return null;
+      }
+      if (indexNode instanceof IntegerLiteralNode) {
+        IntegerLiteralNode intIndexNode = (IntegerLiteralNode) indexNode;
+        elements.push(new ArrayIndexElement(arrayElement, intIndexNode.getValue()));
+      } else {
+        Element indexElement = getElementFromArrayNode(indexNode);
+        if (indexElement != null) {
+          elements.push(new ArrayIndexElement(arrayElement, indexElement));
+        } else {
+          return null;
+        }
+      }
+      result = buildAccessPathRecursive(stripCasts(arrayNode), elements, apContext, mapKey);
     } else if (node instanceof MethodInvocationNode) {
       MethodInvocationNode invocation = (MethodInvocationNode) node;
       AccessPathElement accessPathElement;
@@ -399,7 +440,7 @@
           // a zero-argument static method call can be the root of an access path
           return new AccessPath(symbol, ImmutableList.copyOf(elements), mapKey);
         } else {
-          accessPathElement = new AccessPathElement(accessNode.getMethod());
+          accessPathElement = new FieldOrMethodCallElement(accessNode.getMethod());
         }
       } else {
         List<String> constantArgumentValues = new ArrayList<>();
@@ -468,7 +509,8 @@
               return null; // Not an AP
           }
         }
-        accessPathElement = new AccessPathElement(accessNode.getMethod(), constantArgumentValues);
+        accessPathElement =
+            new FieldOrMethodCallElement(accessNode.getMethod(), constantArgumentValues);
       }
       elements.push(accessPathElement);
       result =

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

@@ -1,62 +1,16 @@
 package com.uber.nullaway.dataflow;
 
-import com.google.common.collect.ImmutableList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Objects;
-import javax.annotation.Nullable;
 import javax.lang.model.element.Element;
 
-/**
- * Represents a (non-root) element of an AccessPath.
- *
- * <p>This is just a java Element (field, method, etc) in the access-path chain (e.g. f or g() in
- * x.f.g()). Plus, optionally, a list of constant arguments, allowing access path elements for
- * method calls with constant values (e.g. h(3) or k("STR_KEY") in x.h(3).g().k("STR_KEY")).
- */
-public final class AccessPathElement {
-  private final Element javaElement;
-  @Nullable private final ImmutableList<String> constantArguments;
-
-  public AccessPathElement(Element javaElement, List<String> constantArguments) {
-    this.javaElement = javaElement;
-    this.constantArguments = ImmutableList.copyOf(constantArguments);
-  }
-
-  public AccessPathElement(Element javaElement) {
-    this.javaElement = javaElement;
-    this.constantArguments = null;
-  }
-
-  public Element getJavaElement() {
-    return this.javaElement;
-  }
+public interface AccessPathElement {
+  Element getJavaElement();
 
   @Override
-  public boolean equals(Object obj) {
-    if (obj instanceof AccessPathElement) {
-      AccessPathElement otherNode = (AccessPathElement) obj;
-      return this.javaElement.equals(otherNode.javaElement)
-          && Objects.equals(constantArguments, otherNode.constantArguments);
-    } else {
-      return false;
-    }
-  }
+  String toString();
 
   @Override
-  public int hashCode() {
-    int result = javaElement.hashCode();
-    result = 31 * result + (constantArguments != null ? constantArguments.hashCode() : 0);
-    return result;
-  }
+  boolean equals(Object obj);
 
   @Override
-  public String toString() {
-    return "APElement{"
-        + "javaElement="
-        + javaElement.toString()
-        + ", constantArguments="
-        + Arrays.deepToString(constantArguments != null ? constantArguments.toArray() : null)
-        + '}';
-  }
+  int hashCode();
 }

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

@@ -788,17 +788,30 @@ public TransferResult<Nullness, NullnessStore> visitArrayAccess(
       ArrayAccessNode node, TransferInput<Nullness, NullnessStore> input) {
     ReadableUpdates updates = new ReadableUpdates();
     setNonnullIfAnalyzeable(updates, node.getArray());
-    Nullness resultNullness;
+    Nullness resultNullness = defaultAssumption;
     // Unsoundly assume @NonNull, except in JSpecify mode where we check the type
-    boolean isElementNullable = false;
     if (config.isJSpecifyMode()) {
       Symbol arraySymbol = ASTHelpers.getSymbol(node.getArray().getTree());
+      boolean isElementNullable = false;
       if (arraySymbol != null) {
         isElementNullable = NullabilityUtil.isArrayElementNullable(arraySymbol, config);
       }
-    }
 
-    resultNullness = isElementNullable ? Nullness.NULLABLE : defaultAssumption;
+      if (isElementNullable) {
+        AccessPath arrayAccessPath = AccessPath.getAccessPathForNode(node, state, apContext);
+        if (arrayAccessPath != null) {
+          @Nullable
+          Nullness accessPathNullness =
+              input.getRegularStore().getNullnessOfAccessPath(arrayAccessPath);
+          if (accessPathNullness == Nullness.NULLABLE) {
+            resultNullness = Nullness.NULLABLE;
+          }
+        }
+
+      } else {
+        resultNullness = Nullness.NONNULL;
+      }
+    }
     return updateRegularStore(resultNullness, input, updates);
   }
 

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

@@ -0,0 +1,61 @@
+package com.uber.nullaway.dataflow;
+
+import java.util.Objects;
+import javax.annotation.Nullable;
+import javax.lang.model.element.Element;
+
+public class ArrayIndexElement implements AccessPathElement {
+  private final Element javaElement;
+  @Nullable private final Integer constantIndex;
+  @Nullable private final Element variableElement;
+
+  public ArrayIndexElement(Element javaElement, Integer constantIndex) {
+    this.javaElement = javaElement;
+    this.constantIndex = constantIndex;
+    this.variableElement = null;
+  }
+
+  public ArrayIndexElement(Element javaElement, Element variableElement) {
+    this.javaElement = javaElement;
+    this.variableElement = variableElement;
+    this.constantIndex = null;
+  }
+
+  @Override
+  public Element getJavaElement() {
+    return this.javaElement;
+  }
+
+  @Override
+  public String toString() {
+    return "ArrayIndexElement{"
+        + "javaElement="
+        + javaElement
+        + ", constantIndex="
+        + constantIndex
+        + ", variableElement="
+        + (variableElement != null ? variableElement.getSimpleName() : null)
+        + '}';
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj instanceof ArrayIndexElement) {
+      ArrayIndexElement other = (ArrayIndexElement) obj;
+      return Objects.equals(javaElement, other.javaElement)
+          && Objects.equals(constantIndex, other.constantIndex)
+          && Objects.equals(variableElement, other.variableElement);
+    }
+    return false;
+  }
+
+  @Override
+  public int hashCode() {
+    int result = javaElement.hashCode();
+    result =
+        31 * result
+            + (constantIndex != null ? constantIndex.hashCode() : 0)
+            + (variableElement != null ? variableElement.hashCode() : 0);
+    return result;
+  }
+}

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

@@ -0,0 +1,55 @@
+package com.uber.nullaway.dataflow;
+
+import com.google.common.collect.ImmutableList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import javax.lang.model.element.Element;
+
+public class FieldOrMethodCallElement implements AccessPathElement {
+  private final Element javaElement;
+  @Nullable private final ImmutableList<String> constantArguments;
+
+  public FieldOrMethodCallElement(Element javaElement, List<String> constantArguments) {
+    this.javaElement = javaElement;
+    this.constantArguments = ImmutableList.copyOf(constantArguments);
+  }
+
+  public FieldOrMethodCallElement(Element javaElement) {
+    this.javaElement = javaElement;
+    this.constantArguments = null;
+  }
+
+  @Override
+  public Element getJavaElement() {
+    return this.javaElement;
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj instanceof FieldOrMethodCallElement) {
+      FieldOrMethodCallElement other = (FieldOrMethodCallElement) obj;
+      return this.javaElement.equals(other.javaElement)
+          && Objects.equals(this.constantArguments, other.constantArguments);
+    }
+    return false;
+  }
+
+  @Override
+  public int hashCode() {
+    int result = javaElement.hashCode();
+    result = 31 * result + (constantArguments != null ? constantArguments.hashCode() : 0);
+    return result;
+  }
+
+  @Override
+  public String toString() {
+    return "FieldOrMethodCallElement{"
+        + "javaElement="
+        + javaElement
+        + ", constantArguments="
+        + (constantArguments != null ? Arrays.toString(constantArguments.toArray()) : "null")
+        + '}';
+  }
+}