[PECOFF] Add a pass to ensure the correct ordering of grouped sections.

Reviewers: shankarke

CC: llvm-commits

Differential Revision: http://llvm-reviews.chandlerc.com/D998

llvm-svn: 184327

lld/lib/ReaderWriter/PECOFF/Atoms.cpp

@@ -0,0 +1,30 @@
+//===- lib/ReaderWriter/PECOFF/Atoms.cpp ----------------------------------===//
+//
+//                             The LLVM Linker
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Atoms.h"
+
+namespace lld {
+namespace coff {
+
+namespace {
+void addEdge(COFFDefinedAtom *a, COFFDefinedAtom *b,
+             lld::Reference::Kind kind) {
+  auto ref = new COFFReference(kind);
+  ref->setTarget(b);
+  a->addReference(std::unique_ptr<COFFReference>(ref));
+}
+}
+
+void connectAtomsWithLayoutEdge(COFFDefinedAtom *a, COFFDefinedAtom *b) {
+  addEdge(a, b, lld::Reference::kindLayoutAfter);
+  addEdge(b, a, lld::Reference::kindLayoutBefore);
+}
+
+} // namespace coff
+} // namespace lld

lld/lib/ReaderWriter/PECOFF/Atoms.h

@@ -15,10 +15,14 @@
 
 namespace lld {
 namespace coff {
+class COFFDefinedAtom;
 
+using llvm::object::COFFObjectFile;
 using llvm::object::coff_section;
 using llvm::object::coff_symbol;
 
+void connectAtomsWithLayoutEdge(COFFDefinedAtom *a, COFFDefinedAtom *b);
+
 /// A COFFReference represents relocation information for an atom. For
 /// example, if atom X has a reference to atom Y with offsetInAtom=8, that
 /// means that the address starting at 8th byte of the content of atom X needs
@@ -93,23 +97,23 @@ private:
 class COFFDefinedAtom : public DefinedAtom {
 public:
   COFFDefinedAtom(const File &f, llvm::StringRef n, const coff_symbol *symb,
-                  const coff_section *sec, llvm::ArrayRef<uint8_t> d)
-      : _owningFile(f), _name(n), _symbol(symb), _section(sec), _data(d) {}
+                  const coff_section *sec, llvm::ArrayRef<uint8_t> d,
+                  StringRef sectionName, uint64_t ordinal)
+      : _owningFile(f), _name(n), _symbol(symb), _section(sec), _data(d),
+        _sectionName(sectionName), _ordinal(ordinal) {}
 
   virtual const class File &file() const { return _owningFile; }
 
   virtual llvm::StringRef name() const { return _name; }
 
-  virtual uint64_t ordinal() const {
-    return reinterpret_cast<intptr_t>(_symbol);
-  }
+  virtual uint64_t ordinal() const { return _ordinal; }
 
   virtual uint64_t size() const { return _data.size(); }
 
   uint64_t originalOffset() const { return _symbol->Value; }
 
-  void addReference(COFFReference *reference) {
-    _references.push_back(reference);
+  void addReference(std::unique_ptr<COFFReference> reference) {
+    _references.push_back(std::move(reference));
   }
 
   virtual Scope scope() const {
@@ -163,6 +167,8 @@ public:
 
   virtual bool isAlias() const { return false; }
 
+  virtual StringRef getSectionName() const { return _sectionName; }
+
   virtual llvm::ArrayRef<uint8_t> rawContent() const { return _data; }
 
   virtual reference_iterator begin() const {
@@ -177,7 +183,7 @@ public:
 private:
   virtual const Reference *derefIterator(const void *iter) const {
     size_t index = reinterpret_cast<size_t>(iter);
-    return _references[index];
+    return _references[index].get();
   }
 
   virtual void incrementIterator(const void *&iter) const {
@@ -189,8 +195,10 @@ private:
   llvm::StringRef _name;
   const coff_symbol *_symbol;
   const coff_section *_section;
-  std::vector<COFFReference *> _references;
+  std::vector<std::unique_ptr<COFFReference> > _references;
   llvm::ArrayRef<uint8_t> _data;
+  StringRef _sectionName;
+  uint64_t _ordinal;
 };
 
 } // namespace coff

lld/lib/ReaderWriter/PECOFF/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_lld_library(lldPECOFF
+  Atoms.cpp
   PECOFFTargetInfo.cpp
   ReaderCOFF.cpp
   ReaderImportHeader.cpp

lld/lib/ReaderWriter/PECOFF/GroupedSectionsPass.h

@@ -0,0 +1,122 @@
+//===- lib/ReaderWriter/PECOFF/GroupedSectionsPass.h-----------------------===//
+//
+//                             The LLVM Linker
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file \brief This pass adds layout-{before,after} references to atoms in
+/// grouped sections, so that they will appear in the correct order in the
+/// output.
+///
+/// In COFF, sections will be merged into one section by the linker if their
+/// names are the same after discarding the "$" character and all characters
+/// follow it from their names. The characters following the "$" character
+/// determines the merge order. Assume there's an object file containing four
+/// data sections in the following order.
+///
+///   - .data$2
+///   - .data$3
+///   - .data$1
+///   - .data
+///
+/// In this case, the resulting binary should have ".data" section with the
+/// contents of ".data", ".data$1", ".data$2" and ".data$3" in that order.
+///
+/// In lld, this contraint is modeled by the atom model using
+/// layout-{before,after} references. Atoms in the same (unmerged-)section have
+/// already been connected with layout-{before,after} edges each other when the
+/// control reaches to this pass. We pick the head atom from each section that
+/// needs to merged, and connects them with layout-{before,after} edges in the
+/// right order, so that they'll be sorted correctly in the layout pass.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLD_READER_WRITER_PE_COFF_GROUPED_SECTIONS_PASS_H_
+#define LLD_READER_WRITER_PE_COFF_GROUPED_SECTIONS_PASS_H_
+
+#include "Atoms.h"
+#include "lld/Core/Pass.h"
+#include "lld/Core/Pass.h"
+#include "llvm/Support/Debug.h"
+
+#include <algorithm>
+#include <map>
+
+using lld::coff::COFFDefinedAtom;
+
+namespace lld {
+namespace pecoff {
+
+namespace {
+bool compareByFileOrdinal(const DefinedAtom *a, const DefinedAtom *b) {
+  return a->file().ordinal() > b->file().ordinal();
+}
+}
+
+class GroupedSectionsPass : public lld::Pass {
+public:
+  GroupedSectionsPass() {}
+
+  virtual void perform(MutableFile &mergedFile) {
+    std::map<StringRef, std::vector<COFFDefinedAtom *>> sectionToHeadAtoms(
+        filterHeadAtoms(mergedFile));
+    std::vector<std::vector<COFFDefinedAtom *>> groupedAtomsList(
+        groupBySectionName(sectionToHeadAtoms));
+    for (auto &groupedAtoms : groupedAtomsList)
+      connectAtoms(groupedAtoms);
+  }
+
+private:
+  typedef std::map<StringRef, std::vector<COFFDefinedAtom *>> SectionToAtomsT;
+
+  /// Returns the list of atoms that appeared at the beginning of sections.
+  SectionToAtomsT filterHeadAtoms(MutableFile &mutableFile) const {
+    SectionToAtomsT result;
+    for (const DefinedAtom *atom : mutableFile.defined()) {
+      auto *coffAtom = (COFFDefinedAtom *)atom;
+      if (coffAtom->ordinal() == 0)
+        result[coffAtom->getSectionName()].push_back(coffAtom);
+    }
+    return std::move(result);
+  }
+
+  /// Group atoms that needs to be merged. Returned atoms are sorted by section
+  /// name and file ordinal.
+  std::vector<std::vector<COFFDefinedAtom *>>
+  groupBySectionName(SectionToAtomsT sectionToHeadAtoms) const {
+    SectionToAtomsT res;
+    // Note that the atoms are already sorted by section name because std::map
+    // is a sorted map.
+    for (auto &i : sectionToHeadAtoms) {
+      StringRef sectionName = i.first;
+      std::vector<COFFDefinedAtom*> &atoms = i.second;
+      // Sections with the same name could exist if they are from different
+      // files. If that's the case, the sections needs to be sorted in the same
+      // order as they appeared in the command line.
+      std::stable_sort(atoms.begin(), atoms.end(), compareByFileOrdinal);
+      for (COFFDefinedAtom *atom : atoms) {
+        StringRef baseName = sectionName.split('$').first;
+        res[baseName].push_back(atom);
+      }
+    }
+    std::vector<std::vector<COFFDefinedAtom *>> vec;
+    for (auto &i : res)
+      vec.push_back(std::move(i.second));
+    return std::move(vec);
+  }
+
+  void connectAtoms(std::vector<COFFDefinedAtom *> atoms) const {
+    if (atoms.size() < 2)
+      return;
+    for (auto it = atoms.begin(), e = atoms.end(); it + 1 != e; ++it)
+      connectAtomsWithLayoutEdge(*it, *(it + 1));
+  }
+};
+
+} // namespace pecoff
+} // namespace lld
+
+#endif

lld/lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp

@@ -7,15 +7,14 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "lld/ReaderWriter/PECOFFTargetInfo.h"
+#include "GroupedSectionsPass.h"
 
 #include "lld/Core/PassManager.h"
 #include "lld/Passes/LayoutPass.h"
+#include "lld/ReaderWriter/PECOFFTargetInfo.h"
 #include "lld/ReaderWriter/Reader.h"
 #include "lld/ReaderWriter/Writer.h"
 
-#include "llvm/Support/Debug.h"
-
 namespace lld {
 
 error_code PECOFFTargetInfo::parseFile(
@@ -59,6 +58,7 @@ PECOFFTargetInfo::stringFromRelocKind(Reference::Kind kind) const {
 }
 
 void PECOFFTargetInfo::addPasses(PassManager &pm) const {
+  pm.add(std::unique_ptr<Pass>(new pecoff::GroupedSectionsPass()));
   pm.add(std::unique_ptr<Pass>(new LayoutPass()));
 }
 

lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp

@@ -184,14 +184,19 @@ private:
     }));
 
     llvm::ArrayRef<uint8_t> SecData;
+    StringRef sectionName;
     if (error_code ec = Obj->getSectionContents(section, SecData))
       return ec;
+    if (error_code ec = Obj->getSectionName(section, sectionName))
+      return ec;
+    uint64_t ordinal = 0;
 
     // Create an atom for the entire section.
     if (symbols.empty()) {
       llvm::ArrayRef<uint8_t> Data(SecData.data(), SecData.size());
       atoms.push_back(new (AtomStorage.Allocate<COFFDefinedAtom>())
-                      COFFDefinedAtom(*this, "", nullptr, section, Data));
+                      COFFDefinedAtom(*this, "", nullptr, section, Data,
+                                      sectionName, ordinal++));
       return error_code::success();
     }
 
@@ -201,7 +206,8 @@ private:
       uint64_t Size = symbols[0]->Value;
       llvm::ArrayRef<uint8_t> Data(SecData.data(), Size);
       atoms.push_back(new (AtomStorage.Allocate<COFFDefinedAtom>())
-                      COFFDefinedAtom(*this, "", nullptr, section, Data));
+                      COFFDefinedAtom(*this, "", nullptr, section, Data,
+                                      sectionName, ordinal++));
     }
 
     for (auto si = symbols.begin(), se = symbols.end(); si != se; ++si) {
@@ -215,24 +221,12 @@ private:
       if (error_code ec = Obj->getSymbolName(*si, name))
         return ec;
       atoms.push_back(new (AtomStorage.Allocate<COFFDefinedAtom>())
-                      COFFDefinedAtom(*this, name, *si, section, Data));
+                      COFFDefinedAtom(*this, name, *si, section, Data,
+                                      sectionName, ordinal++));
     }
     return error_code::success();
   }
 
