[mach-o] Add support for -order_file option

The darwin linker lets you rearrange functions and data for better locality
(less paging).  You do this with the -order_file option which supplies a text
file containing one symbol per line.

Implementing this required a small change to LayoutPass to add a custom sorter
hook.

llvm-svn: 221545

8 files changed, 282 insertions, 12 deletions

diff --git a/lld/include/lld/Passes/LayoutPass.h b/lld/include/lld/Passes/LayoutPass.h
index e59a3cc6236..2196eb5a1b3 100644
--- a/lld/include/lld/Passes/LayoutPass.h
+++ b/lld/include/lld/Passes/LayoutPass.h
@@ -37,7 +37,10 @@ public:
     uint64_t _override;
   };
 
-  LayoutPass(const Registry &registry);
+  typedef std::function<bool (const DefinedAtom *left, const DefinedAtom *right,
+                              bool &leftBeforeRight)> SortOverride;
+
+  LayoutPass(const Registry &registry, SortOverride sorter=nullptr);
 
   /// Sorts atoms in mergedFile by content type then by command line order.
   void perform(std::unique_ptr<MutableFile> &mergedFile) override;
@@ -57,6 +60,7 @@ private:
   void buildOrdinalOverrideMap(MutableFile::DefinedAtomRange &range);
 
   const Registry &_registry;
+  SortOverride _customSorter;
 
   typedef llvm::DenseMap<const DefinedAtom *, const DefinedAtom *> AtomToAtomT;
   typedef llvm::DenseMap<const DefinedAtom *, uint64_t> AtomToOrdinalT;
diff --git a/lld/include/lld/ReaderWriter/MachOLinkingContext.h b/lld/include/lld/ReaderWriter/MachOLinkingContext.h
index 0792bd2e51b..ef4de3b47c8 100644
--- a/lld/include/lld/ReaderWriter/MachOLinkingContext.h
+++ b/lld/include/lld/ReaderWriter/MachOLinkingContext.h
@@ -103,6 +103,8 @@ public:
     _debugInfoMode = mode;
   }
 
+  void appendOrderedSymbol(StringRef symbol, StringRef filename);
+
   bool keepPrivateExterns() const { return _keepPrivateExterns; }
   void setKeepPrivateExterns(bool v) { _keepPrivateExterns = v; }
   bool demangleSymbols() const { return _demangle; }
@@ -282,8 +284,8 @@ private:
   mach_o::MachODylibFile* loadIndirectDylib(StringRef path);
   void checkExportWhiteList(const DefinedAtom *atom) const;
   void checkExportBlackList(const DefinedAtom *atom) const;
-
-
+  bool customAtomOrderer(const DefinedAtom *left, const DefinedAtom *right,
+                         bool &leftBeforeRight);
   struct ArchInfo {
     StringRef                 archName;
     MachOLinkingContext::Arch arch;
@@ -298,6 +300,14 @@ private:
     uint8_t   align2;
   };
 
+  struct OrderFileNode {
+    StringRef fileFilter;
+    unsigned  order;
+  };
+
+  static bool findOrderOrdinal(const std::vector<OrderFileNode> &nodes,
+                             const DefinedAtom *atom, unsigned &ordinal);
+
   static ArchInfo _s_archInfos[];
 
   std::set<StringRef> _existingPaths; // For testing only.
@@ -334,6 +344,8 @@ private:
   llvm::StringSet<> _exportedSymbols;
   DebugInfoMode _debugInfoMode;
   std::unique_ptr<llvm::raw_fd_ostream> _dependencyInfo;
+  llvm::StringMap<std::vector<OrderFileNode>> _orderFiles;
+  unsigned _orderFileEntries;
 };
 
 } // end namespace lld
diff --git a/lld/lib/Driver/DarwinLdDriver.cpp b/lld/lib/Driver/DarwinLdDriver.cpp
index 86f19e31cac..21cd691a578 100644
--- a/lld/lib/Driver/DarwinLdDriver.cpp
+++ b/lld/lib/Driver/DarwinLdDriver.cpp
@@ -120,6 +120,59 @@ static std::error_code parseExportsList(StringRef exportFilePath,
   return std::error_code();
 }
 
