[llvm-objcopy] Refactor code to include initialize method

This change refactors some of the code to allow for some code
deduplication in later diffs as well as just to make adding a new
section type more self contained to the class itself. The idea for this
was first mentioned by James in D 37915 and will be used in that change
as recommended.

This change follows changes for dynamic sections but precedes support
for dynamic relocations.

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@314148 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 92 insertions, 34 deletions

diff --git a/tools/llvm-objcopy/Object.cpp b/tools/llvm-objcopy/Object.cpp
index 28323f47500..a8cd321086b 100644
--- a/tools/llvm-objcopy/Object.cpp
+++ b/tools/llvm-objcopy/Object.cpp
@@ -49,6 +49,7 @@ void Segment::writeSegment(FileOutputBuffer &Out) const {
   std::copy(std::begin(Contents), std::end(Contents), Buf);
 }
 
+void SectionBase::initialize(SectionTableRef SecTable) {}
 void SectionBase::finalize() {}
 
 template <class ELFT>
@@ -149,6 +150,16 @@ void SymbolTableSection::addSymbol(StringRef Name, uint8_t Bind, uint8_t Type,
   Size += this->EntrySize;
 }
 
+void SymbolTableSection::initialize(SectionTableRef SecTable) {
+  Size = 0;
+  setStrTab(SecTable.getSectionOfType<StringTableSection>(
+      Link,
+      "Symbol table has link index of " + Twine(Link) +
+          " which is not a valid index",
+      "Symbol table has link index of " + Twine(Link) +
+          " which is not a string table"));
+}
+
 void SymbolTableSection::finalize() {
   // Make sure SymbolNames is finalized before getting name indexes.
   SymbolNames->finalize();
@@ -195,6 +206,19 @@ void SymbolTableSectionImpl<ELFT>::writeSection(
   }
 }
 
