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 @@ -124,6 +124,11 @@ return Raw; } + MemoryBufferRef getBinary() const { + assert(getKind() == Type::Binary); + return MemoryBufferRef(Raw, ""); + } + /// Get an ArrayDocNode for an array node. If Convert, convert the node to an /// array node if necessary. ArrayDocNode &getArray(bool Convert = false) { @@ -201,6 +206,7 @@ /// that restriction. DocNode &operator=(const char *Val) { return *this = StringRef(Val); } DocNode &operator=(StringRef Val); + DocNode &operator=(MemoryBufferRef Val); DocNode &operator=(bool Val); DocNode &operator=(int Val); DocNode &operator=(unsigned Val); @@ -368,6 +374,17 @@ return getNode(StringRef(V), Copy); } + /// Create a Binary node associated with this Document. If !Copy, the passed + /// buffer must remain valid for the lifetime of the Document. + DocNode getNode(MemoryBufferRef V, bool Copy = false) { + auto Raw = V.getBuffer(); + if (Copy) + Raw = addString(Raw); + auto N = DocNode(&KindAndDocs[size_t(Type::Binary)]); + N.Raw = Raw; + return N; + } + /// Create an empty Map node associated with this Document. MapDocNode getMapNode() { auto N = DocNode(&KindAndDocs[size_t(Type::Map)]); 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 @@ -80,6 +80,10 @@ *this = getDocument()->getNode(Val); return *this; } +DocNode &DocNode::operator=(MemoryBufferRef Val) { + *this = getDocument()->getNode(Val); + return *this; +} DocNode &DocNode::operator=(bool Val) { *this = getDocument()->getNode(Val); return *this; @@ -167,6 +171,9 @@ case Type::String: Node = getNode(Obj.Raw); break; + case Type::Binary: + Node = getNode(MemoryBufferRef(Obj.Raw, "")); + break; case Type::Map: Node = getMapNode(); break; @@ -277,6 +284,9 @@ case Type::String: MPWriter.write(Node.getString()); break; + case Type::Binary: + MPWriter.write(Node.getBinary()); + break; case Type::Empty: llvm_unreachable("unhandled empty msgpack node"); default: 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 @@ -30,6 +30,17 @@ ASSERT_EQ(Doc.getRoot().getInt(), 0); } +TEST(MsgPackDocument, TestReadBinary) { + Document Doc; + uint8_t data[] = {1, 2, 3, 4}; + bool Ok = + Doc.readFromBlob(StringRef("\xC4\x4\x1\x2\x3\x4", 6), /*Multi=*/false); + ASSERT_TRUE(Ok); + ASSERT_EQ(Doc.getRoot().getKind(), Type::Binary); + ASSERT_EQ(Doc.getRoot().getBinary().getBuffer(), + StringRef(reinterpret_cast(data), 4)); +} + TEST(MsgPackDocument, TestReadMergeArray) { Document Doc; bool Ok = Doc.readFromBlob(StringRef("\x92\xd0\x01\xc0"), /*Multi=*/false); @@ -189,6 +200,16 @@ ASSERT_EQ(Buffer, "\x01"); } +TEST(MsgPackDocument, TestWriteBinary) { + uint8_t data[] = {1, 2, 3, 4}; + Document Doc; + Doc.getRoot() = MemoryBufferRef( + StringRef(reinterpret_cast(data), sizeof(data)), ""); + std::string Buffer; + Doc.writeToBlob(Buffer); + ASSERT_EQ(Buffer, "\xC4\x4\x1\x2\x3\x4"); +} + TEST(MsgPackDocument, TestWriteArray) { Document Doc; auto A = Doc.getRoot().getArray(/*Convert=*/true);