1 files changed, 298 insertions, 19 deletions

diff --git a/src/google/protobuf/util/field_mask_util.cc b/src/google/protobuf/util/field_mask_util.cc
index c59f43aa..a2e2a388 100644
--- a/src/google/protobuf/util/field_mask_util.cc
+++ b/src/google/protobuf/util/field_mask_util.cc
@@ -31,6 +31,7 @@
 #include <google/protobuf/util/field_mask_util.h>
 
 #include <google/protobuf/stubs/strutil.h>
+
 #include <google/protobuf/stubs/map_util.h>
 
 namespace google {
@@ -45,30 +46,113 @@ string FieldMaskUtil::ToString(const FieldMask& mask) {
 
 void FieldMaskUtil::FromString(StringPiece str, FieldMask* out) {
   out->Clear();
-  vector<string> paths = Split(str, ",");
+  std::vector<string> paths = Split(str, ",");
   for (int i = 0; i < paths.size(); ++i) {
     if (paths[i].empty()) continue;
     out->add_paths(paths[i]);
   }
 }
 
-bool FieldMaskUtil::InternalIsValidPath(const Descriptor* descriptor,
-                                        StringPiece path) {
-  vector<string> parts = Split(path, ".");
+bool FieldMaskUtil::SnakeCaseToCamelCase(StringPiece input, string* output) {
+  output->clear();
+  bool after_underscore = false;
+  for (int i = 0; i < input.size(); ++i) {
+    if (input[i] >= 'A' && input[i] <= 'Z') {
+      // The field name must not contain uppercase letters.
+      return false;
+    }
+    if (after_underscore) {
+      if (input[i] >= 'a' && input[i] <= 'z') {
+        output->push_back(input[i] + 'A' - 'a');
+        after_underscore = false;
+      } else {
+        // The character after a "_" must be a lowercase letter.
+        return false;
+      }
+    } else if (input[i] == '_') {
+      after_underscore = true;
+    } else {
+      output->push_back(input[i]);
+    }
+  }
+  if (after_underscore) {
+    // Trailing "_".
+    return false;
+  }
+  return true;
+}
+
+bool FieldMaskUtil::CamelCaseToSnakeCase(StringPiece input, string* output) {
+  output->clear();
+  for (int i = 0; i < input.size(); ++i) {
+    if (input[i] == '_') {
+      // The field name must not contain "_"s.
+      return false;
+    }
+    if (input[i] >= 'A' && input[i] <= 'Z') {
+      output->push_back('_');
+      output->push_back(input[i] + 'a' - 'A');
+    } else {
+      output->push_back(input[i]);
+    }
+  }
+  return true;
+}
+
+bool FieldMaskUtil::ToJsonString(const FieldMask& mask, string* out) {
+  out->clear();
+  for (int i = 0; i < mask.paths_size(); ++i) {
+    const string& path = mask.paths(i);
+    string camelcase_path;
+    if (!SnakeCaseToCamelCase(path, &camelcase_path)) {
+      return false;
+    }
+    if (i > 0) {
+      out->push_back(',');
+    }
+    out->append(camelcase_path);
+  }
+  return true;
+}
+
+bool FieldMaskUtil::FromJsonString(StringPiece str, FieldMask* out) {
+  out->Clear();
+  std::vector<string> paths = Split(str, ",");
+  for (int i = 0; i < paths.size(); ++i) {
+    if (paths[i].empty()) continue;
+    string snakecase_path;
+    if (!CamelCaseToSnakeCase(paths[i], &snakecase_path)) {
+      return false;
+    }
+    out->add_paths(snakecase_path);
+  }
+  return true;
+}
+
+bool FieldMaskUtil::GetFieldDescriptors(
+    const Descriptor* descriptor, StringPiece path,
+    std::vector<const FieldDescriptor*>* field_descriptors) {
+  if (field_descriptors != nullptr) {
+    field_descriptors->clear();
+  }
+  std::vector<string> parts = Split(path, ".");
   for (int i = 0; i < parts.size(); ++i) {
     const string& field_name = parts[i];
-    if (descriptor == NULL) {
+    if (descriptor == nullptr) {
       return false;
     }
     const FieldDescriptor* field = descriptor->FindFieldByName(field_name);
-    if (field == NULL) {
+    if (field == nullptr) {
       return false;
     }
+    if (field_descriptors != nullptr) {
+      field_descriptors->push_back(field);
+    }
     if (!field->is_repeated() &&
         field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
       descriptor = field->message_type();
     } else {
-      descriptor = NULL;
+      descriptor = nullptr;
     }
   }
   return true;
@@ -109,6 +193,13 @@ class FieldMaskTree {
   // children removed because the path matches all the node's children.
   void AddPath(const string& path);
 
+  // Remove a path from the tree.
+  // If the path is a sub-path of an existing field path in the tree, it means
+  // we need remove the existing fied path and add all sub-paths except
+  // specified path. If the path matches an existing node in the tree, this node
+  // will be moved.
+  void RemovePath(const string& path, const Descriptor* descriptor);
+
   // Calculate the intersection part of a field path with this tree and add
   // the intersection field path into out.
   void IntersectPath(const string& path, FieldMaskTree* out);
@@ -124,6 +215,27 @@ 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.
+    if (root_.children.empty()) {
+      return;
+    }
+    TrimMessage(&root_, message);
+  }
+
  private:
   struct Node {
     Node() {}
@@ -131,14 +243,14 @@ class FieldMaskTree {
     ~Node() { ClearChildren(); }
 
     void ClearChildren() {
-      for (map<string, Node*>::iterator it = children.begin();
+      for (std::map<string, Node*>::iterator it = children.begin();
            it != children.end(); ++it) {
         delete it->second;
       }
       children.clear();
     }
 
-    map<string, Node*> children;
+    std::map<string, Node*> children;
 
    private:
     GOOGLE_DISALLOW_EVIL_CONSTRUCTORS(Node);
@@ -157,6 +269,15 @@ 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);
+
   Node root_;
 
   GOOGLE_DISALLOW_EVIL_CONSTRUCTORS(FieldMaskTree);
@@ -186,7 +307,7 @@ void FieldMaskTree::MergeToFieldMask(const string& prefix, const Node* node,
     out->add_paths(prefix);
     return;
   }
-  for (map<string, Node*>::const_iterator it = node->children.begin();
+  for (std::map<string, Node*>::const_iterator it = node->children.begin();
        it != node->children.end(); ++it) {
     string current_path = prefix.empty() ? it->first : prefix + "." + it->first;
     MergeToFieldMask(current_path, it->second, out);
@@ -194,7 +315,7 @@ void FieldMaskTree::MergeToFieldMask(const string& prefix, const Node* node,
 }
 
 void FieldMaskTree::AddPath(const string& path) {
-  vector<string> parts = Split(path, ".");
+  std::vector<string> parts = Split(path, ".");
   if (parts.empty()) {
     return;
   }
@@ -220,8 +341,67 @@ void FieldMaskTree::AddPath(const string& path) {
   }
 }
 
+void FieldMaskTree::RemovePath(const string& path,
+                               const Descriptor* descriptor) {
+  if (root_.children.empty()) {
+    // Nothing to be removed from an empty tree. We shortcut it here so an empty
+    // tree won't be interpreted as a field mask containing all fields by the
+    // code below.
+    return;
+  }
+  std::vector<string> parts = Split(path, ".");
+  if (parts.empty()) {
+    return;
+  }
+  std::vector<Node*> nodes(parts.size());
+  Node* node = &root_;
+  const Descriptor* current_descriptor = descriptor;
+  Node* new_branch_node = nullptr;
+  for (int i = 0; i < parts.size(); ++i) {
+    nodes[i] = node;
+    const FieldDescriptor* field_descriptor =
+        current_descriptor->FindFieldByName(parts[i]);
+    if (field_descriptor == nullptr ||
+        (field_descriptor->cpp_type() != FieldDescriptor::CPPTYPE_MESSAGE &&
+         i != parts.size() - 1)) {
+      // Invalid path.
+      if (new_branch_node != nullptr) {
+        // If add any new nodes, cleanup.
+        new_branch_node->ClearChildren();
+      }
+      return;
+    }
+
+    if (node->children.empty()) {
+      if (new_branch_node == nullptr) {
+        new_branch_node = node;
+      }
+      for (int i = 0; i < current_descriptor->field_count(); ++i) {
+        node->children[current_descriptor->field(i)->name()] = new Node();
+      }
+    }
+    if (ContainsKey(node->children, parts[i])) {
+      node = node->children[parts[i]];
+      if (field_descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
+        current_descriptor = field_descriptor->message_type();
+      }
+    } else {
+      // Path does not exist.
+      return;
+    }
+  }
+  // Remove path.
+  for (int i = parts.size() - 1; i >= 0; i--) {
+    delete nodes[i]->children[parts[i]];
+    nodes[i]->children.erase(parts[i]);
+    if (!nodes[i]->children.empty()) {
+      break;
+    }
+  }
+}
+
 void FieldMaskTree::IntersectPath(const string& path, FieldMaskTree* out) {
-  vector<string> parts = Split(path, ".");
+  std::vector<string> parts = Split(path, ".");
   if (parts.empty()) {
     return;
   }
@@ -251,7 +431,7 @@ void FieldMaskTree::MergeLeafNodesToTree(const string& prefix, const Node* node,
   if (node->children.empty()) {
     out->AddPath(prefix);
   }
-  for (map<string, Node*>::const_iterator it = node->children.begin();
+  for (std::map<string, Node*>::const_iterator it = node->children.begin();
        it != node->children.end(); ++it) {
     string current_path = prefix.empty() ? it->first : prefix + "." + it->first;
     MergeLeafNodesToTree(current_path, it->second, out);
@@ -265,7 +445,7 @@ void FieldMaskTree::MergeMessage(const Node* node, const Message& source,
   const Reflection* source_reflection = source.GetReflection();
   const Reflection* destination_reflection = destination->GetReflection();
   const Descriptor* descriptor = source.GetDescriptor();
-  for (map<string, Node*>::const_iterator it = node->children.begin();
+  for (std::map<string, Node*>::const_iterator it = node->children.begin();
        it != node->children.end(); ++it) {
     const string& field_name = it->first;
     const Node* child = it->second;
@@ -291,11 +471,15 @@ void FieldMaskTree::MergeMessage(const Node* node, const Message& source,
     }
     if (!field->is_repeated()) {
       switch (field->cpp_type()) {
-#define COPY_VALUE(TYPE, Name)                                            \
-  case FieldDescriptor::CPPTYPE_##TYPE: {                                 \
-    destination_reflection->Set##Name(                                    \
-        destination, field, source_reflection->Get##Name(source, field)); \
-    break;                                                                \
+#define COPY_VALUE(TYPE, Name)                                              \
+  case FieldDescriptor::CPPTYPE_##TYPE: {                                   \
+    if (source_reflection->HasField(source, field)) {                       \
+      destination_reflection->Set##Name(                                    \
+          destination, field, source_reflection->Get##Name(source, field)); \
+    } else {                                                                \
+      destination_reflection->ClearField(destination, field);               \
+    }                                                                       \
+    break;                                                                  \
   }
         COPY_VALUE(BOOL, Bool)
         COPY_VALUE(INT32, Int32)
@@ -357,6 +541,63 @@ 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 == nullptr) {
+        // 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();
+  const Descriptor* descriptor = message->GetDescriptor();
+  const int32 field_count = descriptor->field_count();
+  for (int index = 0; index < field_count; ++index) {
+    const FieldDescriptor* field = descriptor->field(index);
+    std::map<string, Node*>::const_iterator it =
+        node->children.find(field->name());
+    if (it == node->children.end()) {
+      reflection->ClearField(message, field);
+    } else {
+      if (field->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
+        Node* child = it->second;
+        if (!child->children.empty() && reflection->HasField(*message, field)) {
+          TrimMessage(child, reflection->MutableMessage(message, field));
+        }
+      }
+    }
+  }
+}
+
 }  // namespace
 
 void FieldMaskUtil::ToCanonicalForm(const FieldMask& mask, FieldMask* out) {
@@ -386,6 +627,22 @@ void FieldMaskUtil::Intersect(const FieldMask& mask1, const FieldMask& mask2,
   intersection.MergeToFieldMask(out);
 }
 
+void FieldMaskUtil::InternalSubtract(const Descriptor* descriptor,
+                                     const FieldMask& mask1,
+                                     const FieldMask& mask2, FieldMask* out) {
+  if (mask1.paths().empty()) {
+    out->Clear();
+    return;
+  }
+  FieldMaskTree tree;
+  tree.MergeFromFieldMask(mask1);
+  for (int i = 0; i < mask2.paths_size(); ++i) {
+    tree.RemovePath(mask2.paths(i), descriptor);
+  }
+  out->Clear();
+  tree.MergeToFieldMask(out);
+}
+
 bool FieldMaskUtil::IsPathInFieldMask(StringPiece path, const FieldMask& mask) {
   for (int i = 0; i < mask.paths_size(); ++i) {
     const string& mask_path = mask.paths(i);
@@ -413,6 +670,28 @@ void FieldMaskUtil::MergeMessageTo(const Message& source, const FieldMask& mask,
   tree.MergeMessage(source, options, destination);
 }
 
+void FieldMaskUtil::TrimMessage(const FieldMask& mask, Message* destination) {
+  // Build a FieldMaskTree and walk through the tree to merge all specified
+  // fields.
+  FieldMaskTree tree;
+  tree.MergeFromFieldMask(mask);
+  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