+template <class ELFT>
+void RelocationSection<ELFT>::initialize(SectionTableRef SecTable) {
+  setSymTab(SecTable.getSectionOfType<SymbolTableSection>(
+      Link,
+      "Link field value " + Twine(Link) + " in section " + Name + " is invalid",
+      "Link field value " + Twine(Link) + " in section " + Name +
+          " is not a symbol table"));
+
+  setSection(SecTable.getSection(Info,
+                                 "Info field value " + Twine(Info) +
+                                     " in section " + Name + " is invalid"));
+}
+
 template <class ELFT> void RelocationSection<ELFT>::finalize() {
   this->Link = Symbols->Index;
   this->Info = SecToApplyRel->Index;
@@ -232,6 +256,14 @@ bool SectionWithStrTab::classof(const SectionBase *S) {
   return isa<DynamicSymbolTableSection>(S) || isa<DynamicSection>(S);
 }
 
+void SectionWithStrTab::initialize(SectionTableRef SecTable) {
+  setStrTab(SecTable.getSectionOfType<StringTableSection>(
+      Link,
+      "Link field value " + Twine(Link) + " in section " + Name + " is invalid",
+      "Link field value " + Twine(Link) + " in section " + Name +
+          " is not a string table"));
+}
+
 void SectionWithStrTab::finalize() { this->Link = StrTab->Index; }
 
 // Returns true IFF a section is wholly inside the range of a segment
@@ -311,15 +343,8 @@ void Object<ELFT>::readProgramHeaders(const ELFFile<ELFT> &ElfFile) {
 
 template <class ELFT>
 void Object<ELFT>::initSymbolTable(const llvm::object::ELFFile<ELFT> &ElfFile,
-                                   SymbolTableSection *SymTab) {
-
-  SymTab->Size = 0;
-  SymTab->setStrTab(getSectionOfType<StringTableSection>(
-      SymbolTable->Link,
-      "Symbol table has link index of " + Twine(SymTab->Link) +
-          " which is not a valid index",
-      "Symbol table has link index of " + Twine(SymTab->Link) +
-          " which is not a string table"));
+                                   SymbolTableSection *SymTab,
+                                   SectionTableRef SecTable) {
 
   const Elf_Shdr &Shdr = *unwrapOrError(ElfFile.getSection(SymTab->Index));
   StringRef StrTabData = unwrapOrError(ElfFile.getStringTableForSymtab(Shdr));
@@ -336,7 +361,7 @@ void Object<ELFT>::initSymbolTable(const llvm::object::ELFFile<ELFT> &ElfFile,
             Twine(Sym.st_shndx));
       }
     } else if (Sym.st_shndx != SHN_UNDEF) {
-      DefSection = getSection(
+      DefSection = SecTable.getSection(
           Sym.st_shndx,
           "Symbol '" + Name + "' is defined in invalid section with index " +
               Twine(Sym.st_shndx));
@@ -368,6 +393,20 @@ void initRelocations(RelocationSection<ELFT> *Relocs,
   }
 }
 
+SectionBase *SectionTableRef::getSection(uint16_t Index, Twine ErrMsg) {
+  if (Index == SHN_UNDEF || Index > Sections.size())
+    error(ErrMsg);
+  return Sections[Index - 1].get();
+}
+
+template <class T>
+T *SectionTableRef::getSectionOfType(uint16_t Index, Twine IndexErrMsg,
+                                  Twine TypeErrMsg) {
+  if (T *Sec = llvm::dyn_cast<T>(getSection(Index, IndexErrMsg)))
+    return Sec;
+  error(TypeErrMsg);
+}
+
 template <class ELFT>
 SectionBase *Object<ELFT>::getSection(uint16_t Index, Twine ErrMsg) {
   if (Index == SHN_UNDEF || Index > Sections.size())
@@ -428,7 +467,7 @@ Object<ELFT>::makeSection(const llvm::object::ELFFile<ELFT> &ElfFile,
 }
 
 template <class ELFT>
-void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
+SectionTableRef Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
   uint32_t Index = 0;
   for (const auto &Shdr : unwrapOrError(ElfFile.sections())) {
     if (Index == 0) {
@@ -451,35 +490,30 @@ void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
     Sections.push_back(std::move(Sec));
   }
 
+  SectionTableRef SecTable(Sections);
+
   // Now that all of the sections have been added we can fill out some extra
-  // details about symbol tables.
-  if (SymbolTable)
-    initSymbolTable(ElfFile, SymbolTable);
+  // details about symbol tables. We need the symbol table filled out before
+  // any relocations.
+  if (SymbolTable) {
+    SymbolTable->initialize(SecTable);
+    initSymbolTable(ElfFile, SymbolTable, SecTable);
+  }
 
   // Now that all sections and symbols have been added we can add
   // relocations that reference symbols and set the link and info fields for
   // relocation sections.
   for (auto &Section : Sections) {
+    if (Section.get() == SymbolTable)
+      continue;
+    Section->initialize(SecTable);
     if (auto RelSec = dyn_cast<RelocationSection<ELFT>>(Section.get())) {
-
-      auto SymTab = getSectionOfType<SymbolTableSection>(
-          RelSec->Link,
-          "Link field value " + Twine(RelSec->Link) + " in section " +
-              RelSec->Name + " is invalid",
-          "Link field value " + Twine(RelSec->Link) + " in section " +
-              RelSec->Name + " is not a symbol table");
-      RelSec->setSymTab(SymTab);
-
-      RelSec->setSection(getSection(RelSec->Info,
-                                    "Info field value " + Twine(RelSec->Link) +
-                                        " in section " + RelSec->Name +
-                                        " is invalid"));
-
       auto Shdr = unwrapOrError(ElfFile.sections()).begin() + RelSec->Index;
       if (RelSec->Type == SHT_REL)
-        initRelocations(RelSec, SymTab, unwrapOrError(ElfFile.rels(Shdr)));
+        initRelocations(RelSec, SymbolTable, unwrapOrError(ElfFile.rels(Shdr)));
       else
-        initRelocations(RelSec, SymTab, unwrapOrError(ElfFile.relas(Shdr)));
+        initRelocations(RelSec, SymbolTable,
+                        unwrapOrError(ElfFile.relas(Shdr)));
     }
 
     if (auto Sec = dyn_cast<SectionWithStrTab>(Section.get())) {
@@ -491,6 +525,8 @@ void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
               " is not a string table"));
     }
   }
+
+  return SecTable;
 }
 
 template <class ELFT> Object<ELFT>::Object(const ELFObjectFile<ELFT> &Obj) {
@@ -504,10 +540,10 @@ template <class ELFT> Object<ELFT>::Object(const ELFObjectFile<ELFT> &Obj) {
   Entry = Ehdr.e_entry;
   Flags = Ehdr.e_flags;
 
-  readSectionHeaders(ElfFile);
+  SectionTableRef SecTable = readSectionHeaders(ElfFile);
   readProgramHeaders(ElfFile);
 
-  SectionNames = getSectionOfType<StringTableSection>(
+  SectionNames = SecTable.getSectionOfType<StringTableSection>(
       Ehdr.e_shstrndx,
       "e_shstrndx field value " + Twine(Ehdr.e_shstrndx) + " in elf header " +
           " is invalid",
diff --git a/tools/llvm-objcopy/Object.h b/tools/llvm-objcopy/Object.h
index 9ee8b7377a0..6e849e47006 100644
--- a/tools/llvm-objcopy/Object.h
+++ b/tools/llvm-objcopy/Object.h
@@ -18,6 +18,24 @@
 #include <set>
 
 class Segment;
+class SectionBase;
+
+class SectionTableRef {
+private:
+  llvm::ArrayRef<std::unique_ptr<SectionBase>> Sections;
+
+public:
+  SectionTableRef(llvm::ArrayRef<std::unique_ptr<SectionBase>> Secs)
+      : Sections(Secs) {}
+  SectionTableRef(const SectionTableRef &) = default;
+
+  SectionBase *getSection(uint16_t Index, llvm::Twine ErrMsg);
+
+  template <class T>
+
+  T *getSectionOfType(uint16_t Index, llvm::Twine IndexErrMsg,
+                      llvm::Twine TypeErrMsg);
+};
 
 class SectionBase {
 public:
@@ -39,6 +57,7 @@ public:
   uint64_t Type = llvm::ELF::SHT_NULL;
 
   virtual ~SectionBase() {}
+  virtual void initialize(SectionTableRef SecTable);
   virtual void finalize();
   template <class ELFT> void writeHeader(llvm::FileOutputBuffer &Out) const;
   virtual void writeSection(llvm::FileOutputBuffer &Out) const = 0;
@@ -154,6 +173,7 @@ public:
                  uint64_t Sz);
   void addSymbolNames();
   const Symbol *getSymbolByIndex(uint32_t Index) const;
+  void initialize(SectionTableRef SecTable) override;
   void finalize() override;
   static bool classof(const SectionBase *S) {
     return S->Type == llvm::ELF::SHT_SYMTAB;
@@ -187,6 +207,7 @@ public:
   void setSymTab(SymbolTableSection *StrTab) { Symbols = StrTab; }
   void setSection(SectionBase *Sec) { SecToApplyRel = Sec; }
   void addRelocation(Relocation Rel) { Relocations.push_back(Rel); }
+  void initialize(SectionTableRef SecTable) override;
   void finalize() override;
   void writeSection(llvm::FileOutputBuffer &Out) const override;
   static bool classof(const SectionBase *S) {
@@ -201,6 +222,7 @@ private:
 public:
   SectionWithStrTab(llvm::ArrayRef<uint8_t> Data) : Section(Data) {}
   void setStrTab(StringTableSection *StringTable) { StrTab = StringTable; }
+  void initialize(SectionTableRef SecTable) override;
   void finalize() override;
   static bool classof(const SectionBase *S);
 };
@@ -232,11 +254,11 @@ private:
   typedef typename ELFT::Phdr Elf_Phdr;
 
   void initSymbolTable(const llvm::object::ELFFile<ELFT> &ElfFile,
-                       SymbolTableSection *SymTab);
+                       SymbolTableSection *SymTab, SectionTableRef SecTable);
   SecPtr makeSection(const llvm::object::ELFFile<ELFT> &ElfFile,
                      const Elf_Shdr &Shdr);
   void readProgramHeaders(const llvm::object::ELFFile<ELFT> &ElfFile);
-  void readSectionHeaders(const llvm::object::ELFFile<ELFT> &ElfFile);
+  SectionTableRef readSectionHeaders(const llvm::object::ELFFile<ELFT> &ElfFile);
 
   SectionBase *getSection(uint16_t Index, llvm::Twine ErrMsg);