-  void addEdge(COFFDefinedAtom *a, COFFDefinedAtom *b,
-               lld::Reference::Kind kind) const {
-    auto ref = new (AtomStorage.Allocate<COFFReference>()) COFFReference(kind);
-    ref->setTarget(b);
-    a->addReference(ref);
-  }
-
-  void connectAtomsWithLayoutEdge(COFFDefinedAtom *a,
-                                  COFFDefinedAtom *b) const {
-    addEdge(a, b, lld::Reference::kindLayoutAfter);
-    addEdge(b, a, lld::Reference::kindLayoutBefore);
-  }
-
   /// Connect atoms appeared in the same section with layout-{before,after}
   /// edges. It has two purposes.
   ///
@@ -335,9 +329,8 @@ private:
     COFFDefinedAtom *atom = findAtomAt(rel->VirtualAddress, section, atoms);
     uint32_t offsetInAtom = itemAddress - atom->originalOffset();
     assert(offsetInAtom < atom->size());
-    COFFReference *ref = new (AtomStorage.Allocate<COFFReference>())
-        COFFReference(targetAtom, offsetInAtom, rel->Type);
-    atom->addReference(ref);
+    atom->addReference(std::unique_ptr<COFFReference>(
+        new COFFReference(targetAtom, offsetInAtom, rel->Type)));
     return error_code::success();
   }
 

lld/test/pecoff/Inputs/grouped-sections.asm

@@ -0,0 +1,17 @@
+.386
+.model flat, c
+
+_data$2 SEGMENT BYTE alias(".data$2")
+	db "orld", 0
+_data$2 ends
+
+_data$1 SEGMENT BYTE alias(".data$1")
+	db "o, w"
+_data$1 ends
+
+.data
+	db "Hell"
+
+.code
+main:
+end main

lld/test/pecoff/grouped-sections.test

@@ -0,0 +1,15 @@
+# RUN: lld -flavor link -out %t1 -subsystem console -force \
+# RUN:   -- %p/Inputs/grouped-sections.obj && llvm-objdump -s %t1 | FileCheck %s
+#
+# The file "grouped-sections.obj" has three data sections in the following
+# order:
+#
+#   .data$2
+#   .data$1
+#   .data
+#
+# If all the sections will be merged correctly, the resulting ".data"
+# section will have the string "Hello, world".
+
+CHECK: Contents of section .rdata:
+CHECK-NEXT: Hello, world