5 files changed, 110 insertions, 75 deletions

diff --git a/lld/lib/ReaderWriter/ELF/DefaultTargetHandler.h b/lld/lib/ReaderWriter/ELF/DefaultTargetHandler.h
index b2cb5d5e001..0852bc854b2 100644
--- a/lld/lib/ReaderWriter/ELF/DefaultTargetHandler.h
+++ b/lld/lib/ReaderWriter/ELF/DefaultTargetHandler.h
@@ -13,6 +13,8 @@
 #include "DefaultLayout.h"
 #include "TargetHandler.h"
 #include "ELFReader.h"
+#include "DynamicLibraryWriter.h"
+#include "ExecutableWriter.h"
 
 #include "lld/ReaderWriter/ELFLinkingContext.h"
 
@@ -78,10 +80,64 @@ public:
     return std::unique_ptr<Reader>(new ELFDSOReader(useShlibUndefines));
   }
 
+  virtual std::unique_ptr<Writer> getWriter();
+
 protected:
   llvm::BumpPtrAllocator _alloc;
 };
+
+template <class ELFT>
+std::unique_ptr<Writer> DefaultTargetHandler<ELFT>::getWriter() {
+  switch (this->_context.getOutputELFType()) {
+  case llvm::ELF::ET_EXEC:
+    if (this->_context.is64Bits()) {
+      if (this->_context.isLittleEndian())
+        return std::unique_ptr<Writer>(
+            new elf::ExecutableWriter<ELFType<support::little, 8, true>>(
+                this->_context));
+      else
+        return std::unique_ptr<Writer>(
+            new elf::ExecutableWriter<ELFType<support::big, 8, true>>(
+                this->_context));
+    } else {
+      if (this->_context.isLittleEndian())
+        return std::unique_ptr<Writer>(
+            new elf::ExecutableWriter<ELFType<support::little, 4, false>>(
+                this->_context));
+      else
+        return std::unique_ptr<Writer>(
+            new elf::ExecutableWriter<ELFType<support::big, 4, false>>(
+                this->_context));
+    }
+    break;
+  case llvm::ELF::ET_DYN:
+    if (this->_context.is64Bits()) {
+      if (this->_context.isLittleEndian())
+        return std::unique_ptr<Writer>(
+            new elf::DynamicLibraryWriter<ELFType<support::little, 8, true>>(
+                this->_context));
+      else
+        return std::unique_ptr<Writer>(
+            new elf::DynamicLibraryWriter<ELFType<support::big, 8, true>>(
+                this->_context));
+    } else {
+      if (this->_context.isLittleEndian())
+        return std::unique_ptr<Writer>(
+            new elf::DynamicLibraryWriter<ELFType<support::little, 4, false>>(
+                this->_context));
+      else
+        return std::unique_ptr<Writer>(
+            new elf::DynamicLibraryWriter<ELFType<support::big, 4, false>>(
+                this->_context));
+    }
+    break;
+  case llvm::ELF::ET_REL:
+    llvm_unreachable("TODO: support -r mode");
+  default:
+    llvm_unreachable("unsupported output type");
+  }
+}
+
 } // end namespace elf
 } // end namespace lld
-
 #endif
diff --git a/lld/lib/ReaderWriter/ELF/ELFLinkingContext.cpp b/lld/lib/ReaderWriter/ELF/ELFLinkingContext.cpp
index d2a4c4e3102..c3b68909311 100644
--- a/lld/lib/ReaderWriter/ELF/ELFLinkingContext.cpp
+++ b/lld/lib/ReaderWriter/ELF/ELFLinkingContext.cpp
@@ -90,7 +90,7 @@ bool ELFLinkingContext::validateImpl(raw_ostream &diagnostics) {
     llvm_unreachable("Unimplemented");
     break;
   default:
-    _writer = createWriterELF(*this);
+    _writer = createWriterELF(this->targetHandler());
     break;
   }
   return true;
