1 files changed, 25 insertions, 22 deletions

diff --git a/lib/ReaderWriter/MachO/LayoutPass.cpp b/lib/ReaderWriter/MachO/LayoutPass.cpp
index 0c14ee9d39ab..dd2ee8567ec9 100644
--- a/lib/ReaderWriter/MachO/LayoutPass.cpp
+++ b/lib/ReaderWriter/MachO/LayoutPass.cpp
@@ -17,6 +17,7 @@
 #include "llvm/Support/Debug.h"
 #include <algorithm>
 #include <set>
+#include <utility>
 
 using namespace lld;
 
@@ -133,7 +134,7 @@ static void checkReachabilityFromRoot(AtomToAtomT &followOnRoots,
   }
 }
 
-static void printDefinedAtoms(const SimpleFile::DefinedAtomRange &atomRange) {
+static void printDefinedAtoms(const File::AtomRange<DefinedAtom> &atomRange) {
   for (const DefinedAtom *atom : atomRange) {
     llvm::dbgs() << "  file=" << atom->file().path()
                  << ", name=" << atom->name()
@@ -146,7 +147,7 @@ static void printDefinedAtoms(const SimpleFile::DefinedAtomRange &atomRange) {
 
 /// Verify that the followon chain is sane. Should not be called in
 /// release binary.
-void LayoutPass::checkFollowonChain(SimpleFile::DefinedAtomRange &range) {
+void LayoutPass::checkFollowonChain(const File::AtomRange<DefinedAtom> &range) {
   ScopedTask task(getDefaultDomain(), "LayoutPass::checkFollowonChain");
 
   // Verify that there's no cycle in follow-on chain.
@@ -176,8 +177,8 @@ 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;
+  const DefinedAtom *left = lc._atom.get();
+  const DefinedAtom *right = rc._atom.get();
   if (left == right) {
     reason = "same";
     return false;
@@ -252,14 +253,15 @@ static bool compareAtoms(const LayoutPass::SortKey &lc,
   bool result = compareAtomsSub(lc, rc, customSorter, reason);
   DEBUG({
     StringRef comp = result ? "<" : ">=";
-    llvm::dbgs() << "Layout: '" << lc._atom->name() << "' " << comp << " '"
-                 << rc._atom->name() << "' (" << reason << ")\n";
+    llvm::dbgs() << "Layout: '" << lc._atom.get()->name()
+                 << "' " << comp << " '"
+                 << rc._atom.get()->name() << "' (" << reason << ")\n";
   });
   return result;
 }
 
 LayoutPass::LayoutPass(const Registry &registry, SortOverride sorter)
-  : _registry(registry), _customSorter(sorter) {}
+    : _registry(registry), _customSorter(std::move(sorter)) {}
 
 // Returns the atom immediately followed by the given atom in the followon
 // chain.
@@ -329,12 +331,12 @@ void LayoutPass::setChainRoot(const DefinedAtom *targetAtom,
 /// d) If the targetAtom is part of a different chain and the root of the
 ///    targetAtom until the targetAtom has all atoms of size 0, then chain the
 ///    targetAtoms and its tree to the current chain
-void LayoutPass::buildFollowOnTable(SimpleFile::DefinedAtomRange &range) {
+void LayoutPass::buildFollowOnTable(const File::AtomRange<DefinedAtom> &range) {
   ScopedTask task(getDefaultDomain(), "LayoutPass::buildFollowOnTable");
   // Set the initial size of the followon and the followonNext hash to the
   // number of atoms that we have.
-  _followOnRoots.resize(range.size());
-  _followOnNexts.resize(range.size());
+  _followOnRoots.reserve(range.size());
+  _followOnNexts.reserve(range.size());
   for (const DefinedAtom *ai : range) {
     for (const Reference *r : *ai) {
       if (r->kindNamespace() != lld::Reference::KindNamespace::all ||
@@ -397,7 +399,8 @@ void LayoutPass::buildFollowOnTable(SimpleFile::DefinedAtomRange &range) {
 /// assigning ordinals to each atom, if the atoms have their ordinals
 /// already assigned skip the atom and move to the next. This is the
 /// main map thats used to sort the atoms while comparing two atoms together
-void LayoutPass::buildOrdinalOverrideMap(SimpleFile::DefinedAtomRange &range) {
+void
+LayoutPass::buildOrdinalOverrideMap(const File::AtomRange<DefinedAtom> &range) {
   ScopedTask task(getDefaultDomain(), "LayoutPass::buildOrdinalOverrideMap");
   uint64_t index = 0;
   for (const DefinedAtom *ai : range) {
@@ -417,31 +420,31 @@ void LayoutPass::buildOrdinalOverrideMap(SimpleFile::DefinedAtomRange &range) {
 }
 
 std::vector<LayoutPass::SortKey>
-LayoutPass::decorate(SimpleFile::DefinedAtomRange &atomRange) const {
+LayoutPass::decorate(File::AtomRange<DefinedAtom> &atomRange) const {
   std::vector<SortKey> ret;
-  for (const DefinedAtom *atom : atomRange) {
-    auto ri = _followOnRoots.find(atom);
-    auto oi = _ordinalOverrideMap.find(atom);
-    const DefinedAtom *root = (ri == _followOnRoots.end()) ? atom : ri->second;
+  for (OwningAtomPtr<DefinedAtom> &atom : atomRange.owning_ptrs()) {
+    auto ri = _followOnRoots.find(atom.get());
+    auto oi = _ordinalOverrideMap.find(atom.get());
+    const auto *root = (ri == _followOnRoots.end()) ? atom.get() : ri->second;
     uint64_t override = (oi == _ordinalOverrideMap.end()) ? 0 : oi->second;
-    ret.push_back(SortKey(atom, root, override));
+    ret.push_back(SortKey(std::move(atom), root, override));
   }
   return ret;
 }
 
-void LayoutPass::undecorate(SimpleFile::DefinedAtomRange &atomRange,
+void LayoutPass::undecorate(File::AtomRange<DefinedAtom> &atomRange,
                             std::vector<SortKey> &keys) const {
   size_t i = 0;
   for (SortKey &k : keys)
-    atomRange[i++] = k._atom;
+    atomRange[i++] = std::move(k._atom);
 }
 
 /// Perform the actual pass
-std::error_code LayoutPass::perform(SimpleFile &mergedFile) {
+llvm::Error LayoutPass::perform(SimpleFile &mergedFile) {
   DEBUG(llvm::dbgs() << "******** Laying out atoms:\n");
   // sort the atoms
   ScopedTask task(getDefaultDomain(), "LayoutPass");
-  SimpleFile::DefinedAtomRange atomRange = mergedFile.definedAtoms();
+  File::AtomRange<DefinedAtom> atomRange = mergedFile.defined();
 
   // Build follow on tables
   buildFollowOnTable(atomRange);
@@ -471,7 +474,7 @@ std::error_code LayoutPass::perform(SimpleFile &mergedFile) {
   });
 
   DEBUG(llvm::dbgs() << "******** Finished laying out atoms\n");
-  return std::error_code();
+  return llvm::Error();
 }
 
 void addLayoutPass(PassManager &pm, const MachOLinkingContext &ctx) {