PE/COFF: rework how we handle base relocations

[compnerd]

Summary

Generalise the base relocation handling slightly to support multiple base
relocation types in PE/COFF. This is necessary to generate proper executables
for WoA.

Track the base relocation type from the decision that we need a base relocation
to the point where we emit the base relocation into base relocation directory.

[ruiu]

LGTM with these fixes.

lib/ReaderWriter/PECOFF/WriterPECOFF.cpp

769

Not directly related to this change, but can you remove this TODO? It's obsolete.

785

I'd do like this because this switch has only one non-default case.

if (ref->kindValue() == llvm::COFF::IMAGE_REL_I386_DIR32)
  reloc = std::make_pair(...);
break;

795

ditto

799

In other switches default is after all cases. (But my personal preference is if () {...} else if () {...} break.)

[compnerd]

Thanks; committed as SVN r226335 with the requested changes.

lib/ReaderWriter/PECOFF/WriterPECOFF.cpp

private:
   uint64_t _virtualAddress;
   uint32_t _stringTableOffset;
   uint64_t _memAlign;
 };
 
 /// An AtomChunk represents a section containing atoms.
 class AtomChunk : public SectionChunk {
 public:
+  typedef std::pair<uint64_t, llvm::COFF::BaseRelocationType> BaseRelocation;
+  typedef std::vector<BaseRelocation> BaseRelocationList;
+
   AtomChunk(const PECOFFLinkingContext &ctx, StringRef name,
             const std::vector<const DefinedAtom *> &atoms);
 
   void write(uint8_t *buffer) override;
 
   uint64_t memAlign() const override;
   void appendAtom(const DefinedAtom *atom);
   void buildAtomRvaMap(std::map<const Atom *, uint64_t> &atomRva) const;
 
   void applyRelocationsARM(uint8_t *buffer,
                            std::map<const Atom *, uint64_t> &atomRva,
                            std::vector<uint64_t> &sectionRva,
                            uint64_t imageBaseAddress);
   void applyRelocationsX86(uint8_t *buffer,
                            std::map<const Atom *, uint64_t> &atomRva,
                            std::vector<uint64_t> &sectionRva,
                            uint64_t imageBaseAddress);
   void applyRelocationsX64(uint8_t *buffer,
                            std::map<const Atom *, uint64_t> &atomRva,
                            std::vector<uint64_t> &sectionRva,
                            uint64_t imageBaseAddress);
 
   void printAtomAddresses(uint64_t baseAddr) const;
-  void addBaseRelocations(std::vector<uint64_t> &relocSites) const;
+  void addBaseRelocations(BaseRelocationList &relocSites) const;
 
   void setVirtualAddress(uint32_t rva) override;
   uint64_t getAtomVirtualAddress(StringRef name) const;
 
   static bool classof(const Chunk *c) { return c->getKind() == kindAtomChunk; }
 
 protected:
   std::vector<AtomLayout *> _atomLayouts;

 /// specified by .reloc section.
 ///
 /// The executable is almost always loaded at the preferred base address because
 /// it's loaded into an empty address space. The DLL is however an subject of
 /// load-time relocation because it may conflict with other DLLs or the
 /// executable.
 class BaseRelocChunk : public SectionChunk {
   typedef std::vector<std::unique_ptr<Chunk> > ChunkVectorT;
-  typedef std::map<uint64_t, std::vector<uint16_t> > PageOffsetT;
 
 public:
+  typedef std::pair<uint16_t, llvm::COFF::BaseRelocationType> BaseRelocation;
+  typedef std::vector<BaseRelocation> BaseRelocations;
+  typedef std::map<uint64_t, BaseRelocations> RelocationBlocks;
+
   BaseRelocChunk(ChunkVectorT &chunks, const PECOFFLinkingContext &ctx)
       : SectionChunk(kindSection, ".reloc", characteristics, ctx),
         _ctx(ctx), _contents(createContents(chunks)) {}
 
   void write(uint8_t *buffer) override {
     std::memcpy(buffer, &_contents[0], _contents.size());
   }
 
   uint64_t size() const override { return _contents.size(); }
 
 private:
   // When loaded into memory, reloc section should be readable and writable.
   static const uint32_t characteristics =
       llvm::COFF::IMAGE_SCN_MEM_READ |
       llvm::COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
       llvm::COFF::IMAGE_SCN_MEM_DISCARDABLE;
 
   std::vector<uint8_t> createContents(ChunkVectorT &chunks) const;
 
   // Returns a list of RVAs that needs to be relocated if the binary is loaded
   // at an address different from its preferred one.
-  std::vector<uint64_t> listRelocSites(ChunkVectorT &chunks) const;
+  AtomChunk::BaseRelocationList listRelocSites(ChunkVectorT &chunks) const;
 
   // Divide the given RVAs into blocks.
-  PageOffsetT groupByPage(const std::vector<uint64_t> &relocSites) const;
+  RelocationBlocks
+  groupByPage(const AtomChunk::BaseRelocationList &relocSites) const;
 
   // Create the content of a relocation block.
   std::vector<uint8_t>
-  createBaseRelocBlock(uint64_t pageAddr,
-                       const std::vector<uint16_t> &offsets) const;
+  createBaseRelocBlock(uint64_t pageAddr, const BaseRelocations &relocs) const;
 
   const PECOFFLinkingContext &_ctx;
   std::vector<uint8_t> _contents;
 };
 
 template <class PEHeader>
 PEHeaderChunk<PEHeader>::PEHeaderChunk(const PECOFFLinkingContext &ctx)
     : HeaderChunk() {
llvm::dbgs() << llvm::format("0x%08llx: ", addr + baseAddr)
                  << "\n";
   }
 }
 
 /// List all virtual addresses (and not relative virtual addresses) that need
 /// to be fixed up if image base is relocated. The only relocation type that
 /// needs to be fixed is DIR32 on i386. REL32 is not (and should not be)
 /// fixed up because it's PC-relative.
