Addressed PR comments

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

@@ -353,8 +353,6 @@ private static Element getElementFromArrayNode(Node arrayNode) {
       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;
     }
@@ -412,24 +410,24 @@ private static AccessPath buildAccessPathRecursive(
       }
     } else if (node instanceof ArrayAccessNode) {
       ArrayAccessNode arrayAccess = (ArrayAccessNode) node;
-      Node arrayNode = arrayAccess.getArray();
+      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 {
-        Element indexElement = getElementFromArrayNode(indexNode);
         if (indexElement != null) {
           elements.push(new ArrayIndexElement(arrayElement, indexElement));
         } else {
           return null;
         }
       }
-      result = buildAccessPathRecursive(stripCasts(arrayNode), elements, apContext, mapKey);
+      result = buildAccessPathRecursive(arrayNode, elements, apContext, mapKey);
     } else if (node instanceof MethodInvocationNode) {
       MethodInvocationNode invocation = (MethodInvocationNode) node;
       AccessPathElement accessPathElement;

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

@@ -2,15 +2,29 @@
 
 import javax.lang.model.element.Element;
 
+/**
+ * Represents a generic element in an access path used for nullability analysis.
+ *
+ * <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 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();
-
-  @Override
-  String toString();
-
-  @Override
-  boolean equals(Object obj);
-
-  @Override
-  int hashCode();
 }

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

@@ -796,7 +796,6 @@ public TransferResult<Nullness, NullnessStore> visitArrayAccess(
       if (arraySymbol != null) {
         isElementNullable = NullabilityUtil.isArrayElementNullable(arraySymbol, config);
       }
-
       if (isElementNullable) {
         AccessPath arrayAccessPath = AccessPath.getAccessPathForNode(node, state, apContext);
         if (arrayAccessPath != null) {
@@ -806,11 +805,15 @@ public TransferResult<Nullness, NullnessStore> visitArrayAccess(
           if (accessPathNullness == Nullness.NULLABLE) {
             resultNullness = Nullness.NULLABLE;
           }
+        } else {
+          resultNullness = Nullness.NULLABLE;
         }
 
       } else {
         resultNullness = Nullness.NONNULL;
       }
+    } else {
+      resultNullness = Nullness.NONNULL;
     }
     return updateRegularStore(resultNullness, input, updates);
   }

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

@@ -1,24 +1,30 @@
 package com.uber.nullaway.dataflow;
 
 import java.util.Objects;
-import javax.annotation.Nullable;
 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;
-  @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) {
+  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.variableElement = variableElement;
-    this.constantIndex = null;
+    this.index = index;
   }
 
   @Override
@@ -31,31 +37,24 @@ public String toString() {
     return "ArrayIndexElement{"
         + "javaElement="
         + javaElement
-        + ", constantIndex="
-        + constantIndex
-        + ", variableElement="
-        + (variableElement != null ? variableElement.getSimpleName() : null)
+        + ", 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(constantIndex, other.constantIndex)
-          && Objects.equals(variableElement, other.variableElement);
+      return Objects.equals(javaElement, other.javaElement) && Objects.equals(index, other.index);
     }
     return false;
   }
 
   @Override
   public int hashCode() {
     int result = javaElement.hashCode();
-    result =
-        31 * result
-            + (constantIndex != null ? constantIndex.hashCode() : 0)
-            + (variableElement != null ? variableElement.hashCode() : 0);
+    result = 31 * result + (index != null ? index.hashCode() : 0);
     return result;
   }
 }

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

@@ -7,6 +7,13 @@
 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;

nullaway/src/test/java/com/uber/nullaway/jspecify/ArrayTests.java

@@ -368,6 +368,8 @@ public void fieldAccessIndexArray() {
             "  static final Integer i = 0;",
             "  static void foo() {",
             "  if (fizz[i]!=null) { ",
+            "   // field access indexes aren't handled",
+            "   // BUG: Diagnostic contains: dereferenced expression fizz[i] is @Nullable",
             "   fizz[i].toString();",
             "}",
             "    // BUG: Diagnostic contains: dereferenced expression fizz[i] is @Nullable",
@@ -388,6 +390,7 @@ public void constantIndexArray() {
             "  static @Nullable String [] fizz = {\"1\"};",
             "  static void foo() {",
             "  if (fizz[0]!=null) { ",
+            "    // OK: constant integer indexes are handled by dataflow",
             "   fizz[0].toString();",
             "}",
             "    // BUG: Diagnostic contains: dereferenced expression fizz[0] is @Nullable",
@@ -409,6 +412,7 @@ public void localVariableIndexArray() {
             "  static void foo() {",
             "    int index = 1;",
             "    if (fizz[index] != null) {",
+            "    // OK: local variable indexes are handled by dataflow",
             "      fizz[index].toString();",
             "    }",
             "    // BUG: Diagnostic contains: dereferenced expression fizz[index] is @Nullable",
@@ -432,6 +436,8 @@ public void methodInvocationIndexArray() {
             "  }",
             "  static void foo() {",
             "    if (fizz[getIndex()] != null) {",
+            "    // index methods aren't handled by dataflow",
+            "    // BUG: Diagnostic contains: dereferenced expression fizz[getIndex()] is @Nullable",
             "      fizz[getIndex()].toString();",
             "    }",
             "    // BUG: Diagnostic contains: dereferenced expression fizz[getIndex()] is @Nullable",
@@ -442,18 +448,64 @@ public void methodInvocationIndexArray() {
   }
 
   @Test
-  public void forEachLoopOverArray() {
+  public void arithmeticIndexArray() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  static @Nullable String[] fizz = {\"1\", null};",
+            "  static void foo() {",
+            "    int i = 0;",
+            "    if (fizz[i+1] != null) {",
+            "    // index expressions aren't handled by dataflow",
+            "    // BUG: Diagnostic contains: dereferenced expression fizz[i+1] is @Nullable",
+            "      fizz[i+1].toString();",
+            "    }",
+            "    // BUG: Diagnostic contains: dereferenced expression fizz[i+1] is @Nullable",
+            "    fizz[i+1].toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void arrayMethodInvocationIndex() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  static @Nullable String[] getArray() { return new String[] {\"1\", null}; }",
+            "  static void foo() {",
+            "    if (getArray()[0] != null) {",
+            "    // array resulting from method invocation isn't handled by dataflow",
+            "    // BUG: Diagnostic contains: dereferenced expression getArray()[0] is @Nullable",
+            "      getArray()[0].toString();",
+            "    }",
+            "    // BUG: Diagnostic contains: dereferenced expression getArray()[0] is @Nullable",
+            "    getArray()[0].toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void mismatchedIndexUse() {
     makeHelper()
         .addSourceLines(
             "Test.java",
             "package com.uber;",
             "import org.jspecify.annotations.Nullable;",
-            "import java.util.List;",
             "class Test {",
-            "  static String[] arr = {\"1\"};",
+            "  static @Nullable String[] fizz = {\"1\", null};",
             "  static void foo() {",
-            "    for (String s : arr) {",
-            "      s.toString();",
+            "    int i = 0;",
+            "    if (fizz[i] != null) {",
+            "      // BUG: Diagnostic contains: dereferenced expression fizz[i+1] is @Nullable",
+            "      fizz[i+1].toString();",
             "    }",
             "  }",
             "}")