1 files changed, 76 insertions, 58 deletions

diff --git a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
index 8f5243cefa0..0b09184497d 100644
--- a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
+++ b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
@@ -130,42 +130,42 @@ using SectionPred = std::function<bool(const SectionBase &Sec)>;
 
 bool IsDWOSection(const SectionBase &Sec) { return Sec.Name.endswith(".dwo"); }
 
-bool OnlyKeepDWOPred(const Object &Obj, const SectionBase &Sec) {
+template <class ELFT>
+bool OnlyKeepDWOPred(const Object<ELFT> &Obj, const SectionBase &Sec) {
   // We can't remove the section header string table.
-  if (&Sec == Obj.SectionNames)
+  if (&Sec == Obj.getSectionHeaderStrTab())
     return false;
   // Short of keeping the string table we want to keep everything that is a DWO
   // section and remove everything else.
   return !IsDWOSection(Sec);
 }
 
-static ElfType OutputElfType;
-
-std::unique_ptr<Writer> CreateWriter(Object &Obj, StringRef File) {
-  if (OutputFormat == "binary") {
-    return llvm::make_unique<BinaryWriter>(OutputFilename, Obj);
-  }
-  // Depending on the initial ELFT and OutputFormat we need a different Writer.
-  switch (OutputElfType) {
-  case ELFT_ELF32LE:
-    return llvm::make_unique<ELFWriter<ELF32LE>>(File, Obj, !StripSections);
-  case ELFT_ELF64LE:
-    return llvm::make_unique<ELFWriter<ELF64LE>>(File, Obj, !StripSections);
-  case ELFT_ELF32BE:
-    return llvm::make_unique<ELFWriter<ELF32BE>>(File, Obj, !StripSections);
-  case ELFT_ELF64BE:
-    return llvm::make_unique<ELFWriter<ELF64BE>>(File, Obj, !StripSections);
-  }
-  llvm_unreachable("Invalid output format");
+template <class ELFT>
+void WriteObjectFile(const Object<ELFT> &Obj, StringRef File) {
+  std::unique_ptr<FileOutputBuffer> Buffer;
+  Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
+      FileOutputBuffer::create(File, Obj.totalSize(),
+                               FileOutputBuffer::F_executable);
+  handleAllErrors(BufferOrErr.takeError(), [](const ErrorInfoBase &) {
+    error("failed to open " + OutputFilename);
+  });
+  Buffer = std::move(*BufferOrErr);
+
+  Obj.write(*Buffer);
+  if (auto E = Buffer->commit())
+    reportError(File, errorToErrorCode(std::move(E)));
 }
 
-void SplitDWOToFile(const Reader &Reader, StringRef File) {
-  auto DWOFile = Reader.create();
-  DWOFile->removeSections(
-      [&](const SectionBase &Sec) { return OnlyKeepDWOPred(*DWOFile, Sec); });
-  auto Writer = CreateWriter(*DWOFile, File);
-  Writer->finalize();
-  Writer->write();
+template <class ELFT>
+void SplitDWOToFile(const ELFObjectFile<ELFT> &ObjFile, StringRef File) {
+  // Construct a second output file for the DWO sections.
+  ELFObject<ELFT> DWOFile(ObjFile);
+
+  DWOFile.removeSections([&](const SectionBase &Sec) {
+    return OnlyKeepDWOPred<ELFT>(DWOFile, Sec);
+  });
+  DWOFile.finalize();
+  WriteObjectFile(DWOFile, File);
 }
 
 // This function handles the high level operations of GNU objcopy including
@@ -175,16 +175,23 @@ void SplitDWOToFile(const Reader &Reader, StringRef File) {
 // any previous removals. Lastly whether or not something is removed shouldn't
 // depend a) on the order the options occur in or b) on some opaque priority
 // system. The only priority is that keeps/copies overrule removes.
-void HandleArgs(Object &Obj, const Reader &Reader) {
+template <class ELFT> void CopyBinary(const ELFObjectFile<ELFT> &ObjFile) {
+  std::unique_ptr<Object<ELFT>> Obj;
 
-  if (!SplitDWO.empty()) {
-    SplitDWOToFile(Reader, SplitDWO);
-  }
+  if (!OutputFormat.empty() && OutputFormat != "binary")
+    error("invalid output format '" + OutputFormat + "'");
+  if (!OutputFormat.empty() && OutputFormat == "binary")
+    Obj = llvm::make_unique<BinaryObject<ELFT>>(ObjFile);
+  else
+    Obj = llvm::make_unique<ELFObject<ELFT>>(ObjFile);
+
+  if (!SplitDWO.empty())
+    SplitDWOToFile<ELFT>(ObjFile, SplitDWO.getValue());
 
   // Localize:
 
   if (LocalizeHidden) {
-    Obj.SymbolTable->localize([](const Symbol &Sym) {
+    Obj->getSymTab()->localize([](const Symbol &Sym) {
       return Sym.Visibility == STV_HIDDEN || Sym.Visibility == STV_INTERNAL;
     });
   }
@@ -207,7 +214,7 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
 
   if (ExtractDWO)
     RemovePred = [RemovePred, &Obj](const SectionBase &Sec) {
-      return OnlyKeepDWOPred(Obj, Sec) || RemovePred(Sec);
+      return OnlyKeepDWOPred(*Obj, Sec) || RemovePred(Sec);
     };
 
   if (StripAllGNU)
@@ -216,7 +223,7 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
         return true;
       if ((Sec.Flags & SHF_ALLOC) != 0)
         return false;
-      if (&Sec == Obj.SectionNames)
+      if (&Sec == Obj->getSectionHeaderStrTab())
         return false;
       switch (Sec.Type) {
       case SHT_SYMTAB:
@@ -232,6 +239,7 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
     RemovePred = [RemovePred](const SectionBase &Sec) {
       return RemovePred(Sec) || (Sec.Flags & SHF_ALLOC) == 0;
     };
+    Obj->WriteSectionHeaders = false;
   }
 
   if (StripDebug) {
@@ -244,7 +252,7 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
     RemovePred = [RemovePred, &Obj](const SectionBase &Sec) {
       if (RemovePred(Sec))
         return true;
-      if (&Sec == Obj.SectionNames)
+      if (&Sec == Obj->getSectionHeaderStrTab())
         return false;
       return (Sec.Flags & SHF_ALLOC) == 0;
     };
@@ -253,7 +261,7 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
     RemovePred = [RemovePred, &Obj](const SectionBase &Sec) {
       if (RemovePred(Sec))
         return true;
-      if (&Sec == Obj.SectionNames)
+      if (&Sec == Obj->getSectionHeaderStrTab())
         return false;
       if (Sec.Name.startswith(".gnu.warning"))
         return false;
@@ -270,15 +278,17 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
         return false;
 
       // Allow all implicit removes.
-      if (RemovePred(Sec))
+      if (RemovePred(Sec)) {
         return true;
+      }
 
       // Keep special sections.
-      if (Obj.SectionNames == &Sec)
+      if (Obj->getSectionHeaderStrTab() == &Sec) {
         return false;
-      if (Obj.SymbolTable == &Sec || Obj.SymbolTable->getStrTab() == &Sec)
+      }
+      if (Obj->getSymTab() == &Sec || Obj->getSymTab()->getStrTab() == &Sec) {
         return false;
-
+      }
       // Remove everything else.
       return true;
     };
@@ -295,7 +305,7 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
     };
   }
 
-  Obj.removeSections(RemovePred);
+  Obj->removeSections(RemovePred);
 
   if (!AddSection.empty()) {
     for (const auto &Flag : AddSection) {
@@ -308,22 +318,16 @@ void HandleArgs(Object &Obj, const Reader &Reader) {
       auto Buf = std::move(*BufOrErr);
       auto BufPtr = reinterpret_cast<const uint8_t *>(Buf->getBufferStart());
       auto BufSize = Buf->getBufferSize();
-      Obj.addSection<OwnedDataSection>(SecName,
-                                       ArrayRef<uint8_t>(BufPtr, BufSize));
+      Obj->addSection(SecName, ArrayRef<uint8_t>(BufPtr, BufSize));
     }
   }
 
   if (!AddGnuDebugLink.empty()) {
-    Obj.addSection<GnuDebugLinkSection>(StringRef(AddGnuDebugLink));
+    Obj->addGnuDebugLink(AddGnuDebugLink);
   }
-}
 
-std::unique_ptr<Reader> CreateReader() {
-  // Right now we can only read ELF files so there's only one reader;
-  auto Out = llvm::make_unique<ELFReader>(StringRef(InputFilename));
-  // We need to set the default ElfType for output.
-  OutputElfType = Out->getElfType();
-  return std::move(Out);
+  Obj->finalize();
+  WriteObjectFile(*Obj, OutputFilename.getValue());
 }
 
 int main(int argc, char **argv) {
@@ -337,11 +341,25 @@ int main(int argc, char **argv) {
     cl::PrintHelpMessage();
     return 2;
   }
-
-  auto Reader = CreateReader();
-  auto Obj = Reader->create();
-  auto Writer = CreateWriter(*Obj, OutputFilename);
-  HandleArgs(*Obj, *Reader);
-  Writer->finalize();
-  Writer->write();
+  Expected<OwningBinary<Binary>> BinaryOrErr = createBinary(InputFilename);
+  if (!BinaryOrErr)
+    reportError(InputFilename, BinaryOrErr.takeError());
+  Binary &Binary = *BinaryOrErr.get().getBinary();
+  if (auto *o = dyn_cast<ELFObjectFile<ELF64LE>>(&Binary)) {
+    CopyBinary(*o);
+    return 0;
+  }
+  if (auto *o = dyn_cast<ELFObjectFile<ELF32LE>>(&Binary)) {
+    CopyBinary(*o);
+    return 0;
+  }
+  if (auto *o = dyn_cast<ELFObjectFile<ELF64BE>>(&Binary)) {
+    CopyBinary(*o);
+    return 0;
+  }
+  if (auto *o = dyn_cast<ELFObjectFile<ELF32BE>>(&Binary)) {
+    CopyBinary(*o);
+    return 0;
+  }
+  reportError(InputFilename, object_error::invalid_file_type);
 }