diff --git a/llvm/include/llvm/BinaryFormat/MsgPackDocument.h b/llvm/include/llvm/BinaryFormat/MsgPackDocument.h --- a/llvm/include/llvm/BinaryFormat/MsgPackDocument.h +++ b/llvm/include/llvm/BinaryFormat/MsgPackDocument.h @@ -193,6 +193,18 @@ /// 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=(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) {} @@ -217,8 +229,12 @@ /// 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. diff --git a/llvm/lib/BinaryFormat/MsgPackDocument.cpp b/llvm/lib/BinaryFormat/MsgPackDocument.cpp --- a/llvm/lib/BinaryFormat/MsgPackDocument.cpp +++ b/llvm/lib/BinaryFormat/MsgPackDocument.cpp @@ -48,6 +48,20 @@ 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 @@ 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 Length, DocNode *MapEntry = nullptr) diff --git a/llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp b/llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp --- a/llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp +++ b/llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp @@ -135,7 +135,7 @@ 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"); @@ -154,8 +154,8 @@ 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" @@ -194,11 +194,11 @@ 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); @@ -214,11 +214,11 @@ 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);