[MsgPack] Added convenience assignment to MsgPackDocument

This commit increases the convenience of using the MsgPackDocument API,
especially when creating a document for writing out.

It adds direct assignment of bool, integer and string types to a
DocNode, as long as that DocNode is already inside a document, e.g. the
result of a map lookup. It also adds map lookup given an integer type
(it already had that for string).

So, to assign a string to a map element whose key is an int, you can
now write
  MyMap[42] = "towel";
instead of
  MyMap[MyMap.getDocument()->getNode(42)] =
      MyMap.getDocument()->getNode("towel");

Also added MapDocNode::erase methods.

Differential Revision: https://reviews.llvm.org/D80121

Change-Id: I17301fa15bb9802231c52542798af5b54beb583e

llvm/include/llvm/BinaryFormat/MsgPackDocument.h

@@ -189,6 +189,19 @@ public:
   /// not rely on S having a lifetime beyond this call. Tag is "" or a YAML tag.
   StringRef fromString(StringRef S, StringRef Tag = "");
 
+  /// Convenience assignment operators. This only works if the destination
+  /// DocNode has an associated Document, i.e. it was not constructed using the
+  /// default constructor. The string one does not copy, so the string must
+  /// remain valid for the lifetime of the Document. Use fromString to avoid
+  /// that restriction.
+  DocNode &operator=(const char *Val) { return *this = StringRef(Val); }
+  DocNode &operator=(StringRef Val);
+  DocNode &operator=(bool Val);
+  DocNode &operator=(int Val);
+  DocNode &operator=(unsigned Val);
+  DocNode &operator=(int64_t Val);
+  DocNode &operator=(uint64_t Val);
+
 private:
   // Private constructor setting KindAndDoc, used by methods in Document.
   DocNode(const KindAndDocument *KindAndDoc) : KindAndDoc(KindAndDoc) {}
@@ -210,11 +223,21 @@ public:
   MapTy::iterator end() { return Map->end(); }
   MapTy::iterator find(DocNode Key) { return Map->find(Key); }
   MapTy::iterator find(StringRef Key);
+  MapTy::iterator erase(MapTy::const_iterator I) { return Map->erase(I); }
+  size_t erase(DocNode Key) { return Map->erase(Key); }
+  MapTy::iterator erase(MapTy::const_iterator First,
+                        MapTy::const_iterator Second) {
+    return Map->erase(First, Second);
+  }
   /// Member access. The string data must remain valid for the lifetime of the
   /// Document.
   DocNode &operator[](StringRef S);
-  /// Member access.
+  /// Member access, with convenience versions for an integer key.
   DocNode &operator[](DocNode Key);
+  DocNode &operator[](int Key);
+  DocNode &operator[](unsigned Key);
+  DocNode &operator[](int64_t Key);
+  DocNode &operator[](uint64_t Key);
 };
 
 /// A DocNode that is an array.

llvm/lib/BinaryFormat/MsgPackDocument.cpp

@@ -48,6 +48,20 @@ DocNode &MapDocNode::operator[](DocNode Key) {
   return N;
 }
 
