[llvm-objcopy][NFC] replace class Buffer/MemBuffer/FileBuffer with streams.

During D88827 it was requested to remove the local implementation
of Memory/File Buffers:

// TODO: refactor the buffer classes in LLVM to enable us to use them here
// directly.

This patch uses raw_ostream instead of Buffers. Generally, using streams
could allow us to reduce memory usages. No need to load all data into the
memory - the data could be streamed through a smaller buffer.
Thus, this patch uses raw_ostream as an interface for output data:

Error executeObjcopyOnBinary(CopyConfig &Config,
                             object::Binary &In,
                             raw_ostream &Out);

Note 1. This patch does not change the implementation of Writers
so that data would be directly stored into raw_ostream.
This is assumed to be done later.

Note 2. It would be better if Writers would be implemented in a such way
that data could be streamed without seeking/updating. If that would be
inconvenient then raw_ostream could be replaced with raw_pwrite_stream
to have a possibility to seek back and update file headers.
This is assumed to be done later if necessary.

Note 3. Current FileOutputBuffer allows using a memory-mapped file.
The raw_fd_ostream (which could be used if data should be stored in the file)
does not allow us to use a memory-mapped file. Memory map functionality
could be implemented for raw_fd_ostream:

It is possible to add resize() method into raw_ostream.

class raw_ostream {
  void resize(uint64_t size);
}

That method, implemented for raw_fd_ostream, could create a memory-mapped file.
The streamed data would be written into that memory file then.
Thus we would be able to use memory-mapped files with raw_fd_ostream.
This is assumed to be done later if necessary.

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

1 files changed, 76 insertions, 24 deletions

diff --git a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
index 43e3334..bc32473 100644
--- a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
+++ b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "Buffer.h"
+#include "llvm-objcopy.h"
 #include "COFF/COFFObjcopy.h"
 #include "CopyConfig.h"
 #include "ELF/ELFObjcopy.h"
@@ -17,6 +17,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Object/Archive.h"
 #include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
@@ -32,6 +33,7 @@
 #include "llvm/Option/Option.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ErrorOr.h"
@@ -40,6 +42,7 @@
 #include "llvm/Support/Memory.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Process.h"