+
+
+/// Order files are one symbol per line. Blank lines are ignored.
+/// Trailing comments start with #. Symbol names can be prefixed with an
+/// architecture name and/or .o leaf name.  Examples:
+///     _foo
+///     bar.o:_bar
+///     libfrob.a(bar.o):_bar
+///     x86_64:_foo64
+static std::error_code parseOrderFile(StringRef orderFilePath,
+                                      MachOLinkingContext &ctx,
+                                      raw_ostream &diagnostics) {
+  // Map in order file.
+  ErrorOr<std::unique_ptr<MemoryBuffer>> mb =
+                                   MemoryBuffer::getFileOrSTDIN(orderFilePath);
+  if (std::error_code ec = mb.getError())
+    return ec;
+  ctx.addInputFileDependency(orderFilePath);
+  StringRef buffer = mb->get()->getBuffer();
+  while (!buffer.empty()) {
+    // Split off each line in the file.
+    std::pair<StringRef, StringRef> lineAndRest = buffer.split('\n');
+    StringRef line = lineAndRest.first;
+    buffer = lineAndRest.second;
+    // Ignore trailing # comments.
+    std::pair<StringRef, StringRef> symAndComment = line.split('#');
+    if (symAndComment.first.empty())
+      continue;
+    StringRef sym = symAndComment.first.trim();
+    if (sym.empty())
+      continue;
+    // Check for prefix.
+    StringRef prefix;
+    std::pair<StringRef, StringRef> prefixAndSym = sym.split(':');
+    if (!prefixAndSym.second.empty()) {
+      sym = prefixAndSym.second;
+      prefix = prefixAndSym.first;
+      if (!prefix.endswith(".o") && !prefix.endswith(".o)")) {
+        // If arch name prefix does not match arch being linked, ignore symbol.
+        if (!ctx.archName().equals(prefix))
+          continue;
+        prefix = "";
+      }
+    } else
+     sym = prefixAndSym.first;
+    if (!sym.empty()) {
+      ctx.appendOrderedSymbol(sym, prefix);
+      //llvm::errs() << sym << ", prefix=" << prefix << "\n";
+    }
+  }
+  return std::error_code();
+}
+
 //
 // There are two variants of the  -filelist option:
 //
