[PECOFF] Change data directory atom structure.

The data directory in the PE/COFF header consisted of list of data directory
atoms. This patch changes it -- now there's only one data directory entry that
contains former data directories. That's easier to handle in the writer as well
as to write to/read from YAML/Native files. The main purpose of this refactoring
is to enable RoundTrip tests for PE/COFF.

There's no functionality change.

llvm-svn: 193854

lld/lib/ReaderWriter/PECOFF/Atoms.h

@@ -247,24 +247,11 @@ private:
 // COFF header.
 class COFFDataDirectoryAtom : public COFFLinkerInternalAtom {
 public:
-  COFFDataDirectoryAtom(const File &file, uint64_t ordinal, uint32_t entrySize,
-                        uint32_t entryAddr = 0)
-      : COFFLinkerInternalAtom(file, assembleRawContent(entrySize, entryAddr)),
-        _ordinal(ordinal) {}
+  COFFDataDirectoryAtom(const File &file, std::vector<uint8_t> contents)
+      : COFFLinkerInternalAtom(file, contents) {}
 
-  virtual uint64_t ordinal() const { return _ordinal; }
   virtual ContentType contentType() const { return typeDataDirectoryEntry; }
   virtual ContentPermissions permissions() const { return permR__; }
-
-private:
-  std::vector<uint8_t> assembleRawContent(uint32_t entrySize, uint32_t entryAddr) {
-    std::vector<uint8_t> data = std::vector<uint8_t>(8, 0);
-    *(reinterpret_cast<llvm::support::ulittle32_t *>(&data[0])) = entryAddr;
-    *(reinterpret_cast<llvm::support::ulittle32_t *>(&data[4])) = entrySize;
-    return data;
-  }
-
-  uint64_t _ordinal;
 };
 
 // A COFFSharedLibraryAtom represents a symbol for data in an import library.  A

lld/lib/ReaderWriter/PECOFF/IdataPass.h

@@ -312,23 +312,35 @@ private:
 
   /// The addresses of the import dirctory and the import address table needs to
   /// be set to the COFF Optional Data Directory header. A COFFDataDirectoryAtom
-  /// represents an entry in the data directory header. We create atoms of class
-  /// COFFDataDirectoryAtom and set relocations to them, so that the address
-  /// will be set by the writer.
+  /// represents the data directory header. We create a COFFDataDirectoryAtom
+  /// and set relocations to them, so that the address will be set by the
+  /// writer.
   void createDataDirectoryAtoms(Context &context) {
-    auto *dir = new (_alloc) coff::COFFDataDirectoryAtom(
-        context.dummyFile, llvm::COFF::DataDirectoryIndex::IMPORT_TABLE,
-        context.importDirectories.size() *
-            context.importDirectories[0]->size());
-    addDir32NBReloc(dir, context.importDirectories[0]);
-    context.file.addAtom(*dir);
+    // CLR_RUNTIME_HEADER is the last index of the data directory.
+    int nentries = llvm::COFF::CLR_RUNTIME_HEADER + 1;
+    int entSize = sizeof(llvm::object::data_directory);
+    std::vector<uint8_t> contents(nentries * entSize, 0);
 
-    auto *iat = new (_alloc) coff::COFFDataDirectoryAtom(
-        context.dummyFile, llvm::COFF::DataDirectoryIndex::IAT,
-        context.importAddressTables.size() *
-            context.importAddressTables[0]->size());
-    addDir32NBReloc(iat, context.importAddressTables[0]);
-    context.file.addAtom(*iat);
+    auto importTableOffset = llvm::COFF::DataDirectoryIndex::IMPORT_TABLE
+        * entSize;
+    auto iatOffset = llvm::COFF::DataDirectoryIndex::IAT * entSize;
+
+    auto *importTableEntry = reinterpret_cast<llvm::object::data_directory *>(
+        &contents[0] + importTableOffset);
+    auto *iatEntry = reinterpret_cast<llvm::object::data_directory *>(
+        &contents[0] + iatOffset);
+
+    importTableEntry->Size = context.importDirectories.size()
+        * context.importDirectories[0]->size();
+    iatEntry->Size = context.importAddressTables.size()
+        * context.importAddressTables[0]->size();
+
+    auto *dir = new (_alloc) coff::COFFDataDirectoryAtom(
+        context.dummyFile, std::move(contents));
+    addDir32NBReloc(dir, context.importDirectories[0], importTableOffset);
+    addDir32NBReloc(dir, context.importAddressTables[0], iatOffset);
+
+    context.file.addAtom(*dir);
   }
 
   /// Transforms a reference to a COFFSharedLibraryAtom to a real reference.

lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp

@@ -412,13 +412,12 @@ protected:
 /// in memory) and 8 byte entry data size.
 class DataDirectoryChunk : public AtomChunk {
 public:
-  DataDirectoryChunk(const File &linkedFile)
-      : AtomChunk(kindDataDirectory), _file(linkedFile) {
-    // Extract atoms from the linked file and append them to this section.
+  DataDirectoryChunk(const File &linkedFile) : AtomChunk(kindDataDirectory) {
+    // Find the data directory atom.
     for (const DefinedAtom *atom : linkedFile.defined()) {
       if (atom->contentType() == DefinedAtom::typeDataDirectoryEntry) {
-        uint64_t offset = atom->ordinal() * sizeof(llvm::object::data_directory);
-        _atomLayouts.push_back(new (_alloc) AtomLayout(atom, offset, offset));
+        _atomLayouts.push_back(new (_alloc) AtomLayout(atom, 0, 0));
+        return;
       }
     }
   }
@@ -428,24 +427,30 @@ public:
   }
 
   void setBaseRelocField(uint32_t addr, uint32_t size) {
-    auto *atom = new (_alloc) coff::COFFDataDirectoryAtom(
-        _file, llvm::COFF::DataDirectoryIndex::BASE_RELOCATION_TABLE, size,
-        addr);
-    uint64_t offset = atom->ordinal() * sizeof(llvm::object::data_directory);
-    _atomLayouts.push_back(new (_alloc) AtomLayout(atom, offset, offset));
+    _baseRelocAddr = addr;
+    _baseRelocSize = size;
   }
 
   virtual void write(uint8_t *fileBuffer) {
-    for (const AtomLayout *layout : _atomLayouts) {
-      if (!layout)
-        continue;
+    if (!_atomLayouts.empty()) {
+      assert(_atomLayouts.size() == 1);
+      const AtomLayout *layout = _atomLayouts[0];
       ArrayRef<uint8_t> content = static_cast<const DefinedAtom *>(layout->_atom)->rawContent();
-      std::memcpy(fileBuffer + layout->_fileOffset, content.data(), content.size());
+      std::memcpy(fileBuffer + _fileOffset, content.data(), content.size());
     }
+
+    // Write base relocation table entry.
+    int baseRelocOffset = llvm::COFF::DataDirectoryIndex::BASE_RELOCATION_TABLE
+        * sizeof(llvm::object::data_directory);
+    auto *baseReloc = reinterpret_cast<llvm::object::data_directory *>(
+        fileBuffer + _fileOffset + baseRelocOffset);
+    baseReloc->RelativeVirtualAddress = _baseRelocAddr;
+    baseReloc->Size = _baseRelocSize;
   }
 
 private:
-  const File &_file;
+  uint32_t _baseRelocAddr;
+  uint32_t _baseRelocSize;
   mutable llvm::BumpPtrAllocator _alloc;
 };