cel-expr · copybara-service · Aug 21, 2025 · Aug 21, 2025
diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
@@ -910,21 +910,32 @@ private static void unwrapListArgumentsInMacroCallExpr(

    CelMutableExpr loopStepExpr = comprehension.loopStep();
    List<CelMutableExpr> loopStepArgs = loopStepExpr.call().args();
-    if (loopStepArgs.size() != 2) {
+    if (loopStepArgs.size() != 2 && loopStepArgs.size() != 3) {
      throw new IllegalArgumentException(
          String.format(
-              "Expected exactly 2 arguments but got %d instead on expr id: %d",
+              "Expected exactly 2 or 3 arguments but got %d instead on expr id: %d",
              loopStepArgs.size(), loopStepExpr.id()));
    }

    CelMutableCall existingMacroCall = newMacroCallExpr.call();
    CelMutableCall newMacroCall =
        existingMacroCall.target().isPresent()
-            ? CelMutableCall.create(existingMacroCall.target().get(), existingMacroCall.function())
+            ? CelMutableCall.create(existingMacroCall.target().get(),
+            existingMacroCall.function())
            : CelMutableCall.create(existingMacroCall.function());
    newMacroCall.addArgs(
        existingMacroCall.args().get(0)); // iter_var is first argument of the call by convention
-    newMacroCall.addArgs(loopStepArgs.get(1).list().elements());
+
+    CelMutableList extraneousList = null;
+    if (loopStepArgs.size() == 2) {
+      extraneousList = loopStepArgs.get(1).list();
+    } else {
+      newMacroCall.addArgs(loopStepArgs.get(0));
+      // For map(x,y,z), z is wrapped in a _+_(@result, [z])
+      extraneousList = loopStepArgs.get(1).call().args().get(1).list();
+    }
+
+    newMacroCall.addArgs(extraneousList.elements());

    newMacroCallExpr.setCall(newMacroCall);
  }

diff --git a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java
@@ -581,6 +581,34 @@ public void replaceSubtree_replaceExtraneousListCreatedByMacro_unparseSuccess()
        .containsExactly(2L);
  }

+  @Test
+  @SuppressWarnings("unchecked") // Test only
+  public void replaceSubtree_replaceExtraneousListCreatedByThreeArgMacro_unparseSuccess()
+      throws Exception {
+    CelAbstractSyntaxTree ast = CEL.compile("[1].map(x, true, 1)").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+    CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast);
+
+    // These two mutation are equivalent.
+    CelAbstractSyntaxTree mutatedAstWithList =
+        AST_MUTATOR
+            .replaceSubtree(
+                mutableAst,
+                CelMutableExpr.ofList(
+                    CelMutableList.create(CelMutableExpr.ofConstant(CelConstant.ofValue(2L)))),
+                10L)
+            .toParsedAst();
+    CelAbstractSyntaxTree mutatedAstWithConstant =
+        AST_MUTATOR
+            .replaceSubtree(mutableAst2, CelMutableExpr.ofConstant(CelConstant.ofValue(2L)), 6L)
+            .toParsedAst();
+
+    assertThat(CEL_UNPARSER.unparse(mutatedAstWithList)).isEqualTo("[1].map(x, true, 2)");
+    assertThat(CEL_UNPARSER.unparse(mutatedAstWithConstant)).isEqualTo("[1].map(x, true, 2)");
+    assertThat((List<Long>) CEL.createProgram(CEL.check(mutatedAstWithList).getAst()).eval())
+        .containsExactly(2L);
+  }
+
  @Test
  public void globalCallExpr_replaceRoot() throws Exception {
    // Tree shape (brackets are expr IDs):