Implement mutable map support for planner

PiperOrigin-RevId: 901433988

Author: l46kok

common/src/main/java/dev/cel/common/values/BUILD.bazel

@@ -118,6 +118,38 @@ java_library(
     ],
 )
 
+java_library(
+    name = "mutable_map_value",
+    srcs = ["MutableMapValue.java"],
+    tags = [
+    ],
+    deps = [
+        "//common/annotations",
+        "//common/exceptions:attribute_not_found",
+        "//common/types",
+        "//common/types:type_providers",
+        "//common/values",
+        "//common/values:cel_value",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
+)
+
+cel_android_library(
+    name = "mutable_map_value_android",
+    srcs = ["MutableMapValue.java"],
+    tags = [
+    ],
+    deps = [
+        ":cel_value_android",
+        "//common/annotations",
+        "//common/exceptions:attribute_not_found",
+        "//common/types:type_providers_android",
+        "//common/types:types_android",
+        "//common/values:values_android",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
+)
+
 cel_android_library(
     name = "values_android",
     srcs = CEL_VALUES_SOURCES,

common/src/main/java/dev/cel/common/values/MutableMapValue.java

@@ -0,0 +1,146 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.common.values;
+
+import com.google.errorprone.annotations.Immutable;
+import dev.cel.common.annotations.Internal;
+import dev.cel.common.exceptions.CelAttributeNotFoundException;
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.MapType;
+import dev.cel.common.types.SimpleType;
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+/**
+ * A custom CelValue implementation that allows O(1) insertions for maps during comprehension.
+ *
+ * <p>CEL Library Internals. Do Not Use.
+ */
+@Internal
+@Immutable
+@SuppressWarnings("Immutable") // Intentionally mutable for performance reasons
+public final class MutableMapValue extends CelValue
+    implements SelectableValue<Object>, Map<Object, Object> {
+  private final Map<Object, Object> internalMap;
+  private final CelType celType;
+
+  public static MutableMapValue create(Map<?, ?> map) {
+    return new MutableMapValue(map);
+  }
+
+  @Override
+  public int size() {
+    return internalMap.size();
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return internalMap.isEmpty();
+  }
+
+  @Override
+  public boolean containsKey(Object key) {
+    return internalMap.containsKey(key);
+  }
+
+  @Override
+  public boolean containsValue(Object value) {
+    return internalMap.containsValue(value);
+  }
+
+  @Override
+  public Object get(Object key) {
+    return internalMap.get(key);
+  }
+
+  @Override
+  public Object put(Object key, Object value) {
+    return internalMap.put(key, value);
+  }
+
+  @Override
+  public Object remove(Object key) {
+    return internalMap.remove(key);
+  }
+
+  @Override
+  public void putAll(Map<?, ?> m) {
+    internalMap.putAll(m);
+  }
+
+  @Override
+  public void clear() {
+    internalMap.clear();
+  }
+
+  @Override
+  public Set<Object> keySet() {
+    return internalMap.keySet();
+  }
+
+  @Override
+  public Collection<Object> values() {
+    return internalMap.values();
+  }
+
+  @Override
+  public Set<Entry<Object, Object>> entrySet() {
+    return internalMap.entrySet();
+  }
+
+  @Override
+  public Object select(Object field) {
+    Object val = internalMap.get(field);
+    if (val != null) {
+      return val;
+    }
+    if (!internalMap.containsKey(field)) {
+      throw CelAttributeNotFoundException.forMissingMapKey(field.toString());
+    }
+    throw CelAttributeNotFoundException.of(
+        String.format("Map value cannot be null for key: %s", field));
+  }
+
+  @Override
+  public Optional<?> find(Object field) {
+    if (internalMap.containsKey(field)) {
+      return Optional.ofNullable(internalMap.get(field));
+    }
+    return Optional.empty();
+  }
+
+  @Override
+  public Object value() {
+    return this;
+  }
+
+  @Override
+  public boolean isZeroValue() {
+    return internalMap.isEmpty();
+  }
+
+  @Override
+  public CelType celType() {
+    return celType;
+  }
+
+  private MutableMapValue(Map<?, ?> map) {
+    this.internalMap = new LinkedHashMap<>(map);
+    this.celType = MapType.create(SimpleType.DYN, SimpleType.DYN);
+  }
+}

common/values/BUILD.bazel

@@ -47,6 +47,18 @@ cel_android_library(
     exports = ["//common/src/main/java/dev/cel/common/values:values_android"],
 )
 
+java_library(
+    name = "mutable_map_value",
+    visibility = ["//:internal"],
+    exports = ["//common/src/main/java/dev/cel/common/values:mutable_map_value"],
+)
+
+cel_android_library(
+    name = "mutable_map_value_android",
+    visibility = ["//:internal"],
+    exports = ["//common/src/main/java/dev/cel/common/values:mutable_map_value_android"],
+)
+
 java_library(
     name = "base_proto_cel_value_converter",
     exports = ["//common/src/main/java/dev/cel/common/values:base_proto_cel_value_converter"],

extensions/src/main/java/dev/cel/extensions/BUILD.bazel

@@ -307,6 +307,7 @@ java_library(
         "//common:options",
         "//common/ast",
         "//common/types",
+        "//common/values:mutable_map_value",
         "//compiler:compiler_builder",
         "//extensions:extension_library",
         "//parser:macro",

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

@@ -29,6 +29,7 @@
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.types.MapType;
 import dev.cel.common.types.TypeParamType;
+import dev.cel.common.values.MutableMapValue;
 import dev.cel.compiler.CelCompilerLibrary;
 import dev.cel.parser.CelMacro;
 import dev.cel.parser.CelMacroExprFactory;
@@ -171,38 +172,46 @@ public void setParserOptions(CelParserBuilder parserBuilder) {
     parserBuilder.addMacros(macros());
   }
 
-  // TODO: Implement a more efficient map insertion based on mutability once mutable
-  // maps are supported in Java stack.
-  private static ImmutableMap<Object, Object> mapInsertMap(
+  private static Map<Object, Object> mapInsertMap(
       Map<?, ?> targetMap, Map<?, ?> mapToMerge, RuntimeEquality equality) {
-    ImmutableMap.Builder<Object, Object> resultBuilder =
-        ImmutableMap.builderWithExpectedSize(targetMap.size() + mapToMerge.size());
-
-    for (Map.Entry<?, ?> entry : mapToMerge.entrySet()) {
-      if (equality.findInMap(targetMap, entry.getKey()).isPresent()) {
+    for (Object key : mapToMerge.keySet()) {
+      if (equality.findInMap(targetMap, key).isPresent()) {
         throw new IllegalArgumentException(
-            String.format("insert failed: key '%s' already exists", entry.getKey()));
-      } else {
-        resultBuilder.put(entry.getKey(), entry.getValue());
+            String.format("insert failed: key '%s' already exists", key));
       }
     }
-    return resultBuilder.putAll(targetMap).buildOrThrow();
+
+    if (targetMap instanceof MutableMapValue) {
+      MutableMapValue wrapper = (MutableMapValue) targetMap;
+      wrapper.putAll(mapToMerge);
+      return wrapper;
+    }
+
+    return ImmutableMap.builderWithExpectedSize(targetMap.size() + mapToMerge.size())
+        .putAll(targetMap)
+        .putAll(mapToMerge)
+        .buildOrThrow();
   }
 
-  private static ImmutableMap<Object, Object> mapInsertKeyValue(
-      Object[] args, RuntimeEquality equality) {
-    Map<?, ?> map = (Map<?, ?>) args[0];
+  private static Map<Object, Object> mapInsertKeyValue(Object[] args, RuntimeEquality equality) {
+    Map<?, ?> mapArg = (Map<?, ?>) args[0];
     Object key = args[1];
     Object value = args[2];
 
-    if (equality.findInMap(map, key).isPresent()) {
+    if (equality.findInMap(mapArg, key).isPresent()) {
       throw new IllegalArgumentException(
           String.format("insert failed: key '%s' already exists", key));
     }
 
+    if (mapArg instanceof MutableMapValue) {
+      MutableMapValue mutableMap = (MutableMapValue) mapArg;
+      mutableMap.put(key, value);
+      return mutableMap;
+    }
+
     ImmutableMap.Builder<Object, Object> builder =
-        ImmutableMap.builderWithExpectedSize(map.size() + 1);
-    return builder.put(key, value).putAll(map).buildOrThrow();
+        ImmutableMap.builderWithExpectedSize(mapArg.size() + 1);
+    return builder.put(key, value).putAll(mapArg).buildOrThrow();
   }
 
   private static Optional<CelExpr> expandAllMacro(

extensions/src/test/java/dev/cel/extensions/BUILD.bazel

@@ -15,6 +15,7 @@ java_library(
         "//common:compiler_common",
         "//common:container",
         "//common:options",
+        "//common/exceptions:attribute_not_found",
         "//common/exceptions:divide_by_zero",
         "//common/exceptions:index_out_of_bounds",
         "//common/types",

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

@@ -28,6 +28,7 @@
 import dev.cel.common.CelFunctionDecl;
 import dev.cel.common.CelOptions;
 import dev.cel.common.CelValidationException;
+import dev.cel.common.exceptions.CelAttributeNotFoundException;
 import dev.cel.common.exceptions.CelDivideByZeroException;
 import dev.cel.common.exceptions.CelIndexOutOfBoundsException;
 import dev.cel.common.types.SimpleType;
@@ -222,6 +223,7 @@ public void transformMapEntryMacro_twoVarComprehension_success(
                 + " 'key2': 'value2'}",
             // map.transformMapEntry()
             "{'hello': 'world', 'greetings': 'tacocat'}.transformMapEntry(k, v, {}) == {}",
+            "{'a': 1, 'b': 2}.transformMapEntry(k, v, {k: v}) == {'a': 1, 'b': 2}",
             "{'a': 1, 'b': 2}.transformMapEntry(k, v, {k + '_new': v * 2}) == {'a_new': 2,"
                 + " 'b_new': 4}",
             "{'a': 1, 'b': 2, 'c': 3}.transformMapEntry(k, v, v % 2 == 1, {k: v * 10}) == {'a': 10,"
@@ -364,6 +366,15 @@ public void twoVarComprehension_outOfBounds_runtimeError() throws Exception {
     assertThat(e).hasCauseThat().hasMessageThat().contains("Index out of bounds: 1");
   }
 
+  @Test
+  public void mutableMapValue_select_missingKeyException() throws Exception {
+    CelEvaluationException e =
+        assertThrows(
+            CelEvaluationException.class, () -> eval("cel.bind(my_map, {'a': 1}, my_map.b)"));
+    assertThat(e).hasCauseThat().isInstanceOf(CelAttributeNotFoundException.class);
+    assertThat(e).hasCauseThat().hasMessageThat().contains("key 'b' is not present in map.");
+  }
+
   private Object eval(String expression) throws Exception {
     return eval(this.cel, expression, ImmutableMap.of());
   }

runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel

@@ -365,6 +365,7 @@ java_library(
         ":activation_wrapper",
         ":planned_interpretable",
         "//common/exceptions:runtime_exception",
+        "//common/values:mutable_map_value",
         "//runtime:accumulated_unknowns",
         "//runtime:concatenated_list_view",
         "//runtime:evaluation_exception",

runtime/src/main/java/dev/cel/runtime/planner/EvalFold.java

@@ -15,8 +15,10 @@
 package dev.cel.runtime.planner;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.exceptions.CelRuntimeException;
+import dev.cel.common.values.MutableMapValue;
 import dev.cel.runtime.AccumulatedUnknowns;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.ConcatenatedListView;
@@ -131,31 +133,34 @@ private Object evalList(Collection<?> iterRange, Folder folder, ExecutionFrame f
       boolean cond = (boolean) condition.eval(folder, frame);
       if (!cond) {
         folder.computeResult = true;
-        return result.eval(folder, frame);
+        return maybeUnwrapAccumulator(result.eval(folder, frame));
       }
 
       folder.accuVal = loopStep.eval(folder, frame);
       folder.initialized = true;
       index++;
     }
     folder.computeResult = true;
-    return result.eval(folder, frame);
+    return maybeUnwrapAccumulator(result.eval(folder, frame));
   }
 
   private static Object maybeWrapAccumulator(Object val) {
     if (val instanceof Collection) {
       return new ConcatenatedListView<>((Collection<?>) val);
     }
-    // TODO: Introduce mutable map support (for comp v2)
+    if (val instanceof Map) {
+      return MutableMapValue.create((Map<?, ?>) val);
+    }
     return val;
   }
 
   private static Object maybeUnwrapAccumulator(Object val) {
     if (val instanceof ConcatenatedListView) {
       return ImmutableList.copyOf((ConcatenatedListView<?>) val);
     }
-
-    // TODO: Introduce mutable map support (for comp v2)
+    if (val instanceof MutableMapValue) {
+      return ImmutableMap.copyOf((MutableMapValue) val);
+    }
     return val;
   }