Fix two-variable comprehension cel.@mapInsert overload IDs

PiperOrigin-RevId: 999641973

Author: l46kok

checker/src/test/resources/standardEnvDump.baseline

@@ -5,8 +5,8 @@ Source: 'redundant expression so the env is constructed and can be printed'
 
 Standard environment:
 declare cel.@mapInsert {
-  function cel_@mapInsert_map_map (map(K, V), map(K, V)) -> map(K, V)
-  function cel_@mapInsert_map_key_value (map(K, V), K, V) -> map(K, V)
+  function @mapInsert_map_map (map(K, V), map(K, V)) -> map(K, V)
+  function @mapInsert_map_key_value (map(K, V), K, V) -> map(K, V)
 }
 declare !_ {
   function logical_not (bool) -> bool

extensions/src/main/java/dev/cel/extensions/CelComprehensionsCompilerLibrary.java

@@ -42,8 +42,8 @@ public final class CelComprehensionsCompilerLibrary
     implements CelCompilerLibrary, CelExtensionLibrary.FeatureSet {
 
   private static final String MAP_INSERT_FUNCTION = "cel.@mapInsert";
-  private static final String MAP_INSERT_OVERLOAD_MAP_MAP = "cel_@mapInsert_map_map";
-  private static final String MAP_INSERT_OVERLOAD_KEY_VALUE = "cel_@mapInsert_map_key_value";
+  private static final String MAP_INSERT_OVERLOAD_MAP_MAP = "@mapInsert_map_map";
+  private static final String MAP_INSERT_OVERLOAD_KEY_VALUE = "@mapInsert_map_key_value";
 
   private static final class Types {
     private static final TypeParamType TYPE_PARAM_K = TypeParamType.create("K");

extensions/src/main/java/dev/cel/extensions/CelComprehensionsRuntimeLibrary.java

@@ -34,8 +34,10 @@
 public final class CelComprehensionsRuntimeLibrary implements CelInternalLiteRuntimeLibrary {
 
   private static final String MAP_INSERT_FUNCTION = "cel.@mapInsert";
-  private static final String MAP_INSERT_OVERLOAD_MAP_MAP = "cel_@mapInsert_map_map";
-  private static final String MAP_INSERT_OVERLOAD_KEY_VALUE = "cel_@mapInsert_map_key_value";
+  private static final String MAP_INSERT_OVERLOAD_MAP_MAP = "@mapInsert_map_map";
+  private static final String MAP_INSERT_OVERLOAD_KEY_VALUE = "@mapInsert_map_key_value";
+  private static final String LEGACY_MAP_INSERT_OVERLOAD_MAP_MAP = "cel_@mapInsert_map_map";
+  private static final String LEGACY_MAP_INSERT_OVERLOAD_KEY_VALUE = "cel_@mapInsert_map_key_value";
 
   /** Enumeration of functions for Comprehensions runtime extension. */
   public enum Function {
@@ -126,6 +128,16 @@ public ImmutableList<CelFunctionBinding> newFunctionBindings(RuntimeEquality run
               CelFunctionBinding.from(
                   MAP_INSERT_OVERLOAD_KEY_VALUE,
                   ImmutableList.of(Map.class, Object.class, Object.class),
+                  args -> mapInsertKeyValue(args, runtimeEquality)),
+              // Legacy overload IDs retained to support evaluation of cached ASTs (b/573697356).
+              CelFunctionBinding.from(
+                  LEGACY_MAP_INSERT_OVERLOAD_MAP_MAP,
+                  Map.class,
+                  Map.class,
+                  (map1, map2) -> mapInsertMap(map1, map2, runtimeEquality)),
+              CelFunctionBinding.from(
+                  LEGACY_MAP_INSERT_OVERLOAD_KEY_VALUE,
+                  ImmutableList.of(Map.class, Object.class, Object.class),
                   args -> mapInsertKeyValue(args, runtimeEquality))));
     }
     return bindings.build();

extensions/src/test/java/dev/cel/extensions/CelComprehensionsExtensionsTest.java

@@ -21,6 +21,7 @@
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import com.google.testing.junit.testparameterinjector.TestParameters;
@@ -31,6 +32,7 @@
 import dev.cel.common.CelOptions;
 import dev.cel.common.CelValidationException;
 import dev.cel.common.CelValidationResult;
+import dev.cel.common.ast.CelReference;
 import dev.cel.common.exceptions.CelAttributeNotFoundException;
 import dev.cel.common.exceptions.CelDivideByZeroException;
 import dev.cel.common.exceptions.CelIndexOutOfBoundsException;
@@ -552,4 +554,36 @@ public void runtimeOptions_unsupported_throws() {
         UnsupportedOperationException.class,
         () -> runtimeLibrary.setRuntimeOptions(liteRuntimeBuilder));
   }
+
+  @Test
+  @TestParameters(
+      "{expr: '[1, 2].transformMap(i, v, v + 1) == {0: 2, 1: 3}',"
+          + " canonicalOverloadId: '@mapInsert_map_key_value'}")
+  @TestParameters(
+      "{expr: \"{'a': 1, 'b': 2}.transformMapEntry(k, v, {k: v + 1}) == {'a': 2, 'b': 3}\","
+          + " canonicalOverloadId: '@mapInsert_map_map'}")
+  public void evaluate_legacyMapInsertOverloadIds_success(String expr, String canonicalOverloadId)
+      throws Exception {
+    Assume.assumeFalse(isParseOnly);
+    CelAbstractSyntaxTree ast = cel.compile(expr).getAst();
+    CelAbstractSyntaxTree legacyAst =
+        CelAbstractSyntaxTree.newCheckedAst(
+            ast.getExpr(),
+            ast.getSource(),
+            Maps.transformValues(
+                ast.getReferenceMap(),
+                ref ->
+                    ref.overloadIds().contains(canonicalOverloadId)
+                        ? CelReference.newBuilder()
+                            .addOverloadIds("cel_" + canonicalOverloadId)
+                            .build()
+                        : ref),
+            ast.getTypeMap());
+
+    Object result = cel.createProgram(legacyAst).eval();
+
+    assertThat(ast.getReferenceMap().values())
+        .contains(CelReference.newBuilder().addOverloadIds(canonicalOverloadId).build());
+    assertThat(result).isEqualTo(true);
+  }
 }

runtime/src/test/java/dev/cel/runtime/planner/ProgramPlannerTest.java

@@ -255,7 +255,7 @@ private static DefaultDispatcher newDispatcher(CelOptions celOptions) {
         CelFunctionBinding.fromOverloads(
             "cel.@mapInsert",
             CelFunctionBinding.from(
-                "cel_@mapInsert_map_key_value",
+                "@mapInsert_map_key_value",
                 ImmutableList.of(Map.class, Object.class, Object.class),
                 args -> args[0])));
 

verifier/src/main/java/dev/cel/verifier/axioms/MapInsertAxiom.java

@@ -29,7 +29,7 @@ final class MapInsertAxiom {
   static final CelZ3FunctionAxiom INSTANCE =
       CelZ3FunctionAxiom.newBuilder(CelComprehensionsExtensions.Function.MAP_INSERT.functionDecl())
           .addTernaryOverloadTranslator(
-              "cel_@mapInsert_map_key_value",
+              "@mapInsert_map_key_value",
               (ctx, typeSystem, constraintSink, mapVal, keyVal, valueVal) -> {
                 BoolExpr isMap = typeSystem.isMap(mapVal);
                 BoolExpr isValidKey =