+#include "llvm/Support/SmallVectorMemoryBuffer.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Support/raw_ostream.h"
@@ -54,6 +57,44 @@
 namespace llvm {
 namespace objcopy {
 
+Error writeToFile(StringRef OutputFileName,
+                  std::function<Error(raw_ostream &)> Write, bool KeepOwnership,
+                  unsigned UserID, unsigned GroupID) {
+  if (OutputFileName == "-")
+    return Write(outs());
+
+  if (OutputFileName == "/dev/null") {
+    raw_null_ostream Out;
+    return Write(Out);
+  }
+
+  unsigned Mode = sys::fs::all_read | sys::fs::all_write | sys::fs::all_exe;
+  Expected<sys::fs::TempFile> Temp =
+      sys::fs::TempFile::create(OutputFileName + ".temp-objcopy-%%%%%%", Mode);
+  if (!Temp)
+    return createFileError(OutputFileName, Temp.takeError());
+
+#ifndef _WIN32
+  // Try to preserve file ownership if requested.
+  if (KeepOwnership) {
+    sys::fs::file_status Stat;
+    if (!sys::fs::status(Temp->FD, Stat) && Stat.getUser() == 0)
+      sys::fs::changeFileOwnership(Temp->FD, UserID, GroupID);
+  }
+#endif
+
+  raw_fd_ostream Out(Temp->FD, false);
+
+  if (Error E = Write(Out)) {
+    if (Error DiscardError = Temp->discard())
+      return joinErrors(std::move(E), std::move(DiscardError));
+    return E;
+  }
+  Out.flush();
+
+  return Temp->keep(OutputFileName);
+}
+
 // The name this program was invoked as.
 StringRef ToolName;
 
@@ -108,20 +149,21 @@ static Error deepWriteArchive(StringRef ArcName,
     return Error::success();
 
   for (const NewArchiveMember &Member : NewMembers) {
-    // Internally, FileBuffer will use the buffer created by
-    // FileOutputBuffer::create, for regular files (that is the case for
-    // deepWriteArchive) FileOutputBuffer::create will return OnDiskBuffer.
+    // For regular files (as is the case for deepWriteArchive),
+    // FileOutputBuffer::create will return OnDiskBuffer.
     // OnDiskBuffer uses a temporary file and then renames it. So in reality
     // there is no inefficiency / duplicated in-memory buffers in this case. For
     // now in-memory buffers can not be completely avoided since
     // NewArchiveMember still requires them even though writeArchive does not
     // write them on disk.
-    FileBuffer FB(Member.MemberName);
-    if (Error E = FB.allocate(Member.Buf->getBufferSize()))
-      return E;
+    Expected<std::unique_ptr<FileOutputBuffer>> FB = FileOutputBuffer::create(
+        Member.MemberName, Member.Buf->getBufferSize(),
+        FileOutputBuffer::F_executable | FileOutputBuffer::F_keep_ownership);
+    if (!FB)
+      return FB.takeError();
     std::copy(Member.Buf->getBufferStart(), Member.Buf->getBufferEnd(),
-              FB.getBufferStart());
-    if (Error E = FB.commit())
+              (*FB)->getBufferStart());
+    if (Error E = (*FB)->commit())
       return E;
   }
   return Error::success();
@@ -130,7 +172,7 @@ static Error deepWriteArchive(StringRef ArcName,
 /// The function executeObjcopyOnIHex does the dispatch based on the format
 /// of the output specified by the command line options.
 static Error executeObjcopyOnIHex(CopyConfig &Config, MemoryBuffer &In,
-                                  Buffer &Out) {
+                                  raw_ostream &Out) {
   // TODO: support output formats other than ELF.
   if (Error E = Config.parseELFConfig())
     return E;
@@ -140,7 +182,7 @@ static Error executeObjcopyOnIHex(CopyConfig &Config, MemoryBuffer &In,
 /// The function executeObjcopyOnRawBinary does the dispatch based on the format
 /// of the output specified by the command line options.
 static Error executeObjcopyOnRawBinary(CopyConfig &Config, MemoryBuffer &In,
-                                       Buffer &Out) {
+                                       raw_ostream &Out) {
   switch (Config.OutputFormat) {
   case FileFormat::ELF:
   // FIXME: Currently, we call elf::executeObjcopyOnRawBinary even if the
@@ -160,7 +202,7 @@ static Error executeObjcopyOnRawBinary(CopyConfig &Config, MemoryBuffer &In,
 /// The function executeObjcopyOnBinary does the dispatch based on the format
 /// of the input binary (ELF, MachO or COFF).
 static Error executeObjcopyOnBinary(CopyConfig &Config, object::Binary &In,
-                                    Buffer &Out) {
+                                    raw_ostream &Out) {
   if (auto *ELFBinary = dyn_cast<object::ELFObjectFileBase>(&In)) {
     if (Error E = Config.parseELFConfig())
       return E;
@@ -197,15 +239,19 @@ createNewArchiveMembers(CopyConfig &Config, const Archive &Ar) {
       return createFileError(Ar.getFileName() + "(" + *ChildNameOrErr + ")",
                              ChildOrErr.takeError());
 
-    MemBuffer MB(ChildNameOrErr.get());
-    if (Error E = executeObjcopyOnBinary(Config, *ChildOrErr->get(), MB))
+    SmallVector<char, 0> Buffer;
+    raw_svector_ostream MemStream(Buffer);
+
+    if (Error E = executeObjcopyOnBinary(Config, *ChildOrErr->get(), MemStream))
       return std::move(E);
 
     Expected<NewArchiveMember> Member =
         NewArchiveMember::getOldMember(Child, Config.DeterministicArchives);
     if (!Member)
       return createFileError(Ar.getFileName(), Member.takeError());
-    Member->Buf = MB.releaseMemoryBuffer();
+
+    Member->Buf = std::make_unique<SmallVectorMemoryBuffer>(
+        std::move(Buffer), ChildNameOrErr.get());
     Member->MemberName = Member->Buf->getBufferIdentifier();
     NewArchiveMembers.push_back(std::move(*Member));
   }
@@ -280,7 +326,7 @@ static Error executeObjcopy(CopyConfig &Config) {
     Stat.permissions(static_cast<sys::fs::perms>(0777));
   }
 
-  using ProcessRawFn = Error (*)(CopyConfig &, MemoryBuffer &, Buffer &);
+  using ProcessRawFn = Error (*)(CopyConfig &, MemoryBuffer &, raw_ostream &);
   ProcessRawFn ProcessRaw;
   switch (Config.InputFormat) {
   case FileFormat::Binary:
@@ -297,8 +343,11 @@ static Error executeObjcopy(CopyConfig &Config) {
     auto BufOrErr = MemoryBuffer::getFileOrSTDIN(Config.InputFilename);
     if (!BufOrErr)
       return createFileError(Config.InputFilename, BufOrErr.getError());
-    FileBuffer FB(Config.OutputFilename);
-    if (Error E = ProcessRaw(Config, *BufOrErr->get(), FB))
+
+    if (Error E = writeToFile(
+            Config.OutputFilename, [&](raw_ostream &OutFile) -> Error {
+              return ProcessRaw(Config, *BufOrErr->get(), OutFile);
+            }))
       return E;
   } else {
     Expected<OwningBinary<llvm::object::Binary>> BinaryOrErr =
@@ -310,12 +359,15 @@ static Error executeObjcopy(CopyConfig &Config) {
       if (Error E = executeObjcopyOnArchive(Config, *Ar))
         return E;
     } else {
-      FileBuffer FB(Config.OutputFilename,
-                    Config.InputFilename != "-" &&
-                        Config.InputFilename == Config.OutputFilename,
-                    Stat.getUser(), Stat.getGroup());
-      if (Error E = executeObjcopyOnBinary(Config,
-                                           *BinaryOrErr.get().getBinary(), FB))
+      if (Error E = writeToFile(
+              Config.OutputFilename,
+              [&](raw_ostream &OutFile) -> Error {
+                return executeObjcopyOnBinary(
+                    Config, *BinaryOrErr.get().getBinary(), OutFile);
+              },
+              Config.InputFilename != "-" &&
+                  Config.InputFilename == Config.OutputFilename,
+              Stat.getUser(), Stat.getGroup()))
         return E;
     }
   }