-void AtomChunk::addBaseRelocations(std::vector<uint64_t> &relocSites) const {
+void AtomChunk::addBaseRelocations(BaseRelocationList &relocSites) const {
   // TODO: llvm-objdump doesn't support parsing the base relocation table, so
   // we can't really test this at the moment. As a temporary solution, we
   // should output debug messages with atom names and addresses so that we
   // can inspect relocations, and fix the tests (base-reloc.test, maybe
   // others) to use those messages.
 
-  int16_t relType = 0; /* IMAGE_REL_*_ABSOLUTE */
+  for (const auto *layout : _atomLayouts) {
+    const DefinedAtom *atom = cast<DefinedAtom>(layout->_atom);
+    for (const Reference *ref : *atom) {
+      if (ref->kindNamespace() != Reference::KindNamespace::COFF)
+        continue;
+
+      uint64_t address = layout->_virtualAddr + ref->offsetInAtom();
+      AtomChunk::BaseRelocation reloc;
       switch (_machineType) {
       default: llvm_unreachable("unsupported machine type");
       case llvm::COFF::IMAGE_FILE_MACHINE_I386:
-    relType = llvm::COFF::IMAGE_REL_I386_DIR32;
+        switch (ref->kindValue()) {
+        default: continue;
+        case llvm::COFF::IMAGE_REL_I386_DIR32:
+          reloc = std::make_pair(address, llvm::COFF::IMAGE_REL_BASED_HIGHLOW);
+          break;
+        }
         break;
       case llvm::COFF::IMAGE_FILE_MACHINE_AMD64:
-    relType = llvm::COFF::IMAGE_REL_AMD64_ADDR64;
+        switch (ref->kindValue()) {
+        default: continue;
+        case llvm::COFF::IMAGE_REL_AMD64_ADDR64:
+          reloc = std::make_pair(address, llvm::COFF::IMAGE_REL_BASED_DIR64);
+          break;
+        }
         break;
       case llvm::COFF::IMAGE_FILE_MACHINE_ARMNT:
-    relType = llvm::COFF::IMAGE_REL_ARM_ADDR32;
+        switch (ref->kindValue()) {
+        default: continue;
+        case llvm::COFF::IMAGE_REL_ARM_ADDR32:
+          reloc = std::make_pair(address, llvm::COFF::IMAGE_REL_BASED_HIGHLOW);
+          break;
+        case llvm::COFF::IMAGE_REL_ARM_MOV32T:
+          reloc =
+              std::make_pair(address, llvm::COFF::IMAGE_REL_BASED_ARM_MOV32T);
           break;
         }
-
-  for (const auto *layout : _atomLayouts) {
-    const DefinedAtom *atom = cast<DefinedAtom>(layout->_atom);
-    for (const Reference *ref : *atom)
-      if ((ref->kindNamespace() == Reference::KindNamespace::COFF) &&
-          (ref->kindValue() == relType))
-        relocSites.push_back(layout->_virtualAddr + ref->offsetInAtom());
+        break;
+      }
+      relocSites.push_back(reloc);
+    }
   }
 }
 
 void AtomChunk::setVirtualAddress(uint32_t rva) {
   SectionChunk::setVirtualAddress(rva);
   for (AtomLayout *layout : _atomLayouts)
     layout->_virtualAddr += rva;
 }

 ///
 /// By dividing 32 bit RVAs into blocks, COFF saves disk and memory space for
 /// the base relocation. A block consists of a 32 bit page RVA and 16 bit
 /// relocation entries which represent offsets in the page. That is a more
 /// compact representation than a simple vector of 32 bit RVAs.
 std::vector<uint8_t>
 BaseRelocChunk::createContents(ChunkVectorT &chunks) const {
   std::vector<uint8_t> contents;
-  std::vector<uint64_t> relocSites = listRelocSites(chunks);
-  PageOffsetT blocks = groupByPage(relocSites);
+  AtomChunk::BaseRelocationList relocSites = listRelocSites(chunks);
+  RelocationBlocks blocks = groupByPage(relocSites);
   for (auto &i : blocks) {
     uint64_t pageAddr = i.first;
-    const std::vector<uint16_t> &offsetsInPage = i.second;
-    std::vector<uint8_t> block = createBaseRelocBlock(pageAddr, offsetsInPage);
+    const BaseRelocations &relocs = i.second;
+    std::vector<uint8_t> block = createBaseRelocBlock(pageAddr, relocs);
     contents.insert(contents.end(), block.begin(), block.end());
   }
   return contents;
 }
 
 // Returns a list of RVAs that needs to be relocated if the binary is loaded
 // at an address different from its preferred one.
-std::vector<uint64_t>
+AtomChunk::BaseRelocationList
 BaseRelocChunk::listRelocSites(ChunkVectorT &chunks) const {
-  std::vector<uint64_t> ret;
+  AtomChunk::BaseRelocationList ret;
   for (auto &cp : chunks)
     if (AtomChunk *chunk = dyn_cast<AtomChunk>(&*cp))
       chunk->addBaseRelocations(ret);
   return ret;
 }
 
 // Divide the given RVAs into blocks.
-BaseRelocChunk::PageOffsetT
-BaseRelocChunk::groupByPage(const std::vector<uint64_t> &relocSites) const {
-  PageOffsetT blocks;
+BaseRelocChunk::RelocationBlocks BaseRelocChunk::groupByPage(
+    const AtomChunk::BaseRelocationList &relocSites) const {
+  RelocationBlocks blocks;
   uint64_t mask = _ctx.getPageSize() - 1;
-  for (uint64_t addr : relocSites)
-    blocks[addr & ~mask].push_back(addr & mask);
+  for (const auto &reloc : relocSites)
+    blocks[reloc.first & ~mask].push_back(
+        std::make_pair(reloc.first & mask, reloc.second));
   return blocks;
 }
 
 // Create the content of a relocation block.
-std::vector<uint8_t> BaseRelocChunk::createBaseRelocBlock(
-    uint64_t pageAddr, const std::vector<uint16_t> &offsets) const {
+std::vector<uint8_t>
+BaseRelocChunk::createBaseRelocBlock(uint64_t pageAddr,
+                                     const BaseRelocations &relocs) const {
   // Relocation blocks should be padded with IMAGE_REL_I386_ABSOLUTE to be
   // aligned to a DWORD size boundary.
   uint32_t size = llvm::RoundUpToAlignment(
-      sizeof(ulittle32_t) * 2 + sizeof(ulittle16_t) * offsets.size(),
+      sizeof(ulittle32_t) * 2 + sizeof(ulittle16_t) * relocs.size(),
       sizeof(ulittle32_t));
   std::vector<uint8_t> contents(size);
   uint8_t *ptr = &contents[0];
 
   // The first four bytes is the page RVA.
   *reinterpret_cast<ulittle32_t *>(ptr) = pageAddr;
   ptr += sizeof(ulittle32_t);
 
   // The second four bytes is the size of the block, including the the page
   // RVA and this size field.
   *reinterpret_cast<ulittle32_t *>(ptr) = size;
   ptr += sizeof(ulittle32_t);
 
-  // The rest of the block consists of offsets in the page.
-  uint16_t type = _ctx.is64Bit() ? llvm::COFF::IMAGE_REL_BASED_DIR64
-                                 : llvm::COFF::IMAGE_REL_BASED_HIGHLOW;
-  for (uint16_t offset : offsets) {
-    assert(offset < _ctx.getPageSize());
-    *reinterpret_cast<ulittle16_t *>(ptr) = (type << 12) | offset;
+  for (const auto &reloc : relocs) {
+    assert(reloc.first < _ctx.getPageSize());
+    *reinterpret_cast<ulittle16_t *>(ptr) = (reloc.second << 12) | reloc.first;
     ptr += sizeof(ulittle16_t);
   }
   return contents;
 }
 
 } // end anonymous namespace
 
 class PECOFFWriter : public Writer {