[Object] Change OffloadBinary::write to return SmallString<0>

SmallString<0> is more flexible and avoids an unneeded copy in
ObjectYAML/OffloadEmitter.cpp.

Reviewed By: jhuber6

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

Author: MaskRay

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

@@ -184,7 +184,8 @@ Expected<OffloadFile> getInputBitcodeLibrary(StringRef Input) {
   Image.StringData["arch"] = Arch;
   Image.Image = std::move(*ImageOrError);
 
-  std::unique_ptr<MemoryBuffer> Binary = OffloadBinary::write(Image);
+  std::unique_ptr<MemoryBuffer> Binary =
+      MemoryBuffer::getMemBufferCopy(OffloadBinary::write(Image));
   auto NewBinaryOrErr = OffloadBinary::create(*Binary);
   if (!NewBinaryOrErr)
     return NewBinaryOrErr.takeError();
@@ -909,7 +910,8 @@ Expected<SmallVector<std::unique_ptr<MemoryBuffer>>>
 bundleOpenMP(ArrayRef<OffloadingImage> Images) {
   SmallVector<std::unique_ptr<MemoryBuffer>> Buffers;
   for (const OffloadingImage &Image : Images)
-    Buffers.emplace_back(OffloadBinary::write(Image));
+    Buffers.emplace_back(
+        MemoryBuffer::getMemBufferCopy(OffloadBinary::write(Image)));
 
   return std::move(Buffers);
 }

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

@@ -126,11 +126,11 @@ static Error bundleImages() {
         ImageBinary.StringData[Key] = Value;
       }
     }
-    std::unique_ptr<MemoryBuffer> Buffer = OffloadBinary::write(ImageBinary);
-    if (Buffer->getBufferSize() % OffloadBinary::getAlignment() != 0)
+    llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary);
+    if (Buffer.size() % OffloadBinary::getAlignment() != 0)
       return createStringError(inconvertibleErrorCode(),
                                "Offload binary has invalid size alignment");
-    OS << Buffer->getBuffer();
+    OS << Buffer;
   }
 
   if (Error E = writeFile(OutputFile,

llvm/include/llvm/Object/OffloadBinary.h

@@ -18,6 +18,7 @@
 #define LLVM_OBJECT_OFFLOADBINARY_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Support/Error.h"
@@ -78,7 +79,7 @@ class OffloadBinary : public Binary {
   static Expected<std::unique_ptr<OffloadBinary>> create(MemoryBufferRef);
 
   /// Serialize the contents of \p File to a binary buffer to be read later.
-  static std::unique_ptr<MemoryBuffer> write(const OffloadingImage &);
+  static SmallString<0> write(const OffloadingImage &);
 
   static uint64_t getAlignment() { return 8; }
 

llvm/lib/Object/OffloadBinary.cpp

@@ -204,8 +204,7 @@ OffloadBinary::create(MemoryBufferRef Buf) {
       new OffloadBinary(Buf, TheHeader, TheEntry));
 }
 
-std::unique_ptr<MemoryBuffer>
-OffloadBinary::write(const OffloadingImage &OffloadingData) {
+SmallString<0> OffloadBinary::write(const OffloadingImage &OffloadingData) {
   // Create a null-terminated string table with all the used strings.
   StringTableBuilder StrTab(StringTableBuilder::ELF);
   for (auto &KeyAndValue : OffloadingData.StringData) {
@@ -243,7 +242,7 @@ OffloadBinary::write(const OffloadingImage &OffloadingData) {
   TheEntry.ImageOffset = BinaryDataSize;
   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
 
-  SmallVector<char> Data;
+  SmallString<0> Data;
   Data.reserve(TheHeader.Size);
   raw_svector_ostream OS(Data);
   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
@@ -264,7 +263,7 @@ OffloadBinary::write(const OffloadingImage &OffloadingData) {
   OS.write_zeros(TheHeader.Size - OS.tell());
   assert(TheHeader.Size == OS.tell() && "Size mismatch");
 
-  return MemoryBuffer::getMemBufferCopy(OS.str());
+  return Data;
 }
 
 Error object::extractOffloadBinaries(MemoryBufferRef Buffer,

llvm/lib/ObjectYAML/OffloadEmitter.cpp

@@ -38,14 +38,10 @@ bool yaml2offload(Binary &Doc, raw_ostream &Out, ErrorHandler EH) {
       Member.Content->writeAsBinary(OS);
     Image.Image = MemoryBuffer::getMemBufferCopy(OS.str());
 
-    std::unique_ptr<MemoryBuffer> Binary = object::OffloadBinary::write(Image);
-
     // Copy the data to a new buffer so we can modify the bytes directly.
-    SmallVector<char> NewBuffer;
-    std::copy(Binary->getBufferStart(), Binary->getBufferEnd(),
-              std::back_inserter(NewBuffer));
+    auto Buffer = object::OffloadBinary::write(Image);
     auto *TheHeader =
-        reinterpret_cast<object::OffloadBinary::Header *>(&NewBuffer[0]);
+        reinterpret_cast<object::OffloadBinary::Header *>(&Buffer[0]);
     if (Doc.Version)
       TheHeader->Version = *Doc.Version;
     if (Doc.Size)
@@ -55,7 +51,7 @@ bool yaml2offload(Binary &Doc, raw_ostream &Out, ErrorHandler EH) {
     if (Doc.EntrySize)
       TheHeader->EntrySize = *Doc.EntrySize;
 
-    Out.write(NewBuffer.begin(), NewBuffer.size());
+    Out.write(Buffer.begin(), Buffer.size());
   }
 
   return true;

llvm/unittests/Object/OffloadingTest.cpp

@@ -43,8 +43,8 @@ TEST(OffloadingTest, checkOffloadingBinary) {
   Data.StringData = StringData;
   Data.Image = std::move(ImageData);
 
-  auto BinaryBuffer = OffloadBinary::write(Data);
-
+  auto BinaryBuffer =
+      MemoryBuffer::getMemBufferCopy(OffloadBinary::write(Data));
   auto BinaryOrErr = OffloadBinary::create(*BinaryBuffer);
   if (!BinaryOrErr)
     FAIL();