1 files changed, 54 insertions, 17 deletions

diff --git a/llvm/tools/llvm-objcopy/Object.h b/llvm/tools/llvm-objcopy/Object.h
index 48de0e543a0..2e20b5b299f 100644
--- a/llvm/tools/llvm-objcopy/Object.h
+++ b/llvm/tools/llvm-objcopy/Object.h
@@ -27,6 +27,7 @@
 
 namespace llvm {
 
+class Buffer;
 class SectionBase;
 class Section;
 class OwnedDataSection;
@@ -77,7 +78,7 @@ public:
 
 class SectionWriter : public SectionVisitor {
 protected:
-  FileOutputBuffer &Out;
+  Buffer &Out;
 
 public:
   virtual ~SectionWriter(){};
@@ -91,7 +92,7 @@ public:
   virtual void visit(const GnuDebugLinkSection &Sec) override = 0;
   virtual void visit(const GroupSection &Sec) override = 0;
 
-  SectionWriter(FileOutputBuffer &Buf) : Out(Buf) {}
+  explicit SectionWriter(Buffer &Buf) : Out(Buf) {}
 };
 
 template <class ELFT> class ELFSectionWriter : public SectionWriter {
@@ -107,7 +108,7 @@ public:
   void visit(const GnuDebugLinkSection &Sec) override;
   void visit(const GroupSection &Sec) override;
 
-  ELFSectionWriter(FileOutputBuffer &Buf) : SectionWriter(Buf) {}
+  explicit ELFSectionWriter(Buffer &Buf) : SectionWriter(Buf) {}
 };
 
 #define MAKE_SEC_WRITER_FRIEND                                                 \
@@ -123,24 +124,62 @@ public:
   void visit(const GnuDebugLinkSection &Sec) override;
   void visit(const GroupSection &Sec) override;
 
-  BinarySectionWriter(FileOutputBuffer &Buf) : SectionWriter(Buf) {}
+  explicit BinarySectionWriter(Buffer &Buf) : SectionWriter(Buf) {}
+};
+
+// The class Buffer abstracts out the common interface of FileOutputBuffer and
+// WritableMemoryBuffer so that the hierarchy of Writers depends on this
+// abstract interface and doesn't depend on a particular implementation.
+// TODO: refactor the buffer classes in LLVM to enable us to use them here
+// directly.
+class Buffer {
+  StringRef Name;
+
+public:
+  virtual ~Buffer();
+  virtual void allocate(size_t Size) = 0;
+  virtual uint8_t *getBufferStart() = 0;
+  virtual Error commit() = 0;
+
+  explicit Buffer(StringRef Name) : Name(Name) {}
+  StringRef getName() const { return Name; }
+};
+
+class FileBuffer : public Buffer {
+  std::unique_ptr<FileOutputBuffer> Buf;
+
+public:
+  void allocate(size_t Size) override;
+  uint8_t *getBufferStart() override;
+  Error commit() override;
+
+  explicit FileBuffer(StringRef FileName) : Buffer(FileName) {}
+};
+
+class MemBuffer : public Buffer {
+  std::unique_ptr<WritableMemoryBuffer> Buf;
+
+public:
+  void allocate(size_t Size) override;
+  uint8_t *getBufferStart() override;
+  Error commit() override;
+
+  explicit MemBuffer(StringRef Name) : Buffer(Name) {}
+
+  std::unique_ptr<WritableMemoryBuffer> releaseMemoryBuffer();
 };
 
 class Writer {
 protected:
-  StringRef File;
   Object &Obj;
-  std::unique_ptr<FileOutputBuffer> BufPtr;
-
-  void createBuffer(uint64_t Size);
+  Buffer &Buf;
 
 public:
   virtual ~Writer();
-
   virtual void finalize() = 0;
   virtual void write() = 0;
 
-  Writer(StringRef File, Object &Obj) : File(File), Obj(Obj) {}
+  Writer(Object &O, Buffer &B) : Obj(O), Buf(B) {}
 };
 
 template <class ELFT> class ELFWriter : public Writer {
@@ -169,8 +208,8 @@ public:
 
   void finalize() override;
   void write() override;
-  ELFWriter(StringRef File, Object &Obj, bool WSH)
-      : Writer(File, Obj), WriteSectionHeaders(WSH) {}
+  ELFWriter(Object &Obj, Buffer &Buf, bool WSH)
+      : Writer(Obj, Buf), WriteSectionHeaders(WSH) {}
 };
 
 class BinaryWriter : public Writer {
@@ -183,7 +222,7 @@ public:
   ~BinaryWriter() {}
   void finalize() override;
   void write() override;
-  BinaryWriter(StringRef File, Object &Obj) : Writer(File, Obj) {}
+  BinaryWriter(Object &Obj, Buffer &Buf) : Writer(Obj, Buf) {}
 };
 
 class SectionBase {
@@ -569,14 +608,12 @@ public:
 };
 
 class ELFReader : public Reader {
-private:
-  std::unique_ptr<Binary> Bin;
-  std::shared_ptr<MemoryBuffer> Data;
+  Binary *Bin;
 
 public:
   ElfType getElfType() const;
   std::unique_ptr<Object> create() const override;
-  ELFReader(StringRef File);
+  explicit ELFReader(Binary *B) : Bin(B){};
 };
 
 class Object {