Fix replaceSubtree to properly populate three arg map macro source

PiperOrigin-RevId: 797536828

Author: l46kok

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);
   }

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):