4 files changed, 51 insertions, 1 deletions

diff --git a/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h b/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h
index 06445f7d55f..3e6e649b4d6 100644
--- a/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h
+++ b/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h
@@ -41,7 +41,7 @@ public:
   virtual Writer &writer() const;
   virtual bool validateImpl(raw_ostream &diagnostics);
 
-  virtual void addPasses(PassManager &pm) const {}
+  virtual void addPasses(PassManager &pm) const;
 
   void setStackReserve(uint64_t size) { _stackReserve = size; }
   void setStackCommit(uint64_t size) { _stackCommit = size; }
diff --git a/lld/lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp b/lld/lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp
index 566161516b0..40ac53d068c 100644
--- a/lld/lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp
@@ -9,6 +9,8 @@
 
 #include "lld/ReaderWriter/PECOFFTargetInfo.h"
 
+#include "lld/Core/PassManager.h"
+#include "lld/Passes/LayoutPass.h"
 #include "lld/ReaderWriter/Reader.h"
 #include "lld/ReaderWriter/Writer.h"
 
@@ -56,4 +58,8 @@ PECOFFTargetInfo::stringFromRelocKind(Reference::Kind kind) const {
   return make_error_code(yaml_reader_error::illegal_value);
 }
 
+void PECOFFTargetInfo::addPasses(PassManager &pm) const {
+  pm.add(std::unique_ptr<Pass>(new LayoutPass()));
+}
+
 } // end namespace lld
diff --git a/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp b/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
index aa2b782fed1..1b43a1517fb 100644
--- a/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
@@ -41,6 +41,10 @@ namespace { // anonymous
 /// to be fixed up so that the address points to atom Y's address.
 class COFFReference LLVM_FINAL : public Reference {
 public:
+  COFFReference(Kind kind) : _target(nullptr), _offsetInAtom(0) {
+    _kind = kind;
+  }
+
   COFFReference(const Atom *target, uint32_t offsetInAtom, uint16_t relocType)
       : _target(target), _offsetInAtom(offsetInAtom) {
     setKind(static_cast<Reference::Kind>(relocType));
@@ -431,6 +435,38 @@ private:
     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.
+  ///
+  ///   - To prevent atoms from being GC'ed (aka dead-stripped) if there is a
+  ///     reference to one of the atoms. In that case we want to emit all the
+  ///     atoms appeared in the same section, because the referenced "live"
+  ///     atom may reference other atoms in the same section. If we don't add
+  ///     edges between atoms, unreferenced atoms in the same section would be
+  ///     GC'ed.
+  ///   - To preserve the order of atmos. We want to emit the atoms in the
+  ///     same order as they appeared in the input object file.
+  void addLayoutEdges(vector<COFFDefinedAtom *> &definedAtoms) const {
+    if (definedAtoms.size() <= 1)
+      return;
+    for (auto it = definedAtoms.begin(), e = definedAtoms.end(); it + 1 != e;
+         ++it)
+      connectAtomsWithLayoutEdge(*it, *(it + 1));
+  }
+
   error_code AtomizeDefinedSymbols(SectionToSymbolsT &definedSymbols,
                                    vector<const DefinedAtom *> &definedAtoms,
                                    SymbolNameToAtomT &symbolToAtom,
@@ -445,6 +481,9 @@ private:
               AtomizeDefinedSymbolsInSection(section, symbols, atoms))
         return ec;
 
+      // Connect atoms with layout-before/layout-after edges.
+      addLayoutEdges(atoms);
+
       for (COFFDefinedAtom *atom : atoms) {
         if (!atom->name().empty())
           symbolToAtom[atom->name()] = atom;
diff --git a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
index 38dd519dda0..2e4e230590c 100644
--- a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -336,6 +336,11 @@ public:
         auto relocSite = reinterpret_cast<llvm::support::ulittle32_t *>(
             fileBuffer + layout->_fileOffset + ref->offsetInAtom());
         uint64_t targetAddr = atomToVirtualAddr[ref->target()];
+
+        // Skip if this reference is not for relocation.
+        if (ref->kind() < lld::Reference::kindTargetLow)
+          continue;
+
         switch (ref->kind()) {
         case llvm::COFF::IMAGE_REL_I386_DIR32:
           *relocSite = targetAddr;