Implement mutable map support for planner

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",