From 238e276df6ad4d7deb7c9196ebf2a4c6904811de Mon Sep 17 00:00:00 2001
From: Sascha Kehrli <84683849+skehrli@users.noreply.github.com>
Date: Tue, 8 Oct 2024 20:03:26 +0200
Subject: [PATCH] Augment arrayaccess node

Co-authored-by: Sascha Kehrli <skehrli@ethz.ch>
---
 .../cfg/builder/CFGTranslationPhaseOne.java   |  2 ++
 .../dataflow/cfg/node/ArrayAccessNode.java    | 30 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/dataflow/src/main/java/org/checkerframework/dataflow/cfg/builder/CFGTranslationPhaseOne.java b/dataflow/src/main/java/org/checkerframework/dataflow/cfg/builder/CFGTranslationPhaseOne.java
index 1207988f0b2..8bb453ac57c 100644
--- a/dataflow/src/main/java/org/checkerframework/dataflow/cfg/builder/CFGTranslationPhaseOne.java
+++ b/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.
diff --git a/dataflow/src/main/java/org/checkerframework/dataflow/cfg/node/ArrayAccessNode.java b/dataflow/src/main/java/org/checkerframework/dataflow/cfg/node/ArrayAccessNode.java
index e2113f5c304..cc5e15f8aff 100644
--- a/dataflow/src/main/java/org/checkerframework/dataflow/cfg/node/ArrayAccessNode.java
+++ b/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.
    *