[MsgPack] Added convenience assignment to MsgPackDocument

[tpr]

Summary

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");

Change-Id: I17301fa15bb9802231c52542798af5b54beb583e

[tpr]

V2: Added tests

[tpr]

V3: Added assignment for const char *, otherwise it defaults to using

the bool one rather than the StringRef one.

[tpr]

V4: Added MapDocNode::erase methods.

llvm/include/llvm/BinaryFormat/MsgPackDocument.h

public:
   /// Convert this node to a string, assuming it is scalar.
   std::string toString() const;
 
   /// Convert the StringRef and use it to set this DocNode (assuming scalar). If
   /// it is a string, copy the string into the Document's strings list so we do
   /// 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) {}
 
   void convertToArray();
   void convertToMap();
 };
 
 /// A DocNode that is a map.
 class MapDocNode : public DocNode {
 public:
   MapDocNode() {}
   MapDocNode(DocNode &N) : DocNode(N) { assert(getKind() == Type::Map); }
 
   // Map access methods.
   size_t size() const { return Map->size(); }
   bool empty() const { return !size(); }
   MapTy::iterator begin() { return Map->begin(); }
   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.
 class ArrayDocNode : public DocNode {
 public:
   ArrayDocNode() {}
   ArrayDocNode(DocNode &N) : DocNode(N) { assert(getKind() == Type::Array); }
 

llvm/lib/BinaryFormat/MsgPackDocument.cpp

DocNode &MapDocNode::operator[](DocNode Key) {
   DocNode &N = (*Map)[Key];
   if (N.isEmpty()) {
     // Ensure a new element has its KindAndDoc initialized.
     N = getDocument()->getEmptyNode();
   }
   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) {
     // Ensure new elements have their KindAndDoc initialized.
     Array->resize(Index + 1, getDocument()->getEmptyNode());
   }
   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,
              DocNode *MapEntry = nullptr)
       : Node(Node), Index(StartIndex), End(StartIndex + Length),
         MapEntry(MapEntry) {}
   DocNode Node;
   size_t Index;

llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp

TEST(MsgPackDocument, TestReadMergeMap) {
   ASSERT_EQ(BazS.getInt(), 4);
   auto BayS = M["bay"];
   ASSERT_EQ(BayS.getKind(), Type::Int);
   ASSERT_EQ(BayS.getInt(), 8);
 }
 
 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");
 }
 
 TEST(MsgPackDocument, TestWriteArray) {
   Document Doc;
   auto A = Doc.getRoot().getArray(/*Convert=*/true);
   A.push_back(Doc.getNode(int64_t(1)));
   A.push_back(Doc.getNode());
   std::string Buffer;
   Doc.writeToBlob(Buffer);
   ASSERT_EQ(Buffer, "\x92\x01\xc0");
 }
 
 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"
                     "bar"
                     "\x02\xa3"
                     "foo"
                     "\x01");
 }
TEST(MsgPackDocument, TestInputYAMLArray) {
   auto SS = A[1];
   ASSERT_EQ(SS.getKind(), Type::String);
   ASSERT_EQ(SS.getString(), "2");
 }
 
 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);
   ASSERT_EQ(OStream.str(), "---\n"
                            "bar:             2\n"
                            "foo:             1\n"
                            "qux:\n"
                            "  baz:             true\n"
                            "...\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);
   ASSERT_EQ(OStream.str(), "---\n"
                            "bar:             0x2\n"
                            "foo:             1\n"
                            "qux:\n"
                            "  baz:             true\n"