1 files changed, 68 insertions, 1 deletions

diff --git a/src/google/protobuf/util/field_mask_util.cc b/src/google/protobuf/util/field_mask_util.cc
index 85cecec5..982d6407 100644
--- a/src/google/protobuf/util/field_mask_util.cc
+++ b/src/google/protobuf/util/field_mask_util.cc
@@ -207,6 +207,18 @@ class FieldMaskTree {
     MergeMessage(&root_, source, options, destination);
   }
 
+  // Add required field path of the message to this tree based on current tree
+  // structure. If a message is present in the tree, add the path of its
+  // required field to the tree. This is to make sure that after trimming a
+  // message with required fields are set, check IsInitialized() will not fail.
+  void AddRequiredFieldPath(const Descriptor* descriptor) {
+    // Do nothing if the tree is empty.
+    if (root_.children.empty()) {
+      return;
+    }
+    AddRequiredFieldPath(&root_, descriptor);
+  }
+
   // Trims all fields not specified by this tree from the given message.
   void TrimMessage(Message* message) {
     // Do nothing if the tree is empty.
@@ -249,6 +261,12 @@ class FieldMaskTree {
                     const FieldMaskUtil::MergeOptions& options,
                     Message* destination);
 
+  // Add required field path of the message to this tree based on current tree
+  // structure. If a message is present in the tree, add the path of its
+  // required field to the tree. This is to make sure that after trimming a
+  // message with required fields are set, check IsInitialized() will not fail.
+  void AddRequiredFieldPath(Node* node, const Descriptor* descriptor);
+
   // Trims all fields not specified by this sub-tree from the given message.
   void TrimMessage(const Node* node, Message* message);
 
@@ -456,6 +474,41 @@ void FieldMaskTree::MergeMessage(const Node* node, const Message& source,
   }
 }
 
+void FieldMaskTree::AddRequiredFieldPath(
+    Node* node, const Descriptor* descriptor) {
+  const int32 field_count = descriptor->field_count();
+  for (int index = 0; index < field_count; ++index) {
+    const FieldDescriptor* field = descriptor->field(index);
+    if (field->is_required()) {
+      const string& node_name = field->name();
+      Node*& child = node->children[node_name];
+      if (child == NULL) {
+        // Add required field path to the tree
+        child = new Node();
+      } else if (child->children.empty()){
+        // If the required field is in the tree and does not have any children,
+        // do nothing.
+        continue;
+      }
+      // Add required field in the children to the tree if the field is message.
+      if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
+        AddRequiredFieldPath(child, field->message_type());
+      }
+    } else if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
+      std::map<string, Node*>::const_iterator it =
+          node->children.find(field->name());
+      if (it != node->children.end()) {
+        // Add required fields in the children to the
+        // tree if the field is a message and present in the tree.
+        Node* child = it->second;
+        if (!child->children.empty()) {
+          AddRequiredFieldPath(child, field->message_type());
+        }
+      }
+    }
+  }
+}
+
 void FieldMaskTree::TrimMessage(const Node* node, Message* message) {
   GOOGLE_DCHECK(!node->children.empty());
   const Reflection* reflection = message->GetReflection();
@@ -470,7 +523,7 @@ void FieldMaskTree::TrimMessage(const Node* node, Message* message) {
     } else {
       if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
         Node* child = it->second;
-        if (!child->children.empty()) {
+        if (!child->children.empty() && reflection->HasField(*message, field)) {
           TrimMessage(child, reflection->MutableMessage(message, field));
         }
       }
@@ -542,6 +595,20 @@ void FieldMaskUtil::TrimMessage(const FieldMask& mask, Message* destination) {
   tree.TrimMessage(GOOGLE_CHECK_NOTNULL(destination));
 }
 
+void FieldMaskUtil::TrimMessage(const FieldMask& mask, Message* destination,
+                                const TrimOptions& options) {
+  // Build a FieldMaskTree and walk through the tree to merge all specified
+  // fields.
+  FieldMaskTree tree;
+  tree.MergeFromFieldMask(mask);
+  // If keep_required_fields is true, implicitely add required fields of
+  // a message present in the tree to prevent from trimming.
+  if (options.keep_required_fields()) {
+    tree.AddRequiredFieldPath(GOOGLE_CHECK_NOTNULL(destination->GetDescriptor()));
+  }
+  tree.TrimMessage(GOOGLE_CHECK_NOTNULL(destination));
+}
+
 }  // namespace util
 }  // namespace protobuf
 }  // namespace google