Fix comprehension variable mangling scopes in AstMutator

PiperOrigin-RevId: 1005927619

Author: l46kok

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(

optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java

@@ -846,6 +846,23 @@ 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(CEL_UNPARSER.unparse(mangledAst)).isEqualTo("cel.bind(x, x, @ac:0:0 + 1)");
+    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 {

optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java

@@ -368,6 +368,22 @@ public void cse_comprehensionReferencingOuterBindVar_notExtractedOutOfBindScope(
                 + " cel.bind(t, 99, @index0.exists(@it:0:0, @it:0:0 == @ac:1:0)))");
   }
 
+  @Test
+  @TestParameters("{source: '[cel.bind(x, x, x + 1), cel.bind(x, x, x + 1)]', expected: '[6, 6]'}")
+  @TestParameters(
+      "{source: '[[1, 2].map(k, cel.bind(k, k, k + x)), [1, 2].map(k, cel.bind(k, k, k + x))]',"
+          + " expected: '[[6, 7], [6, 7]]'}")
+  public void cse_bindInitReferencesOuterVariable_evaluatesCorrectly(String source, String expected)
+      throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile(source).getAst();
+    CelOptimizer celOptimizer = newCseOptimizer(SubexpressionOptimizerOptions.newBuilder().build());
+
+    CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
+
+    assertThat(cel.createProgram(optimizedAst).eval(ImmutableMap.of("x", 5L)))
+        .isEqualTo(cel.createProgram(cel.compile(expected).getAst()).eval());
+  }
+
   @Test
   public void celBlock_astExtensionTagged() throws Exception {
     CelAbstractSyntaxTree ast = cel.compile("size(x) + size(x)").getAst();