@@ -644,6 +697,18 @@ bool DarwinLdDriver::parse(int argc, const char *argv[],
   if (parsedArgs->hasArg(OPT_S))
     ctx.setDebugInfoMode(MachOLinkingContext::DebugInfoMode::noDebugMap);
 
+  // Handle -order_file <file>
+  for (auto orderFile : parsedArgs->filtered(OPT_order_file)) {
+    if (std::error_code ec = parseOrderFile(orderFile->getValue(), ctx,
+                                              diagnostics)) {
+      diagnostics << "error: " << ec.message()
+                  << ", processing '-order_file "
+                  << orderFile->getValue()
+                  << "'\n";
+      return false;
+    }
+  }
+
   // Handle input files
   for (auto &arg : *parsedArgs) {
     bool upward;
diff --git a/lld/lib/Driver/DarwinLdOptions.td b/lld/lib/Driver/DarwinLdOptions.td
index eefb2c7b92d..630f1e787b0 100644
--- a/lld/lib/Driver/DarwinLdOptions.td
+++ b/lld/lib/Driver/DarwinLdOptions.td
@@ -51,6 +51,10 @@ def unexported_symbol : Separate<["-"], "unexported_symbol">,
 def keep_private_externs : Flag<["-"], "keep_private_externs">,
      HelpText<"Private extern (hidden) symbols should not be transformed "
               "into local symbols">, Group<grp_opts>;
+def order_file : Separate<["-"], "order_file">,
+     MetaVarName<"<file-path>">,
+     HelpText<"re-order and move specified symbols to start of their section">,
+     Group<grp_opts>;
 
 // main executable options
 def grp_main : OptionGroup<"opts">, HelpText<"MAIN EXECUTABLE OPTIONS">;
diff --git a/lld/lib/Passes/LayoutPass.cpp b/lld/lib/Passes/LayoutPass.cpp
index cb6d082dad3..44d63b7ea8c 100644
--- a/lld/lib/Passes/LayoutPass.cpp
+++ b/lld/lib/Passes/LayoutPass.cpp
@@ -19,7 +19,8 @@ using namespace lld;
 #define DEBUG_TYPE "LayoutPass"
 
 static bool compareAtoms(const LayoutPass::SortKey &,
-                         const LayoutPass::SortKey &);
+                         const LayoutPass::SortKey &,
+                         LayoutPass::SortOverride customSorter=nullptr);
 
 #ifndef NDEBUG
 // Return "reason (leftval, rightval)"
@@ -161,10 +162,12 @@ void LayoutPass::checkFollowonChain(MutableFile::DefinedAtomRange &range) {
 /// b) Sorts atoms by their ordinal overrides (layout-after/ingroup)
 /// c) Sorts atoms by their permissions
 /// d) Sorts atoms by their content
-/// e) Sorts atoms on how they appear using File Ordinality
-/// f) Sorts atoms on how they appear within the File
+/// e) If custom sorter provided, let it sort
+/// f) Sorts atoms on how they appear using File Ordinality
+/// g) Sorts atoms on how they appear within the File
 static bool compareAtomsSub(const LayoutPass::SortKey &lc,
                             const LayoutPass::SortKey &rc,
+                            LayoutPass::SortOverride customSorter,
                             std::string &reason) {
   const DefinedAtom *left = lc._atom;
   const DefinedAtom *right = rc._atom;
@@ -216,6 +219,13 @@ static bool compareAtomsSub(const LayoutPass::SortKey &lc,
     return leftType < rightType;
   }
 
+  // Use custom sorter if supplied.
+  if (customSorter) {
+    bool leftBeforeRight;
+    if (customSorter(leftRoot, rightRoot, leftBeforeRight))
+      return leftBeforeRight;
+  }
+
   // Sort by .o order.
   const File *leftFile = &leftRoot->file();
   const File *rightFile = &rightRoot->file();
@@ -242,9 +252,10 @@ static bool compareAtomsSub(const LayoutPass::SortKey &lc,
 }
 
 static bool compareAtoms(const LayoutPass::SortKey &lc,
-                         const LayoutPass::SortKey &rc) {
+                         const LayoutPass::SortKey &rc,
+                         LayoutPass::SortOverride customSorter) {
   std::string reason;
-  bool result = compareAtomsSub(lc, rc, reason);
+  bool result = compareAtomsSub(lc, rc, customSorter, reason);
   DEBUG({
     StringRef comp = result ? "<" : ">=";
     llvm::dbgs() << "Layout: '" << lc._atom->name() << "' " << comp << " '"
@@ -253,7 +264,8 @@ static bool compareAtoms(const LayoutPass::SortKey &lc,
   return result;
 }
 
-LayoutPass::LayoutPass(const Registry &registry) : _registry(registry) {}
+LayoutPass::LayoutPass(const Registry &registry, SortOverride sorter)
+  : _registry(registry), _customSorter(sorter) {}
 
 // Returns the atom immediately followed by the given atom in the followon
 // chain.
@@ -522,7 +534,10 @@ void LayoutPass::perform(std::unique_ptr<MutableFile> &mergedFile) {
   });
 
   std::vector<LayoutPass::SortKey> vec = decorate(atomRange);
-  std::sort(vec.begin(), vec.end(), compareAtoms);
+  std::sort(vec.begin(), vec.end(),
+      [&](const LayoutPass::SortKey &l, const LayoutPass::SortKey &r) -> bool {
+        return compareAtoms(l, r, _customSorter);
+      });
   DEBUG(checkTransitivity(vec));
   undecorate(atomRange, vec);
 
diff --git a/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp b/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp
index 85c20b3a932..86b9c3d28e5 100644
--- a/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp
@@ -146,7 +146,7 @@ MachOLinkingContext::MachOLinkingContext()
       _printAtoms(false), _testingFileUsage(false), _keepPrivateExterns(false),
       _demangle(false), _archHandler(nullptr),
       _exportMode(ExportMode::globals),
-      _debugInfoMode(DebugInfoMode::addDebugMap) {}
+      _debugInfoMode(DebugInfoMode::addDebugMap), _orderFileEntries(0) {}
 
 MachOLinkingContext::~MachOLinkingContext() {}
 
@@ -577,7 +577,12 @@ bool MachOLinkingContext::validateImpl(raw_ostream &diagnostics) {
 }
 
 void MachOLinkingContext::addPasses(PassManager &pm) {
-  pm.add(std::unique_ptr<Pass>(new LayoutPass(registry())));
+  pm.add(std::unique_ptr<Pass>(new LayoutPass(
+      registry(), [&](const DefinedAtom * left, const DefinedAtom * right,
+                      bool & leftBeforeRight)
+                      ->bool {
+    return customAtomOrderer(left, right, leftBeforeRight);
+  })));
   if (needsStubsPass())
     mach_o::addStubsPass(pm, *this);
   if (needsCompactUnwindPass())
@@ -825,5 +830,84 @@ void MachOLinkingContext::addOutputFileDependency(StringRef path) const {
   *_dependencyInfo << '\0';
 }
 
+void MachOLinkingContext::appendOrderedSymbol(StringRef symbol,
+                                              StringRef filename) {
+  // To support sorting static functions which may have the same name in
+  // multiple .o files, _orderFiles maps the symbol name to a vector
+  // of OrderFileNode each of which can specify a file prefix.
+  OrderFileNode info;
+  if (!filename.empty())
+    info.fileFilter = copy(filename);
+  info.order = _orderFileEntries++;
+  _orderFiles[symbol].push_back(info);
+}
+
+bool
+MachOLinkingContext::findOrderOrdinal(const std::vector<OrderFileNode> &nodes,
+                                      const DefinedAtom *atom,
+                                      unsigned &ordinal) {
+  const File *objFile = &atom->file();
+  assert(objFile);
+  StringRef objName = objFile->path();
+  std::pair<StringRef, StringRef> dirAndLeaf = objName.rsplit('/');
+  if (!dirAndLeaf.second.empty())
+    objName = dirAndLeaf.second;
+  for (const OrderFileNode &info : nodes) {
+    if (info.fileFilter.empty()) {
+      // Have unprefixed symbol name in order file that matches this atom.
+      ordinal = info.order;
+      llvm::errs() << "ordered " << atom->name() << "\n";
+      return true;
+    }
+    if (info.fileFilter.equals(objName)) {
+      // Have prefixed symbol name in order file that matches atom's path.
+      ordinal = info.order;
+      llvm::errs() << "ordered " << atom->name() << " with prefix '"
+                   << info.fileFilter << "'\n";
+      return true;
+    }
+  }
+  return false;
+}
+
+bool MachOLinkingContext::customAtomOrderer(const DefinedAtom *left,
+                                            const DefinedAtom *right,
+                                            bool &leftBeforeRight) {
+  // No custom sorting if no order file entries.
+  if (!_orderFileEntries)
+    return false;
+
+  // Order files can only order named atoms.
+  StringRef leftName = left->name();
+  StringRef rightName = right->name();
+  if (leftName.empty() || rightName.empty())
+    return false;
+
+  // If neither is in order file list, no custom sorter.
+  auto leftPos = _orderFiles.find(leftName);
+  auto rightPos = _orderFiles.find(rightName);
+  bool leftIsOrdered = (leftPos != _orderFiles.end());
+  bool rightIsOrdered = (rightPos != _orderFiles.end());
+  if (!leftIsOrdered && !rightIsOrdered)
+    return false;
+
+  // There could be multiple symbols with same name but different file prefixes.
+  unsigned leftOrder;
+  unsigned rightOrder;
+  bool foundLeft =
+      leftIsOrdered && findOrderOrdinal(leftPos->getValue(), left, leftOrder);
+  bool foundRight = rightIsOrdered &&
+                    findOrderOrdinal(rightPos->getValue(), right, rightOrder);
+  if (!foundLeft && !foundRight)
+    return false;
+
+  // If only one is in order file list, ordered one goes first.
+  if (foundLeft != foundRight)
+    leftBeforeRight = foundLeft;
+  else
+    leftBeforeRight = (leftOrder < rightOrder);
+
+  return true;
+}
 
 } // end namespace lld
diff --git a/lld/test/mach-o/Inputs/order_file-basic.order b/lld/test/mach-o/Inputs/order_file-basic.order
new file mode 100644
index 00000000000..0ac90cb79da
--- /dev/null
+++ b/lld/test/mach-o/Inputs/order_file-basic.order
@@ -0,0 +1,11 @@
+
+# input file for order_file-basic.yaml
+
+_func2
+libfoo.a(foo.o):_foo  # tests file specific ordering within archive
+i386:_func3           # wrong arch, so ignored
+armv7:_func3          # wrong arch, so ignored
+_func1
+_notfound             # unknown symbol silently ignored
+_data3                # data symbols should be orderable
+
diff --git a/lld/test/mach-o/order_file-basic.yaml b/lld/test/mach-o/order_file-basic.yaml
new file mode 100644
index 00000000000..3fea9be1560
--- /dev/null
+++ b/lld/test/mach-o/order_file-basic.yaml
@@ -0,0 +1,75 @@
+# RUN: lld -flavor darwin -arch x86_64 %s %p/Inputs/libSystem.yaml \
+# RUN:     -order_file %p/Inputs/order_file-basic.order \
+# RUN:     -force_load %p/Inputs/libfoo.a -o %t
+# RUN: llvm-nm -m -n %t | FileCheck %s
+#
+# Test -order_file
+#
+
+--- !mach-o
+arch:            x86_64
+file-type:       MH_OBJECT
+flags:           [ MH_SUBSECTIONS_VIA_SYMBOLS ]
+sections:
+  - segment:         __TEXT
+    section:         __text
+    type:            S_REGULAR
+    attributes:      [ S_ATTR_PURE_INSTRUCTIONS, S_ATTR_SOME_INSTRUCTIONS ]
+    address:         0x0000000000000000
+    content:         [ 0xC3, 0xC3, 0xC3, 0xC3 ]
+  - segment:         __DATA
+    section:         __data
+    type:            S_REGULAR
+    attributes:      [  ]
+    alignment:       2
+    address:         0x0000000000000014
+    content:         [ 0x05, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00,
+                       0x07, 0x00, 0x00, 0x00 ]
+global-symbols:
+  - name:            _data1
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            2
+    value:           0x0000000000000014
+  - name:            _data2
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            2
+    value:           0x0000000000000018
+  - name:            _data3
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            2
+    value:           0x000000000000001C
+  - name:            _func1
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            1
+    value:           0x0000000000000000
+  - name:            _func2
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            1
+    value:           0x0000000000000001
+  - name:            _func3
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            1
+    value:           0x0000000000000002
+  - name:            _main
+    type:            N_SECT
+    scope:           [ N_EXT ]
+    sect:            1
+    value:           0x0000000000000003
+...
+
+
+# CHECK:	{{[0-9a-f]+}} (__TEXT,__text) external _func2
+# CHECK:	{{[0-9a-f]+}} (__TEXT,__text) external _foo
+# CHECK:	{{[0-9a-f]+}} (__TEXT,__text) external _func1
+# CHECK:	{{[0-9a-f]+}} (__TEXT,__text) external _func3
+# CHECK:	{{[0-9a-f]+}} (__TEXT,__text) external _main
+# CHECK:	{{[0-9a-f]+}} (__DATA,__data) external _data3
+# CHECK:	{{[0-9a-f]+}} (__DATA,__data) external _data1
+# CHECK:	{{[0-9a-f]+}} (__DATA,__data) external _data2
+