Add a helper function for finding the declaring comprehension from an identifier

common/src/main/java/dev/cel/common/navigation/CelNavigableExprUtil.java

@@ -27,6 +27,58 @@
 @CheckReturnValue
 public final class CelNavigableExprUtil {
 
+  /**
+   * Returns the nearest enclosing comprehension that declares {@code variableName} in scope for
+   * {@code expr}, or {@code Optional.empty()} if none exists.
+   *
+   * <p>A comprehension declares {@code variableName} in scope for {@code expr} if {@code
+   * variableName} matches {@code iterVar}, {@code iterVar2}, or {@code accuVar}, and {@code expr}
+   * resides within the branch where that variable is active:
+   *
+   * <ul>
+   *   <li>In {@code loopCondition} and {@code loopStep}: {@code iterVar}, {@code iterVar2}, and
+   *       {@code accuVar} are in scope.
+   *   <li>In {@code result}: only {@code accuVar} is in scope.
+   *   <li>In {@code iterRange} and {@code accuInit}: none of the comprehension variables are in
+   *       scope.
+   * </ul>
+   */
+  @SuppressWarnings("ReferenceEquality") // Disambiguates mutable child branches
+  public static <E extends Expression, T extends BaseNavigableExpr<E>>
+      Optional<T> findDeclaringComprehension(T expr, String variableName) {
+    checkNotNull(expr);
+    checkNotNull(variableName);
+    if (variableName.isEmpty()) {
+      return Optional.empty();
+    }
+    T curr = expr;
+    Optional<T> maybeParent = curr.parent();
+    while (maybeParent.isPresent()) {
+      T parent = maybeParent.get();
+      if (parent.getKind() == Kind.COMPREHENSION) {
+        Expression.Comprehension<?> comp = parent.expr().comprehension();
+        Expression currExpr = curr.expr();
+
+        if (currExpr != comp.iterRange() && currExpr != comp.accuInit()) {
+          if (currExpr == comp.result()) {
+            if (comp.accuVar().equals(variableName)) {
+              return Optional.of(parent);
+            }
+          } else {
+            if (comp.iterVar().equals(variableName)
+                || comp.iterVar2().equals(variableName)
+                || comp.accuVar().equals(variableName)) {
+              return Optional.of(parent);
+            }
+          }
+        }
+      }
+      curr = parent;
+      maybeParent = parent.parent();
+    }
+    return Optional.empty();
+  }
+
   /**
    * Returns true if {@code variableName} is in scope and shadowed by an enclosing comprehension
    * above {@code expr}.
@@ -56,7 +108,7 @@ public final class CelNavigableExprUtil {
    * </ul>
    */
   public static boolean isVariableShadowed(BaseNavigableExpr<?> expr, String variableName) {
-    return areVariablesShadowed(expr, Collections.singleton(variableName));
+    return findDeclaringComprehension(expr, variableName).isPresent();
   }
 
   /**
@@ -72,38 +124,14 @@ public static boolean isVariableShadowed(BaseNavigableExpr<?> expr, String varia
    * At {@code y > 0}, {@code areVariablesShadowed(node, ImmutableSet.of("x", "z"))} is {@code true}
    * because {@code x} is in scope from the outer comprehension.
    */
-  @SuppressWarnings("ReferenceEquality") // Required to disambiguate child branches
   public static boolean areVariablesShadowed(
       BaseNavigableExpr<?> expr, Collection<String> variableNames) {
     checkNotNull(expr);
     checkNotNull(variableNames);
-    if (variableNames.isEmpty()) {
-      return false;
-    }
-    BaseNavigableExpr<?> curr = expr;
-    Optional<? extends BaseNavigableExpr<?>> maybeParent = curr.parent();
-    while (maybeParent.isPresent()) {
-      BaseNavigableExpr<?> parent = maybeParent.get();
-      if (parent.getKind() == Kind.COMPREHENSION) {
-        Expression.Comprehension<?> comp = parent.expr().comprehension();
-        Expression currExpr = curr.expr();
-
-        if (currExpr != comp.iterRange() && currExpr != comp.accuInit()) {
-          if (currExpr == comp.result()) {
-            if (variableNames.contains(comp.accuVar())) {
-              return true;
-            }
-          } else {
-            if (variableNames.contains(comp.iterVar())
-                || variableNames.contains(comp.iterVar2())
-                || variableNames.contains(comp.accuVar())) {
-              return true;
-            }
-          }
-        }
+    for (String varName : variableNames) {
+      if (findDeclaringComprehension(expr, varName).isPresent()) {
+        return true;
       }
-      curr = parent;
-      maybeParent = parent.parent();
     }
     return false;
   }

common/src/test/java/dev/cel/common/navigation/CelNavigableExprUtilTest.java

@@ -195,6 +195,76 @@ public void isVariableShadowed_comprehensionResultBranch() throws Exception {
         .isFalse();
   }
 
+  @Test
+  public void isVariableShadowed_twoVarComprehension_resultBranch() throws Exception {
+    CelAbstractSyntaxTree ast =
+        COMPILER.compile("{'k1': 1, 'k2': 2}.all(k, v, k != '' && v > 0)").getAst();
+    CelNavigableMutableAst navigableAst =
+        CelNavigableMutableAst.fromAst(CelMutableAst.fromCelAst(ast));
+
+    CelNavigableMutableExpr comprehensionNode =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(node -> node.getKind() == Kind.COMPREHENSION)
+            .findFirst()
+            .get();
+
+    CelMutableComprehension comprehension = comprehensionNode.expr().comprehension();
+    long resultId = comprehension.result().id();
+
+    CelNavigableMutableExpr resultNode =
+        comprehensionNode.allNodes().filter(node -> node.id() == resultId).findFirst().get();
+
+    // In result branch of two-var comprehension: accuVar is in scope, but iterVar and iterVar2 are
+    // not
+    assertThat(CelNavigableExprUtil.isVariableShadowed(resultNode, comprehension.accuVar()))
+        .isTrue();
+    assertThat(CelNavigableExprUtil.isVariableShadowed(resultNode, comprehension.iterVar()))
+        .isFalse();
+    assertThat(CelNavigableExprUtil.isVariableShadowed(resultNode, comprehension.iterVar2()))
+        .isFalse();
+  }
+
+  @Test
+  public void isVariableShadowed_accuInit_notShadowed() {
+    CelMutableExpr iterRange = CelMutableExpr.ofList(0, CelMutableList.create());
+    CelMutableExpr accuInitIdent = CelMutableExpr.ofIdent(1, "x");
+    CelMutableExpr loopCond = CelMutableExpr.ofConstant(2, CelConstant.ofValue(true));
+    CelMutableExpr loopStep = CelMutableExpr.ofConstant(3, CelConstant.ofValue(true));
+    CelMutableExpr result = CelMutableExpr.ofIdent(4, "accu");
+
+    CelMutableExpr comp =
+        CelMutableExpr.ofComprehension(
+            5,
+            CelMutableComprehension.create(
+                "x", iterRange, "accu", accuInitIdent, loopCond, loopStep, result));
+
+    CelNavigableMutableExpr root = CelNavigableMutableExpr.fromExpr(comp);
+    CelNavigableMutableExpr navAccuInit =
+        root.allNodes().filter(node -> node.id() == 1).findFirst().get();
+
+    assertThat(CelNavigableExprUtil.isVariableShadowed(navAccuInit, "x")).isFalse();
+    assertThat(CelNavigableExprUtil.isVariableShadowed(navAccuInit, "accu")).isFalse();
+  }
+
+  @Test
+  public void findDeclaringComprehension_emptyVariableName_returnsEmpty() throws Exception {
+    CelAbstractSyntaxTree ast = COMPILER.compile("[1, 2].all(x, x > 0)").getAst();
+    CelNavigableAst navigableAst = CelNavigableAst.fromAst(ast);
+
+    CelNavigableExpr identX =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(node -> node.expr().identOrDefault().name().equals("x"))
+            .findFirst()
+            .get();
+
+    assertThat(CelNavigableExprUtil.findDeclaringComprehension(identX, "")).isEmpty();
+    assertThat(CelNavigableExprUtil.isVariableShadowed(identX, "")).isFalse();
+  }
+
   @Test
   public void areVariablesShadowed_multipleVariables() throws Exception {
     CelAbstractSyntaxTree ast = COMPILER.compile("[1, 2].all(x, x > 0)").getAst();
@@ -345,4 +415,64 @@ public void isVariableShadowed_zeroedOutIds_scopedCorrectly() {
     assertThat(CelNavigableExprUtil.isVariableShadowed(navResult, "x")).isFalse();
     assertThat(CelNavigableExprUtil.isVariableShadowed(navResult, "accu")).isTrue();
   }
+
+  @Test
+  public void
+      findDeclaringComprehension_nestedComprehensions_resolvesToInnermostDeclaringComprehension()
+          throws Exception {
+    CelAbstractSyntaxTree ast =
+        COMPILER
+            .compile("[1, 2].all(x, {'k': 1}.exists(k, v, x > 0 && k != '' && v > 0))")
+            .getAst();
+    CelNavigableAst navigableAst = CelNavigableAst.fromAst(ast);
+
+    CelNavigableExpr outerComp =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(
+                node ->
+                    node.getKind() == Kind.COMPREHENSION
+                        && node.expr().comprehension().iterVar().equals("x"))
+            .findFirst()
+            .get();
+
+    CelNavigableExpr innerComp =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(
+                node ->
+                    node.getKind() == Kind.COMPREHENSION
+                        && node.expr().comprehension().iterVar().equals("k"))
+            .findFirst()
+            .get();
+
+    CelNavigableExpr identX =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(node -> node.expr().identOrDefault().name().equals("x"))
+            .findFirst()
+            .get();
+    CelNavigableExpr identK =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(node -> node.expr().identOrDefault().name().equals("k"))
+            .findFirst()
+            .get();
+    CelNavigableExpr identV =
+        navigableAst
+            .getRoot()
+            .allNodes()
+            .filter(node -> node.expr().identOrDefault().name().equals("v"))
+            .findFirst()
+            .get();
+
+    assertThat(CelNavigableExprUtil.findDeclaringComprehension(identX, "x")).hasValue(outerComp);
+    assertThat(CelNavigableExprUtil.findDeclaringComprehension(identK, "k")).hasValue(innerComp);
+    assertThat(CelNavigableExprUtil.findDeclaringComprehension(identV, "v")).hasValue(innerComp);
+    assertThat(CelNavigableExprUtil.findDeclaringComprehension(identX, "unknown")).isEmpty();
+  }
 }