diff --git a/lld/lib/ReaderWriter/ELF/OutputELFWriter.h b/lld/lib/ReaderWriter/ELF/OutputELFWriter.h
index 7702e814260..af19214f2e6 100644
--- a/lld/lib/ReaderWriter/ELF/OutputELFWriter.h
+++ b/lld/lib/ReaderWriter/ELF/OutputELFWriter.h
@@ -53,9 +53,18 @@ protected:
   // Build the output file
   virtual error_code buildOutput(const File &file);
 
+  // Setup the ELF header.
+  virtual error_code setELFHeader();
+
   // Write the file to the path specified
   virtual error_code writeFile(const File &File, StringRef path);
 
+  // Write to the output file.
+  virtual error_code writeOutput(const File &file, StringRef path);
+
+  // Get the size of the output file that the linker would emit.
+  virtual uint64_t outputFileSize() const;
+
   // Build the atom to address map, this has to be called
   // before applying relocations
   virtual void buildAtomToAddressMap(const File &file);
@@ -365,22 +374,7 @@ error_code OutputELFWriter<ELFT>::buildOutput(const File &file) {
   return error_code::success();
 }
 
-template <class ELFT>
-error_code OutputELFWriter<ELFT>::writeFile(const File &file, StringRef path) {
-  buildOutput(file);
-
-  uint64_t totalSize = _shdrtab->fileOffset() + _shdrtab->fileSize();
-
-  OwningPtr<FileOutputBuffer> buffer;
-  ScopedTask createOutputTask(getDefaultDomain(), "ELF Writer Create Output");
-  error_code ec = FileOutputBuffer::create(path,
-                                           totalSize, buffer,
-                                           FileOutputBuffer::F_executable);
-  createOutputTask.end();
-
-  if (ec)
-    return ec;
-
+template <class ELFT> error_code OutputELFWriter<ELFT>::setELFHeader() {
   _elfHeader->e_ident(ELF::EI_CLASS,
                       _context.is64Bits() ? ELF::ELFCLASS64 : ELF::ELFCLASS32);
   _elfHeader->e_ident(ELF::EI_DATA, _context.isLittleEndian()
@@ -388,15 +382,9 @@ error_code OutputELFWriter<ELFT>::writeFile(const File &file, StringRef path) {
                                         : ELF::ELFDATA2MSB);
   _elfHeader->e_type(_context.getOutputELFType());
   _elfHeader->e_machine(_context.getOutputMachine());
-
-  if (!_targetHandler.doesOverrideELFHeader()) {
-    _elfHeader->e_ident(ELF::EI_VERSION, 1);
-    _elfHeader->e_ident(ELF::EI_OSABI, 0);
-    _elfHeader->e_version(1);
-  } else {
-    // override the contents of the ELF Header
-    _targetHandler.setELFHeader(_elfHeader.get());
-  }
+  _elfHeader->e_ident(ELF::EI_VERSION, 1);
+  _elfHeader->e_ident(ELF::EI_OSABI, 0);
+  _elfHeader->e_version(1);
   _elfHeader->e_phoff(_programHeader->fileOffset());
   _elfHeader->e_shoff(_shdrtab->fileOffset());
   _elfHeader->e_phentsize(_programHeader->entsize());
@@ -408,20 +396,53 @@ error_code OutputELFWriter<ELFT>::writeFile(const File &file, StringRef path) {
   _layout->findAtomAddrByName(_context.entrySymbolName(), virtualAddr);
   _elfHeader->e_entry(virtualAddr);
 
+  return error_code::success();
+}
+
+template <class ELFT> uint64_t OutputELFWriter<ELFT>::outputFileSize() const {
+  return _shdrtab->fileOffset() + _shdrtab->fileSize();
+}
+
+template <class ELFT>
+error_code OutputELFWriter<ELFT>::writeOutput(const File &file,
+                                              StringRef path) {
+  OwningPtr<FileOutputBuffer> buffer;
+  ScopedTask createOutputTask(getDefaultDomain(), "ELF Writer Create Output");
+  error_code ec = FileOutputBuffer::create(path, outputFileSize(), buffer,
+                                           FileOutputBuffer::F_executable);
+  createOutputTask.end();
+
+  if (ec)
+    return ec;
+
+  ScopedTask writeTask(getDefaultDomain(), "ELF Writer write to memory");
+
   // HACK: We have to write out the header and program header here even though
   // they are a member of a segment because only sections are written in the
   // following loop.
-  ScopedTask writeTask(getDefaultDomain(), "ELF Writer write to memory");
   _elfHeader->write(this, *buffer);
   _programHeader->write(this, *buffer);
 
   for (auto section : _layout->sections())
-      section->write(this, *buffer);
+    section->write(this, *buffer);
   writeTask.end();
 
   ScopedTask commitTask(getDefaultDomain(), "ELF Writer commit to disk");
   return buffer->commit();
 }
+
+template <class ELFT>
+error_code OutputELFWriter<ELFT>::writeFile(const File &file, StringRef path) {
+  error_code ec = buildOutput(file);
+  if (ec)
+    return ec;
+
+  ec = setELFHeader();
+  if (ec)
+    return ec;
+
+  return writeOutput(file, path);
+}
 } // namespace elf
 } // namespace lld
 
diff --git a/lld/lib/ReaderWriter/ELF/TargetHandler.h b/lld/lib/ReaderWriter/ELF/TargetHandler.h
index eb5c0de7091..1dada951cb7 100644
--- a/lld/lib/ReaderWriter/ELF/TargetHandler.h
+++ b/lld/lib/ReaderWriter/ELF/TargetHandler.h
@@ -96,6 +96,8 @@ public:
 
   virtual std::unique_ptr<Reader> getDSOReader(bool) = 0;
 
+  virtual std::unique_ptr<Writer> getWriter() = 0;
+
 protected:
   ELFLinkingContext &_context;
 };
