From 6624e240bc2a2aa560b0ac2278c06a5a3567a3e0 Mon Sep 17 00:00:00 2001
From: Justin King <jcking@google.com>
Date: Fri, 8 May 2026 15:03:58 -0700
Subject: [PATCH] Fix map field value null pruning for proto2/proto3

PiperOrigin-RevId: 912713146
---
 common/values/struct_value_builder.cc         | 23 +++++++++++++++++++
 conformance/BUILD                             |  6 -----
 .../structs/proto_message_type_adapter.cc     | 16 +++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/common/values/struct_value_builder.cc b/common/values/struct_value_builder.cc
index c342d6478..446b18421 100644
--- a/common/values/struct_value_builder.cc
+++ b/common/values/struct_value_builder.cc
@@ -956,6 +956,19 @@ class MessageValueBuilderImpl {
           if (error_value) {
             return false;
           }
+          if (map_value_field->cpp_type() ==
+                  google::protobuf::FieldDescriptor::CPPTYPE_MESSAGE &&
+              entry_value.IsNull()) {
+            auto well_known_type =
+                map_value_field->message_type()->well_known_type();
+            if (well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_ANY &&
+                well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_VALUE &&
+                well_known_type !=
+                    google::protobuf::Descriptor::WELLKNOWNTYPE_LISTVALUE &&
+                well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_STRUCT) {
+              return true;
+            }
+          }
           google::protobuf::MapValueRef proto_value;
           extensions::protobuf_internal::InsertOrLookupMapValue(
               *reflection_, message_, *field, proto_key, &proto_value);
@@ -989,6 +1002,16 @@ class MessageValueBuilderImpl {
     CEL_RETURN_IF_ERROR(list_value->ForEach(
         [this, field, accessor,
          &error_value](const Value& element) -> absl::StatusOr<bool> {
+          if (field->message_type() != nullptr && element.IsNull()) {
+            auto well_known_type = field->message_type()->well_known_type();
+            if (well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_ANY &&
+                well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_VALUE &&
+                well_known_type !=
+                    google::protobuf::Descriptor::WELLKNOWNTYPE_LISTVALUE &&
+                well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_STRUCT) {
+              return true;
+            }
+          }
           CEL_ASSIGN_OR_RETURN(error_value,
                                (*accessor)(descriptor_pool_, message_factory_,
                                            &well_known_types_, reflection_,
diff --git a/conformance/BUILD b/conformance/BUILD
index 4f9232ab6..ccd2844c9 100644
--- a/conformance/BUILD
+++ b/conformance/BUILD
@@ -203,12 +203,6 @@ _TESTS_TO_SKIP = [
     "conversions/string/double_hard",
 
     # Recent changes
-    "proto2/set_null/map_timestamp_null_pruned",
-    "proto2/set_null/map_duration_null_pruned",
-    "proto2/set_null/map_wrapper_null_pruned",
-    "proto3/set_null/map_timestamp_null_pruned",
-    "proto3/set_null/map_duration_null_pruned",
-    "proto3/set_null/map_wrapper_null_pruned",
     "namespace/namespace_shadowing/basic",
     "namespace/namespace_shadowing/comprehension_shadowing_namespaced_selector_disambiguation",
 ]
diff --git a/eval/public/structs/proto_message_type_adapter.cc b/eval/public/structs/proto_message_type_adapter.cc
index a351890c2..6a3417ba3 100644
--- a/eval/public/structs/proto_message_type_adapter.cc
+++ b/eval/public/structs/proto_message_type_adapter.cc
@@ -582,6 +582,19 @@ absl::Status ProtoMessageTypeAdapter::SetField(
         ValidateSetFieldOp(value_field_descriptor != nullptr, field->name(),
                            "failed to find value field descriptor"));
 
+    bool prune_when_null = false;
+    if (value_field_descriptor->cpp_type() ==
+        google::protobuf::FieldDescriptor::CPPTYPE_MESSAGE) {
+      auto well_known_type =
+          value_field_descriptor->message_type()->well_known_type();
+      if (well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_ANY &&
+          well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_VALUE &&
+          well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_LISTVALUE &&
+          well_known_type != google::protobuf::Descriptor::WELLKNOWNTYPE_STRUCT) {
+        prune_when_null = true;
+      }
+    }
+
     CEL_ASSIGN_OR_RETURN(const CelList* key_list, cel_map->ListKeys(arena));
     for (int i = 0; i < key_list->size(); i++) {
       CelValue key = (*key_list).Get(arena, i);
@@ -589,6 +602,9 @@ absl::Status ProtoMessageTypeAdapter::SetField(
       auto value = (*cel_map).Get(arena, key);
       CEL_RETURN_IF_ERROR(ValidateSetFieldOp(value.has_value(), field->name(),
                                              "error serializing CelMap"));
+      if (prune_when_null && value->IsNull()) {
+        continue;
+      }
       Message* entry_msg = message->GetReflection()->AddMessage(message, field);
       CEL_RETURN_IF_ERROR(internal::SetValueToSingleField(
           key, key_field_descriptor, entry_msg, arena));