Add transformMapEntry macros and a map-merging overload for cel.@mapInsert.

PiperOrigin-RevId: 794350506

Author: dmitriplotnikov

eval/compiler/flat_expr_builder.cc

@@ -357,7 +357,11 @@ const cel::Expr* GetOptimizableListAppendOperand(
 // Returns whether this comprehension appears to be a macro implementation for
 // map transformations. It is not exhaustive, so it is unsafe to use with custom
 // comprehensions outside of the standard macros or hand crafted ASTs.
-bool IsOptimizableMapInsert(const cel::ComprehensionExpr* comprehension) {
+bool IsOptimizableMapInsert(const cel::ComprehensionExpr* comprehension,
+                            bool enable_comprehension_mutable_map) {
+  if (!enable_comprehension_mutable_map) {
+    return false;
+  }
   if (comprehension->iter_var().empty() || comprehension->iter_var2().empty()) {
     return false;
   }
@@ -383,7 +387,7 @@ bool IsOptimizableMapInsert(const cel::ComprehensionExpr* comprehension) {
     call_expr = &(call_expr->args()[1].call_expr());
   }
   return call_expr->function() == "cel.@mapInsert" &&
-         call_expr->args().size() == 3 &&
+         (call_expr->args().size() == 2 || call_expr->args().size() == 3) &&
          call_expr->args()[0].has_ident_expr() &&
          call_expr->args()[0].ident_expr().name() == accu_var;
 }
@@ -1407,7 +1411,9 @@ class FlatExprVisitor : public cel::AstVisitor {
          /*.is_optimizable_list_append=*/
          IsOptimizableListAppend(&comprehension,
                                  options_.enable_comprehension_list_append),
-         /*.is_optimizable_map_insert=*/IsOptimizableMapInsert(&comprehension),
+         /*.is_optimizable_map_insert=*/
+         IsOptimizableMapInsert(&comprehension,
+                                options_.enable_comprehension_mutable_map),
          /*.is_optimizable_bind=*/is_bind,
          /*.iter_var_in_scope=*/false,
          /*.iter_var2_in_scope=*/false,
@@ -1587,21 +1593,6 @@ class FlatExprVisitor : public cel::AstVisitor {
       return;
     }
 
-    if (!comprehension_stack_.empty()) {
-      const ComprehensionStackRecord& comprehension =
-          comprehension_stack_.back();
-      if (comprehension.is_optimizable_map_insert) {
-        if (&(comprehension.comprehension->accu_init()) == &expr) {
-          if (options_.max_recursion_depth != 0) {
-            SetRecursiveStep(CreateDirectMutableMapStep(expr.id()), 1);
-            return;
-          }
-          AddStep(CreateMutableMapStep(expr.id()));
-          return;
-        }
-      }
-    }
-
     auto status_or_resolved_fields =
         ResolveCreateStructFields(struct_expr, expr.id());
     if (!status_or_resolved_fields.ok()) {
@@ -1639,6 +1630,22 @@ class FlatExprVisitor : public cel::AstVisitor {
       ValidateOrError(entry.has_key(), "Map entry missing key");
       ValidateOrError(entry.has_value(), "Map entry missing value");
     }
+
+    if (!comprehension_stack_.empty()) {
+      const ComprehensionStackRecord& comprehension =
+          comprehension_stack_.back();
+      if (comprehension.is_optimizable_map_insert) {
+        if (&(comprehension.comprehension->accu_init()) == &expr) {
+          if (options_.max_recursion_depth != 0) {
+            SetRecursiveStep(CreateDirectMutableMapStep(expr.id()), 1);
+            return;
+          }
+          AddStep(CreateMutableMapStep(expr.id()));
+          return;
+        }
+      }
+    }
+
     auto depth = RecursionEligible();
     if (depth.has_value()) {
       auto deps = ExtractRecursiveDependencies();

eval/public/cel_options.cc

@@ -27,6 +27,7 @@ cel::RuntimeOptions ConvertToRuntimeOptions(const InterpreterOptions& options) {
                              options.enable_comprehension,
                              options.comprehension_max_iterations,
                              options.enable_comprehension_list_append,
+                             options.enable_comprehension_mutable_map,
                              options.enable_regex,
                              options.regex_max_program_size,
                              options.enable_string_conversion,

eval/public/cel_options.h

@@ -17,8 +17,6 @@
 #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_CEL_OPTIONS_H_
 #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_CEL_OPTIONS_H_
 
-#include <string>
-
 #include "absl/base/attributes.h"
 #include "runtime/runtime_options.h"
 #include "google/protobuf/arena.h"
@@ -74,6 +72,10 @@ struct InterpreterOptions {
   // with hand-rolled ASTs.
   bool enable_comprehension_list_append = false;
 
+  // Enable mutable map construction within comprehensions. Note, this option is
+  // not safe with hand-rolled ASTs.
+  bool enable_comprehension_mutable_map = false;
+
   // Enable RE2 match() overload.
   bool enable_regex = true;
 

eval/tests/BUILD

@@ -61,13 +61,17 @@ cc_test(
         "//common:memory",
         "//common:native_type",
         "//common:value",
+        "//extensions:comprehensions_v2_functions",
+        "//extensions:comprehensions_v2_macros",
         "//extensions/protobuf:runtime_adapter",
         "//extensions/protobuf:value",
         "//internal:benchmark",
         "//internal:testing",
         "//internal:testing_descriptor_pool",
         "//internal:testing_message_factory",
         "//parser",
+        "//parser:macro",
+        "//parser:macro_registry",
         "//runtime",
         "//runtime:activation",
         "//runtime:constant_folding",

eval/tests/modern_benchmark_test.cc

@@ -41,12 +41,15 @@
 #include "common/native_type.h"
 #include "common/value.h"
 #include "eval/tests/request_context.pb.h"
+#include "extensions/comprehensions_v2_functions.h"
+#include "extensions/comprehensions_v2_macros.h"
 #include "extensions/protobuf/runtime_adapter.h"
 #include "extensions/protobuf/value.h"
 #include "internal/benchmark.h"
 #include "internal/testing.h"
 #include "internal/testing_descriptor_pool.h"
 #include "internal/testing_message_factory.h"
+#include "parser/macro_registry.h"
 #include "parser/parser.h"
 #include "runtime/activation.h"
 #include "runtime/constant_folding.h"
@@ -70,6 +73,7 @@ using ::cel::extensions::ProtobufRuntimeAdapter;
 using ::cel::expr::Expr;
 using ::cel::expr::ParsedExpr;
 using ::cel::expr::SourceInfo;
+using ::google::api::expr::parser::EnrichedParse;
 using ::google::api::expr::parser::Parse;
 using ::google::api::expr::runtime::RequestContext;
 using ::google::rpc::context::AttributeContext;
@@ -1270,6 +1274,62 @@ void BM_ComprehensionCpp(benchmark::State& state) {
 
 BENCHMARK(BM_ComprehensionCpp)->Range(1, 1 << 20);
 
+void BM_MapTransformComprehension(benchmark::State& state) {
+  ASSERT_OK_AND_ASSIGN(auto source,
+                       NewSource("map_var.transformMapEntry(k, v, {v:k})"));
+
+  MacroRegistry registry;
+  ASSERT_THAT(
+      extensions::RegisterComprehensionsV2Macros(registry, ParserOptions()),
+      IsOk());
+
+  ASSERT_OK_AND_ASSIGN(auto parsed_expr,
+                       EnrichedParse(*source, registry, ParserOptions()));
+
+  RuntimeOptions options = GetOptions();
+  options.comprehension_max_iterations = 10000000;
+
+  // This is a critical optimization: it allows the comprehension to accumulate
+  // results in a mutable map instead of cloning and augmenting an unmodifiable
+  // map on every iteration.
+  options.enable_comprehension_mutable_map = true;
+
+  ASSERT_OK_AND_ASSIGN(auto builder,
+                       CreateStandardRuntimeBuilder(
+                           internal::GetTestingDescriptorPool(), options));
+
+  ASSERT_THAT(extensions::RegisterComprehensionsV2Functions(
+                  builder.function_registry(), options),
+              IsOk());
+
+  ASSERT_OK_AND_ASSIGN(auto runtime, std::move(builder).Build());
+
+  google::protobuf::Arena arena;
+  Activation activation;
+
+  auto map_builder = cel::NewMapValueBuilder(&arena);
+
+  int len = state.range(0);
+  map_builder->Reserve(len);
+  for (int i = 0; i < len; i++) {
+    ASSERT_THAT(map_builder->Put(IntValue(i), IntValue(i)), IsOk());
+  }
+
+  activation.InsertOrAssignValue("map_var", std::move(*map_builder).Build());
+
+  ASSERT_OK_AND_ASSIGN(auto cel_expr, ProtobufRuntimeAdapter::CreateProgram(
+                                          *runtime, parsed_expr.parsed_expr()));
+
+  for (auto _ : state) {
+    ASSERT_OK_AND_ASSIGN(cel::Value result,
+                         cel_expr->Evaluate(&arena, activation));
+    ASSERT_TRUE(InstanceOf<MapValue>(result));
+    ASSERT_THAT(Cast<MapValue>(result).Size(), IsOkAndHolds(len));
+  }
+}
+
+BENCHMARK(BM_MapTransformComprehension)->Range(1, 1 << 16);
+
 }  // namespace
 
 }  // namespace cel

extensions/BUILD

@@ -581,6 +581,7 @@ cc_test(
         ":comprehensions_v2_macros",
         ":strings",
         "//common:source",
+        "//common:value",
         "//common:value_testing",
         "//extensions/protobuf:runtime_adapter",
         "//internal:status_macros",
@@ -596,6 +597,7 @@ cc_test(
         "//runtime:reference_resolver",
         "//runtime:runtime_options",
         "//runtime:standard_runtime_builder_factory",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",

extensions/comprehensions_v2_functions.cc

@@ -35,7 +35,7 @@ namespace cel::extensions {
 
 namespace {
 
-absl::StatusOr<Value> MapInsert(
+absl::StatusOr<Value> MapInsertKeyValue(
     const MapValue& map, const Value& key, const Value& value,
     const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
     google::protobuf::MessageFactory* absl_nonnull message_factory,
@@ -68,6 +68,54 @@ absl::StatusOr<Value> MapInsert(
   return std::move(*builder).Build();
 }
 
+absl::StatusOr<Value> MapInsertMap(
+    const MapValue& map, const MapValue& value,
+    const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
+    google::protobuf::MessageFactory* absl_nonnull message_factory,
+    google::protobuf::Arena* absl_nonnull arena) {
+  if (auto mutable_map_value = common_internal::AsMutableMapValue(map);
+      mutable_map_value) {
+    // Fast path, runtime has given us a mutable map. We can mutate it directly
+    // and return it.
+    CEL_RETURN_IF_ERROR(
+        value.ForEach(
+            [&mutable_map_value](const Value& key,
+                                 const Value& value) -> absl::StatusOr<bool> {
+              CEL_RETURN_IF_ERROR(mutable_map_value->Put(key, value));
+              return true;
+            },
+            descriptor_pool, message_factory, arena))
+        .With(ErrorValueReturn());
+    return map;
+  }
+  // Slow path, we have to make a copy.
+  auto builder = NewMapValueBuilder(arena);
+  if (auto size = map.Size(); size.ok()) {
+    builder->Reserve(*size + 1);
+  } else {
+    size.IgnoreError();
+  }
+  CEL_RETURN_IF_ERROR(
+      map.ForEach(
+          [&builder](const Value& key,
+                     const Value& value) -> absl::StatusOr<bool> {
+            CEL_RETURN_IF_ERROR(builder->Put(key, value));
+            return true;
+          },
+          descriptor_pool, message_factory, arena))
+      .With(ErrorValueReturn());
+  CEL_RETURN_IF_ERROR(
+      value.ForEach(
+          [&builder](const Value& key,
+                     const Value& value) -> absl::StatusOr<bool> {
+            CEL_RETURN_IF_ERROR(builder->Put(key, value));
+            return true;
+          },
+          descriptor_pool, message_factory, arena))
+      .With(ErrorValueReturn());
+  return std::move(*builder).Build();
+}
+
 }  // namespace
 
 absl::Status RegisterComprehensionsV2Functions(FunctionRegistry& registry,
@@ -77,7 +125,15 @@ absl::Status RegisterComprehensionsV2Functions(FunctionRegistry& registry,
                              Value>::CreateDescriptor("cel.@mapInsert",
                                                       /*receiver_style=*/false),
       TernaryFunctionAdapter<absl::StatusOr<Value>, MapValue, Value,
-                             Value>::WrapFunction(&MapInsert)));
+                             Value>::WrapFunction(&MapInsertKeyValue)));
+
+  CEL_RETURN_IF_ERROR(registry.Register(
+      BinaryFunctionAdapter<absl::StatusOr<Value>, MapValue, MapValue>::
+          CreateDescriptor("cel.@mapInsert",
+                           /*receiver_style=*/false),
+      BinaryFunctionAdapter<absl::StatusOr<Value>, MapValue,
+                            MapValue>::WrapFunction(&MapInsertMap)));
+
   return absl::OkStatus();
 }