Augment arrayaccess node

Co-authored-by: Sascha Kehrli <[email protected]>

Author: skehrli

dataflow/src/main/java/org/checkerframework/dataflow/cfg/builder/CFGTranslationPhaseOne.java

@@ -3123,6 +3123,7 @@ public Node visitEnhancedForLoop(EnhancedForLoopTree tree, Void p) {
       handleArtificialTree(arrayAccess);
       ArrayAccessNode arrayAccessNode = new ArrayAccessNode(arrayAccess, arrayNode2, indexNode2);
       arrayAccessNode.setArrayExpression(expression);
+      arrayAccessNode.setEnhancedForLoop(tree);
       arrayAccessNode.setInSource(false);
       extendWithNode(arrayAccessNode);
       AssignmentNode arrayAccessAssignNode =
@@ -3133,6 +3134,7 @@ public Node visitEnhancedForLoop(EnhancedForLoopTree tree, Void p) {
       Node arrayAccessAssignNodeExpr = arrayAccessAssignNode.getExpression();
       if (arrayAccessAssignNodeExpr instanceof ArrayAccessNode) {
         ((ArrayAccessNode) arrayAccessAssignNodeExpr).setArrayExpression(expression);
+        ((ArrayAccessNode) arrayAccessAssignNodeExpr).setEnhancedForLoop(tree);
       } else if (arrayAccessAssignNodeExpr instanceof MethodInvocationNode) {
         // If the array component type is a primitive, there may be a boxing or unboxing
         // conversion. Treat that as an iterator.

dataflow/src/main/java/org/checkerframework/dataflow/cfg/node/ArrayAccessNode.java

@@ -1,6 +1,7 @@
 package org.checkerframework.dataflow.cfg.node;
 
 import com.sun.source.tree.ArrayAccessTree;
+import com.sun.source.tree.EnhancedForLoopTree;
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.Tree;
 import java.util.Arrays;
@@ -39,6 +40,14 @@ public class ArrayAccessNode extends Node {
    */
   protected @Nullable ExpressionTree arrayExpression;
 
+  /**
+   * If this ArrayAccessNode is a node for an array access desugared from an enhanced for loop, then
+   * the {@code enhancedForLoop} field is the {@code EnhancedForLoopTree} AST node.
+   *
+   * <p>Is set by {@link #setEnhancedForLoop}.
+   */
+  protected @Nullable EnhancedForLoopTree enhancedForLoop;
+
   /**
    * Create an ArrayAccessNode.
    *
@@ -75,6 +84,27 @@ public void setArrayExpression(@Nullable ExpressionTree arrayExpression) {
     this.arrayExpression = arrayExpression;
   }
 
+  /**
+   * If this ArrayAccessNode is a node for an array access desugared from an enhanced for loop, then
+   * return the {@code EnhancedForLoopTree} AST node. Otherwise, return null.
+   *
+   * @return the {@code EnhancedForLoopTree}, or null if this is not an array access desugared from
+   *     an enhanced for loop
+   */
+  public @Nullable EnhancedForLoopTree getEnhancedForLoop() {
+    return enhancedForLoop;
+  }
+
+  /**
+   * Set the enhanced for loop from which {@code this} is desugared from.
+   *
+   * @param enhancedForLoop the {@code EnhancedForLoopTree}
+   * @see #getEnhancedForLoop()
+   */
+  public void setEnhancedForLoop(@Nullable EnhancedForLoopTree enhancedForLoop) {
+    this.enhancedForLoop = enhancedForLoop;
+  }
+
   /**
    * Get the node that represents the array expression being accessed.
    *