+/// Member access for MapDocNode for integer key.
+DocNode &MapDocNode::operator[](int Key) {
+  return (*this)[getDocument()->getNode(Key)];
+}
+DocNode &MapDocNode::operator[](unsigned Key) {
+  return (*this)[getDocument()->getNode(Key)];
+}
+DocNode &MapDocNode::operator[](int64_t Key) {
+  return (*this)[getDocument()->getNode(Key)];
+}
+DocNode &MapDocNode::operator[](uint64_t Key) {
+  return (*this)[getDocument()->getNode(Key)];
+}
+
 /// Array element access. This extends the array if necessary.
 DocNode &ArrayDocNode::operator[](size_t Index) {
   if (size() <= Index) {
@@ -57,6 +71,36 @@ DocNode &ArrayDocNode::operator[](size_t Index) {
   return (*Array)[Index];
 }
 
+// Convenience assignment operators. This only works if the destination
+// DocNode has an associated Document, i.e. it was not constructed using the
+// default constructor. The string one does not copy, so the string must
+// remain valid for the lifetime of the Document. Use fromString to avoid
+// that restriction.
+DocNode &DocNode::operator=(StringRef Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
+DocNode &DocNode::operator=(bool Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
+DocNode &DocNode::operator=(int Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
+DocNode &DocNode::operator=(unsigned Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
+DocNode &DocNode::operator=(int64_t Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
+DocNode &DocNode::operator=(uint64_t Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
+
 // A level in the document reading stack.
 struct StackLevel {
   StackLevel(DocNode Node, size_t StartIndex, size_t Length,

llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp

@@ -174,7 +174,7 @@ TEST(MsgPackDocument, TestReadMergeMap) {
 
 TEST(MsgPackDocument, TestWriteInt) {
   Document Doc;
-  Doc.getRoot() = Doc.getNode(int64_t(1));
+  Doc.getRoot() = 1;
   std::string Buffer;
   Doc.writeToBlob(Buffer);
   ASSERT_EQ(Buffer, "\x01");
@@ -193,8 +193,8 @@ TEST(MsgPackDocument, TestWriteArray) {
 TEST(MsgPackDocument, TestWriteMap) {
   Document Doc;
   auto M = Doc.getRoot().getMap(/*Convert=*/true);
-  M["foo"] = Doc.getNode(int64_t(1));
-  M["bar"] = Doc.getNode(int64_t(2));
+  M["foo"] = 1;
+  M["bar"] = 2;
   std::string Buffer;
   Doc.writeToBlob(Buffer);
   ASSERT_EQ(Buffer, "\x82\xa3"
@@ -233,11 +233,11 @@ TEST(MsgPackDocument, TestInputYAMLArray) {
 TEST(MsgPackDocument, TestOutputYAMLMap) {
   Document Doc;
   auto M = Doc.getRoot().getMap(/*Convert=*/true);
-  M["foo"] = Doc.getNode(int64_t(1));
-  M["bar"] = Doc.getNode(uint64_t(2));
+  M["foo"] = 1;
+  M["bar"] = 2U;
   auto N = Doc.getMapNode();
   M["qux"] = N;
-  N["baz"] = Doc.getNode(true);
+  N["baz"] = true;
   std::string Buffer;
   raw_string_ostream OStream(Buffer);
   Doc.toYAML(OStream);
@@ -249,15 +249,34 @@ TEST(MsgPackDocument, TestOutputYAMLMap) {
                            "...\n");
 }
 
+TEST(MsgPackDocument, TestOutputYAMLMapWithErase) {
+  Document Doc;
+  auto M = Doc.getRoot().getMap(/*Convert=*/true);
+  M["foo"] = 1;
+  M["bar"] = 2U;
+  auto N = Doc.getMapNode();
+  M["qux"] = N;
+  N["baz"] = true;
+  M.erase(Doc.getNode("bar"));
+  std::string Buffer;
+  raw_string_ostream OStream(Buffer);
+  Doc.toYAML(OStream);
+  ASSERT_EQ(OStream.str(), "---\n"
+                           "foo:             1\n"
+                           "qux:\n"
+                           "  baz:             true\n"
+                           "...\n");
+}
+
 TEST(MsgPackDocument, TestOutputYAMLMapHex) {
   Document Doc;
   Doc.setHexMode();
   auto M = Doc.getRoot().getMap(/*Convert=*/true);
-  M["foo"] = Doc.getNode(int64_t(1));
-  M["bar"] = Doc.getNode(uint64_t(2));
+  M["foo"] = 1;
+  M["bar"] = 2U;
   auto N = Doc.getMapNode();
   M["qux"] = N;
-  N["baz"] = Doc.getNode(true);
+  N["baz"] = true;
   std::string Buffer;
   raw_string_ostream OStream(Buffer);
   Doc.toYAML(OStream);