diff --git a/lld/lib/ReaderWriter/ELF/Writer.cpp b/lld/lib/ReaderWriter/ELF/Writer.cpp
index 43a6194ee35..94fa1b916ab 100644
--- a/lld/lib/ReaderWriter/ELF/Writer.cpp
+++ b/lld/lib/ReaderWriter/ELF/Writer.cpp
@@ -17,52 +17,8 @@ using namespace llvm;
 using namespace llvm::object;
 namespace lld {
 
-std::unique_ptr<Writer> createWriterELF(const ELFLinkingContext &info) {
-  using llvm::object::ELFType;
-  // Set the default layout to be the static executable layout
-  // We would set the layout to a dynamic executable layout
-  // if we came across any shared libraries in the process
-
-  switch (info.getOutputELFType()) {
-  case llvm::ELF::ET_EXEC:
-    if (info.is64Bits()) {
-      if (info.isLittleEndian())
-        return std::unique_ptr<Writer>(new
-            elf::ExecutableWriter<ELFType<support::little, 8, true>>(info));
-      else
-        return std::unique_ptr<Writer>(new
-                elf::ExecutableWriter<ELFType<support::big, 8, true>>(info));
-    } else {
-      if (info.isLittleEndian())
-        return std::unique_ptr<Writer>(new
-            elf::ExecutableWriter<ELFType<support::little, 4, false>>(info));
-      else
-        return std::unique_ptr<Writer>(new
-                elf::ExecutableWriter<ELFType<support::big, 4, false>>(info));
-    }
-  break;
-  case llvm::ELF::ET_DYN:
-    if (info.is64Bits()) {
-      if (info.isLittleEndian())
-        return std::unique_ptr<Writer>(new
-          elf::DynamicLibraryWriter<ELFType<support::little, 8, true>>(info));
-      else
-        return std::unique_ptr<Writer>(new
-              elf::DynamicLibraryWriter<ELFType<support::big, 8, true>>(info));
-    } else {
-      if (info.isLittleEndian())
-        return std::unique_ptr<Writer>(new
-          elf::DynamicLibraryWriter<ELFType<support::little, 4, false>>(info));
-      else
-        return std::unique_ptr<Writer>(new
-              elf::DynamicLibraryWriter<ELFType<support::big, 4, false>>(info));
-    }
-  break;
-  case llvm::ELF::ET_REL:
-    llvm_unreachable("TODO: support -r mode");
-  default:
-    llvm_unreachable("unsupported output type");
-  }
+std::unique_ptr<Writer> createWriterELF(TargetHandlerBase *handler) {
+  return std::move(handler->getWriter());
 }
 
 } // namespace lld