From 17e5df486e8ea3adf09db28b40353cff05ab3cc4 Mon Sep 17 00:00:00 2001
From: Sean Huh <sokwhan@google.com>
Date: Fri, 23 Oct 2026 15:56:33 -0700
Subject: [PATCH] Fix comprehension variable mangling scopes in AstMutator

PiperOrigin-RevId: 1005927619
---
 .../java/dev/cel/optimizer/AstMutator.java    | 66 ++++++++-----------
 .../dev/cel/optimizer/AstMutatorTest.java     | 16 +++++
 2 files changed, 43 insertions(+), 39 deletions(-)

diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
index 04db3ab99..1f14e01d7 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
@@ -677,55 +677,43 @@ private CelMutableExpr mangleIdentsInComprehensionExpr(
       String originalAccuVar,
       MangledComprehensionName mangledComprehensionName) {
     CelMutableComprehension comprehension = comprehensionExpr.comprehension();
-    replaceIdentName(
-        comprehension.loopStep(), originalIterVar, mangledComprehensionName.iterVarName());
-    replaceIdentName(comprehensionExpr, originalAccuVar, mangledComprehensionName.resultName());
-
+    // Iteration variables are in scope of the loop condition and step, the accumulator also of the
+    // result. Neither is in scope of the iteration range or the accumulator initializer.
+    renameInLoop(comprehension, originalIterVar, mangledComprehensionName.iterVarName());
     comprehension.setIterVar(mangledComprehensionName.iterVarName());
-
-    // Most standard macros set accu_var as @result, but not all (ex: cel.bind).
-    if (comprehension.accuVar().equals(originalAccuVar)) {
-      comprehension.setAccuVar(mangledComprehensionName.resultName());
-    }
-
     if (!originalIterVar2.isEmpty()) {
+      renameInLoop(comprehension, originalIterVar2, mangledComprehensionName.iterVar2Name());
       comprehension.setIterVar2(mangledComprehensionName.iterVar2Name());
-      replaceIdentName(
-          comprehension.loopStep(), originalIterVar2, mangledComprehensionName.iterVar2Name());
     }
+    String resultName = mangledComprehensionName.resultName();
+    renameInLoop(comprehension, originalAccuVar, resultName);
+    comprehension.setResult(replaceIdentName(comprehension.result(), originalAccuVar, resultName));
+    comprehension.setAccuVar(resultName);
 
     return mutateExpr(NO_OP_ID_GENERATOR, root, comprehensionExpr, comprehensionExpr.id());
   }
 
-  private void replaceIdentName(
-      CelMutableExpr comprehensionExpr, String originalIdentName, String newIdentName) {
-    int iterCount;
-    for (iterCount = 0; iterCount < iterationLimit; iterCount++) {
-      CelMutableExpr identToMangle =
-          CelNavigableMutableExpr.fromExpr(comprehensionExpr)
-              .descendants()
-              .map(CelNavigableMutableExpr::expr)
-              .filter(
-                  node ->
-                      node.getKind().equals(ExprKind.Kind.IDENT)
-                          && node.ident().name().equals(originalIdentName))
-              .findAny()
-              .orElse(null);
-      if (identToMangle == null) {
-        break;
-      }
-
-      comprehensionExpr =
-          mutateExpr(
-              NO_OP_ID_GENERATOR,
-              comprehensionExpr,
-              CelMutableExpr.ofIdent(newIdentName),
-              identToMangle.id());
-    }
+  private void renameInLoop(CelMutableComprehension comprehension, String from, String to) {
+    comprehension.setLoopCondition(replaceIdentName(comprehension.loopCondition(), from, to));
+    comprehension.setLoopStep(replaceIdentName(comprehension.loopStep(), from, to));
+  }
 
-    if (iterCount >= iterationLimit) {
-      throw new IllegalStateException("Max iteration count reached.");
+  // Returns the renamed expression, which is a new node if `expr` itself is a matching identifier.
+  private CelMutableExpr replaceIdentName(
+      CelMutableExpr expr, String originalIdentName, String newIdentName) {
+    ImmutableList<Long> identIds =
+        CelNavigableMutableExpr.fromExpr(expr)
+            .allNodes()
+            .filter(
+                node ->
+                    node.getKind().equals(ExprKind.Kind.IDENT)
+                        && node.expr().ident().name().equals(originalIdentName))
+            .map(CelNavigableMutableExpr::id)
+            .collect(toImmutableList());
+    for (long id : identIds) {
+      expr = mutateExpr(NO_OP_ID_GENERATOR, expr, CelMutableExpr.ofIdent(newIdentName), id);
     }
+    return expr;
   }
 
   private CelMutableSource mangleIdentsInMacroSource(
diff --git a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
index ef3dc85b6..6dbe5daa5 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
+++ b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
@@ -846,6 +846,22 @@ public void mangleComprehensionVariable_nonIdentComprehensionResult() throws Exc
     assertConsistentMacroCalls(mangledAst);
   }
 
+  @Test
+  public void mangleComprehensionVariable_bindInitReferencesOuterVariable() throws Exception {
+    // The `x` in the initializer refers to the outer variable, not to the one being bound.
+    CelAbstractSyntaxTree ast = CEL.compile("cel.bind(x, x, x + 1)").getAst();
+
+    CelAbstractSyntaxTree mangledAst =
+        AST_MUTATOR
+            .mangleComprehensionIdentifierNames(CelMutableAst.fromCelAst(ast), "@it", "@it2", "@ac")
+            .mutableAst()
+            .toParsedAst();
+
+    assertThat(mangledAst.getExpr().comprehension().loopStep().ident().name()).isEqualTo("@ac:0:0");
+    assertThat(CEL.createProgram(CEL.check(mangledAst).getAst()).eval(ImmutableMap.of("x", 5L)))
+        .isEqualTo(6L);
+  }
+
   @Test
   public void mangleComprehensionVariable_macroSourceDisabled_macroCallMapIsEmpty()
       throws Exception {