[MsgPack] Add support for writing binary ELF doc nodes

Will be used for a front-end compiler https://github.com/GPUOpen-Drivers/llpc to insert a binary blob value into AMDGPU PAL metadata. Differential Revision: https://reviews.llvm.org/D148607 Change-Id: I90685883d1478335937218e6b1a569b1b26dc05a
llvm · Apr 20, 2023 · 9d52f69 · 9d52f69
1 parent fdaa23a
commit 9d52f69
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 0 deletions.
diff --git a/llvm/include/llvm/BinaryFormat/MsgPackDocument.h b/llvm/include/llvm/BinaryFormat/MsgPackDocument.h
@@ -124,6 +124,11 @@ class DocNode {
     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 @@ class DocNode {
   /// 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 @@ class Document {
     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
@@ -80,6 +80,10 @@ DocNode &DocNode::operator=(StringRef Val) {
   *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 @@ bool Document::readFromBlob(
     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 @@ void Document::writeToBlob(std::string &Blob) {
     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
@@ -30,6 +30,17 @@ TEST(MsgPackDocument, TestReadInt) {
   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<const char *>(data), 4));
+}
+
 TEST(MsgPackDocument, TestReadMergeArray) {
   Document Doc;
   bool Ok = Doc.readFromBlob(StringRef("\x92\xd0\x01\xc0"), /*Multi=*/false);
@@ -189,6 +200,16 @@ TEST(MsgPackDocument, TestWriteInt) {
   ASSERT_EQ(Buffer, "\x01");
 }
 
+TEST(MsgPackDocument, TestWriteBinary) {
+  uint8_t data[] = {1, 2, 3, 4};
+  Document Doc;
+  Doc.getRoot() = MemoryBufferRef(
+      StringRef(reinterpret_cast<const char *>(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);