From cad0258d1758a0ca9f1b9725103bcdbae26697e1 Mon Sep 17 00:00:00 2001
From: Feng Xiao <xfxyjwf@gmail.com>
Date: Tue, 11 Apr 2017 16:08:16 -0700
Subject: Cherry-pick cl/151775298

---
 src/google/protobuf/map_test.cc    | 27 +++++++++++++++++++++++++++
 src/google/protobuf/wire_format.cc | 31 +++++++++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/google/protobuf/map_test.cc b/src/google/protobuf/map_test.cc
index 3fe44512..e10a4278 100644
--- a/src/google/protobuf/map_test.cc
+++ b/src/google/protobuf/map_test.cc
@@ -2811,6 +2811,33 @@ TEST(WireFormatForMapFieldTest, SerializeMap) {
   EXPECT_TRUE(dynamic_data == generated_data);
 }
 
+TEST(WireFormatForMapFieldTest, SerializeMapDynamicMessage) {
+  DynamicMessageFactory factory;
+  google::protobuf::scoped_ptr<Message> dynamic_message;
+  dynamic_message.reset(
+      factory.GetPrototype(unittest::TestMap::descriptor())->New());
+  MapReflectionTester reflection_tester(
+      unittest::TestMap::descriptor());
+  reflection_tester.SetMapFieldsViaReflection(dynamic_message.get());
+  reflection_tester.ExpectMapFieldsSetViaReflection(*dynamic_message);
+
+  unittest::TestMap generated_message;
+  MapTestUtil::SetMapFields(&generated_message);
+  MapTestUtil::ExpectMapFieldsSet(generated_message);
+
+  string generated_data;
+  string dynamic_data;
+
+  // Serialize.
+  generated_message.SerializeToString(&generated_data);
+  dynamic_message->SerializeToString(&dynamic_data);
+
+  // Because map serialization doesn't guarantee order, we just compare
+  // serialized size here. This is enough to tell dynamic message doesn't miss
+  // anything in serialization.
+  EXPECT_TRUE(dynamic_data.size() == generated_data.size());
+}
+
 TEST(WireFormatForMapFieldTest, MapParseHelpers) {
   string data;
 
diff --git a/src/google/protobuf/wire_format.cc b/src/google/protobuf/wire_format.cc
index d3d21c09..7778ecd1 100644
--- a/src/google/protobuf/wire_format.cc
+++ b/src/google/protobuf/wire_format.cc
@@ -797,7 +797,16 @@ void WireFormat::SerializeWithCachedSizes(
   int expected_endpoint = output->ByteCount() + size;
 
   std::vector<const FieldDescriptor*> fields;
-  message_reflection->ListFields(message, &fields);
+
+  // Fields of map entry should always be serialized.
+  if (descriptor->options().map_entry()) {
+    for (int i = 0; i < descriptor->field_count(); i++) {
+      fields.push_back(descriptor->field(i));
+    }
+  } else {
+    message_reflection->ListFields(message, &fields);
+  }
+
   for (int i = 0; i < fields.size(); i++) {
     SerializeFieldWithCachedSizes(fields[i], message, output);
   }
@@ -834,6 +843,9 @@ void WireFormat::SerializeFieldWithCachedSizes(
 
   if (field->is_repeated()) {
     count = message_reflection->FieldSize(message, field);
+  } else if (field->containing_type()->options().map_entry()) {
+    // Map entry fields always need to be serialized.
+    count = 1;
   } else if (message_reflection->HasField(message, field)) {
     count = 1;
   }
@@ -984,7 +996,16 @@ size_t WireFormat::ByteSize(const Message& message) {
   size_t our_size = 0;
 
   std::vector<const FieldDescriptor*> fields;
-  message_reflection->ListFields(message, &fields);
+
+  // Fields of map entry should always be serialized.
+  if (descriptor->options().map_entry()) {
+    for (int i = 0; i < descriptor->field_count(); i++) {
+      fields.push_back(descriptor->field(i));
+    }
+  } else {
+    message_reflection->ListFields(message, &fields);
+  }
+
   for (int i = 0; i < fields.size(); i++) {
     our_size += FieldByteSize(fields[i], message);
   }
@@ -1015,6 +1036,9 @@ size_t WireFormat::FieldByteSize(
   size_t count = 0;
   if (field->is_repeated()) {
     count = FromIntSize(message_reflection->FieldSize(message, field));
+  } else if (field->containing_type()->options().map_entry()) {
+    // Map entry fields always need to be serialized.
+    count = 1;
   } else if (message_reflection->HasField(message, field)) {
     count = 1;
   }
@@ -1044,6 +1068,9 @@ size_t WireFormat::FieldDataOnlyByteSize(
   if (field->is_repeated()) {
     count =
         internal::FromIntSize(message_reflection->FieldSize(message, field));
+  } else if (field->containing_type()->options().map_entry()) {
+    // Map entry fields always need to be serialized.
+    count = 1;
   } else if (message_reflection->HasField(message, field)) {
     count = 1;
   }
-- 
cgit v1.2.3