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 @@ static AccessPath switchRoot(AccessPath origAP, Element newRoot) {
   @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 @@ private static AccessPath fromNodeElementAndContext(
   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,24 @@ public static AccessPath getAccessPathForNode(
       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 {
       return null;
     }
@@ -350,7 +369,7 @@ public static AccessPath fromFieldElement(VariableElement element) {
     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 +403,31 @@ private static AccessPath buildAccessPathRecursive(
         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 = stripCasts(arrayAccess.getArray());
+      Node indexNode = arrayAccess.getIndex();
+      Element arrayElement = getElementFromArrayNode(arrayNode);
+      Element indexElement = getElementFromArrayNode(indexNode);
+      if (arrayElement == null) {
+        return null;
+      }
+      if (indexNode instanceof IntegerLiteralNode) {
+        IntegerLiteralNode intIndexNode = (IntegerLiteralNode) indexNode;
+        elements.push(new ArrayIndexElement(arrayElement, intIndexNode.getValue()));
+      } else {
+        if (indexElement != null) {
+          elements.push(new ArrayIndexElement(arrayElement, indexElement));
+        } else {
+          return null;
+        }
+      }
+      result = buildAccessPathRecursive(arrayNode, elements, apContext, mapKey);
     } else if (node instanceof MethodInvocationNode) {
       MethodInvocationNode invocation = (MethodInvocationNode) node;
       AccessPathElement accessPathElement;
@@ -399,7 +438,7 @@ private static AccessPath buildAccessPathRecursive(
           // 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 +507,8 @@ && isBoxingMethod(ASTHelpers.getSymbol(methodInvocationTree))) {
               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,30 @@
 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.
+ * Represents a generic element in an access path used for nullability analysis.
  *
- * <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")).
+ * <p>This interface abstracts over different kinds of path elements that can be part of an access
+ * path, including fields and methods, or array indices. Implementations of this interface should
+ * specify the type of the access path element:
+ *
+ * <ul>
+ *   <li>{@code FieldOrMethodCallElement} - Represents access to a field or the invocation of a
+ *       method, potentially with constant arguments.
+ *   <li>{@code ArrayIndexElement} - Represents access to an array element either by a constant
+ *       index or via an index that is calculated dynamically.
+ * </ul>
+ *
+ * <p>The {@code getJavaElement()} method returns the corresponding Java {@link Element} that the
+ * access path element refers to.
  */
-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;
-  }
-
-  @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;
-    }
-  }
-
-  @Override
-  public int hashCode() {
-    int result = javaElement.hashCode();
-    result = 31 * result + (constantArguments != null ? constantArguments.hashCode() : 0);
-    return result;
-  }
-
-  @Override
-  public String toString() {
-    return "APElement{"
-        + "javaElement="
-        + javaElement.toString()
-        + ", constantArguments="
-        + Arrays.deepToString(constantArguments != null ? constantArguments.toArray() : null)
-        + '}';
-  }
+public interface AccessPathElement {
+  /**
+   * Returns the Java element associated with this access path element.
+   *
+   * @return the Java {@link Element} related to this path element, such as a field, method, or the
+   *     array itself.
+   */
+  Element getJavaElement();
 }

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

@@ -788,17 +788,33 @@ 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);
       }
-    }
+      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.NULLABLE;
+        }
 
-    resultNullness = isElementNullable ? Nullness.NULLABLE : defaultAssumption;
+      } else {
+        resultNullness = Nullness.NONNULL;
+      }
+    } else {
+      resultNullness = Nullness.NONNULL;
+    }
     return updateRegularStore(resultNullness, input, updates);
   }
 

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

@@ -0,0 +1,60 @@
+package com.uber.nullaway.dataflow;
+
+import java.util.Objects;
+import javax.lang.model.element.Element;
+
+/**
+ * Represents an array index element of an AccessPath, encapsulating access to array elements either
+ * via constant or variable indices.
+ *
+ * <p>This class holds an element that represents the array itself and an index that specifies the
+ * position within the array. The index can be a constant (Integer) if it's statically known, or an
+ * Element representing a variable index.
+ */
+public class ArrayIndexElement implements AccessPathElement {
+  private final Element javaElement;
+  private final Object index;
+
+  /**
+   * Constructs an ArrayIndexElement.
+   *
+   * @param javaElement The element of the array.
+   * @param index The index used to access the array. Must be either an Integer (for constant
+   *     indices) or an Element (for variable indices).
+   */
+  public ArrayIndexElement(Element javaElement, Object index) {
+    this.javaElement = javaElement;
+    this.index = index;
+  }
+
+  @Override
+  public Element getJavaElement() {
+    return this.javaElement;
+  }
+
+  @Override
+  public String toString() {
+    return "ArrayIndexElement{"
+        + "javaElement="
+        + javaElement
+        + ", index="
+        + (index instanceof Element ? ((Element) index).getSimpleName() : index)
+        + '}';
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj instanceof ArrayIndexElement) {
+      ArrayIndexElement other = (ArrayIndexElement) obj;
+      return Objects.equals(javaElement, other.javaElement) && Objects.equals(index, other.index);
+    }
+    return false;
+  }
+
+  @Override
+  public int hashCode() {
+    int result = javaElement.hashCode();
+    result = 31 * result + (index != null ? index.hashCode() : 0);
+    return result;
+  }
+}

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

@@ -0,0 +1,62 @@
+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) field or method call element of an AccessPath.
+ *
+ * <p>This is just a java Element (field or method call) 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 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")
+        + '}';
+  }
+}