Use owning pointers instead of raw pointers for Atom's to fix leaks.

This is a re-commit of r264022 with a fix for MSVC.  The issue there was
that the code was running DefinedAtom::~Atom() for some value and instead
needed to cast to Atom before running ~Atom.  Original commit message follows.

Currently each File contains an BumpPtrAllocator in which Atom's are
allocated.  Some Atom's contain data structures like std::vector which
leak as we don't run ~Atom when they are BumpPtrAllocate'd.

Now each File actually owns its Atom's using an OwningAtomPtr.  This
is analygous to std::unique_ptr and may be replaced by it if possible.

An Atom can therefore only be owned by a single File, so the Resolver now
moves them from one File to another.  The MachOLinkingContext owns the File's
and so clears all the Atom's in ~MachOLinkingContext, then delete's all the
File's.  This makes sure all Atom's have been destructed before any of the
BumpPtrAllocator's in which they run have gone away.

Should hopefully fix the remaining leaks.  Will keep an eye on the bots to
make sure.

llvm-svn: 264067
diff --git a/lld/lib/ReaderWriter/FileArchive.cpp b/lld/lib/ReaderWriter/FileArchive.cpp
index a47024c..73b8612 100644
--- a/lld/lib/ReaderWriter/FileArchive.cpp
+++ b/lld/lib/ReaderWriter/FileArchive.cpp
@@ -88,22 +88,29 @@
     return std::error_code();
   }
 
-  const AtomVector<DefinedAtom> &defined() const override {
+  const AtomRange<DefinedAtom> defined() const override {
     return _noDefinedAtoms;
   }
 
-  const AtomVector<UndefinedAtom> &undefined() const override {
+  const AtomRange<UndefinedAtom> undefined() const override {
     return _noUndefinedAtoms;
   }
 
-  const AtomVector<SharedLibraryAtom> &sharedLibrary() const override {
+  const AtomRange<SharedLibraryAtom> sharedLibrary() const override {
     return _noSharedLibraryAtoms;
   }
 
-  const AtomVector<AbsoluteAtom> &absolute() const override {
+  const AtomRange<AbsoluteAtom> absolute() const override {
     return _noAbsoluteAtoms;
   }
 
+  void clearAtoms() override {
+    _noDefinedAtoms.clear();
+    _noUndefinedAtoms.clear();
+    _noSharedLibraryAtoms.clear();
+    _noAbsoluteAtoms.clear();
+  }
+
 protected:
   std::error_code doParse() override {
     // Make Archive object which will be owned by FileArchive object.
diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp b/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp
index cefa745..7eda42e 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp
@@ -1429,6 +1429,8 @@
     _name = tmp.copy(file.allocator());
   }
 
+  ~Thumb2ToArmShimAtom() override = default;
+
   StringRef name() const override {
     return _name;
   }
@@ -1472,6 +1474,8 @@
     _name = tmp.copy(file.allocator());
   }
 
+  ~ArmToThumbShimAtom() override = default;
+
   StringRef name() const override {
     return _name;
   }
diff --git a/lld/lib/ReaderWriter/MachO/Atoms.h b/lld/lib/ReaderWriter/MachO/Atoms.h
index f50a6cf..c3117d4 100644
--- a/lld/lib/ReaderWriter/MachO/Atoms.h
+++ b/lld/lib/ReaderWriter/MachO/Atoms.h
@@ -32,6 +32,8 @@
         _contentType(type), _scope(scope), _merge(mergeNo), _thumb(false),
         _noDeadStrip(noDeadStrip) {}
 
+  ~MachODefinedAtom() override = default;
+
   uint64_t size() const override { return _content.size(); }
 
   ContentType contentType() const override { return _contentType; }
@@ -83,6 +85,8 @@
                          content, align),
         _sectionName(sectionName) {}
 
+  ~MachODefinedCustomSectionAtom() override = default;
+
   SectionChoice sectionChoice() const override {
     return DefinedAtom::sectionCustomRequired;
   }
@@ -101,6 +105,8 @@
       : SimpleDefinedAtom(f), _name(name), _scope(scope), _size(size),
         _align(align) {}
 
+  ~MachOTentativeDefAtom() override = default;
+
   uint64_t size() const override { return _size; }
 
   Merge merge() const override { return DefinedAtom::mergeAsTentative; }
diff --git a/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp b/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp
index d279303..088f93b 100644
--- a/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/CompactUnwindPass.cpp
@@ -88,6 +88,8 @@
     addSecondLevelPages(pages);
   }
 
+  ~UnwindInfoAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeProcessedUnwindInfo;
   }
diff --git a/lld/lib/ReaderWriter/MachO/ExecutableAtoms.h b/lld/lib/ReaderWriter/MachO/ExecutableAtoms.h
index e1a252b..acced33 100644
--- a/lld/lib/ReaderWriter/MachO/ExecutableAtoms.h
+++ b/lld/lib/ReaderWriter/MachO/ExecutableAtoms.h
@@ -122,21 +122,28 @@
           ArrayRef<uint8_t>(), DefinedAtom::Alignment(1)));
   }
 
-  const AtomVector<DefinedAtom> &defined() const override {
+  const AtomRange<DefinedAtom> defined() const override {
     return _definedAtoms;
   }
-  const AtomVector<UndefinedAtom> &undefined() const override {
+  const AtomRange<UndefinedAtom> undefined() const override {
     return _noUndefinedAtoms;
   }
 
-  const AtomVector<SharedLibraryAtom> &sharedLibrary() const override {
+  const AtomRange<SharedLibraryAtom> sharedLibrary() const override {
     return _noSharedLibraryAtoms;
   }
 
-  const AtomVector<AbsoluteAtom> &absolute() const override {
+  const AtomRange<AbsoluteAtom> absolute() const override {
     return _noAbsoluteAtoms;
   }
 
+  void clearAtoms() override {
+    _definedAtoms.clear();
+    _noUndefinedAtoms.clear();
+    _noSharedLibraryAtoms.clear();
+    _noAbsoluteAtoms.clear();
+  }
+
 
 private:
   mutable AtomVector<DefinedAtom> _definedAtoms;
diff --git a/lld/lib/ReaderWriter/MachO/File.h b/lld/lib/ReaderWriter/MachO/File.h
index a0d20ea..f7262bb 100644
--- a/lld/lib/ReaderWriter/MachO/File.h
+++ b/lld/lib/ReaderWriter/MachO/File.h
@@ -275,7 +275,8 @@
 
   MachODylibFile(StringRef path) : SharedLibraryFile(path) {}
 
-  const SharedLibraryAtom *exports(StringRef name, bool isData) const override {
+  OwningAtomPtr<SharedLibraryAtom> exports(StringRef name,
+                                           bool isData) const override {
     // Pass down _installName so that if this requested symbol
     // is re-exported through this dylib, the SharedLibraryAtom's loadName()
     // is this dylib installName and not the implementation dylib's.
@@ -328,25 +329,30 @@
   }
 
 private:
-  const SharedLibraryAtom *exports(StringRef name,
+  OwningAtomPtr<SharedLibraryAtom> exports(StringRef name,
                                    StringRef installName) const {
     // First, check if requested symbol is directly implemented by this dylib.
     auto entry = _nameToAtom.find(name);
     if (entry != _nameToAtom.end()) {
-      if (!entry->second.atom) {
-        // Lazily create SharedLibraryAtom.
-        entry->second.atom =
-          new (allocator()) MachOSharedLibraryAtom(*this, name, installName,
-                                                   entry->second.weakDef);
-      }
-      return entry->second.atom;
+      // FIXME: Make this map a set and only used in assert builds.
+      // Note, its safe to assert here as the resolver is the only client of
+      // this API and it only requests exports for undefined symbols.
+      // If we return from here we are no longer undefined so we should never
+      // get here again.
+      assert(!entry->second.atom && "Duplicate shared library export");
+      bool weakDef = entry->second.weakDef;
+      auto *atom = new (allocator()) MachOSharedLibraryAtom(*this, name,
+                                                            installName,
+                                                            weakDef);
+      entry->second.atom = atom;
+      return atom;
     }
 
     // Next, check if symbol is implemented in some re-exported dylib.
     for (const ReExportedDylib &dylib : _reExportedDylibs) {
       assert(dylib.file);
       auto atom = dylib.file->exports(name, installName);
-      if (atom)
+      if (atom.get())
         return atom;
     }
 
diff --git a/lld/lib/ReaderWriter/MachO/FlatNamespaceFile.h b/lld/lib/ReaderWriter/MachO/FlatNamespaceFile.h
index 6c6a926..08b28f4 100644
--- a/lld/lib/ReaderWriter/MachO/FlatNamespaceFile.h
+++ b/lld/lib/ReaderWriter/MachO/FlatNamespaceFile.h
@@ -25,34 +25,35 @@
   FlatNamespaceFile(const MachOLinkingContext &context)
     : SharedLibraryFile("flat namespace") { }
 
-  const SharedLibraryAtom *exports(StringRef name,
+  OwningAtomPtr<SharedLibraryAtom> exports(StringRef name,
                                    bool dataSymbolOnly) const override {
-    _sharedLibraryAtoms.push_back(
-      new (allocator()) MachOSharedLibraryAtom(*this, name, getDSOName(),
-                                               false));
-
-    return _sharedLibraryAtoms.back();
+    return new (allocator()) MachOSharedLibraryAtom(*this, name, getDSOName(),
+                                                    false);
   }
 
   StringRef getDSOName() const override { return "flat-namespace"; }
 
-  const AtomVector<DefinedAtom> &defined() const override {
+  const AtomRange<DefinedAtom> defined() const override {
     return _noDefinedAtoms;
   }
-  const AtomVector<UndefinedAtom> &undefined() const override {
+  const AtomRange<UndefinedAtom> undefined() const override {
     return _noUndefinedAtoms;
   }
 
-  const AtomVector<SharedLibraryAtom> &sharedLibrary() const override {
-    return _sharedLibraryAtoms;
+  const AtomRange<SharedLibraryAtom> sharedLibrary() const override {
+    return _noSharedLibraryAtoms;
   }
 
-  const AtomVector<AbsoluteAtom> &absolute() const override {
+  const AtomRange<AbsoluteAtom> absolute() const override {
     return _noAbsoluteAtoms;
   }
 
-private:
-  mutable AtomVector<SharedLibraryAtom> _sharedLibraryAtoms;
+  void clearAtoms() override {
+    _noDefinedAtoms.clear();
+    _noUndefinedAtoms.clear();
+    _noSharedLibraryAtoms.clear();
+    _noAbsoluteAtoms.clear();
+  }
 };
 
 } // namespace mach_o
diff --git a/lld/lib/ReaderWriter/MachO/GOTPass.cpp b/lld/lib/ReaderWriter/MachO/GOTPass.cpp
index eac4864..400dbf7 100644
--- a/lld/lib/ReaderWriter/MachO/GOTPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/GOTPass.cpp
@@ -54,6 +54,8 @@
   GOTEntryAtom(const File &file, bool is64, StringRef name)
     : SimpleDefinedAtom(file), _is64(is64), _name(name) { }
 
+  ~GOTEntryAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeGOT;
   }
diff --git a/lld/lib/ReaderWriter/MachO/LayoutPass.cpp b/lld/lib/ReaderWriter/MachO/LayoutPass.cpp
index e64002e..811f74a 100644
--- a/lld/lib/ReaderWriter/MachO/LayoutPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/LayoutPass.cpp
@@ -146,7 +146,7 @@
 
 /// Verify that the followon chain is sane. Should not be called in
 /// release binary.
-void LayoutPass::checkFollowonChain(SimpleFile::DefinedAtomRange &range) {
+void LayoutPass::checkFollowonChain(const SimpleFile::DefinedAtomRange &range) {
   ScopedTask task(getDefaultDomain(), "LayoutPass::checkFollowonChain");
 
   // Verify that there's no cycle in follow-on chain.
@@ -176,8 +176,8 @@
                             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,8 +252,9 @@
   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;
 }
@@ -329,7 +330,7 @@
 /// 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 SimpleFile::DefinedAtomRange &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.
@@ -397,7 +398,8 @@
 /// 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 SimpleFile::DefinedAtomRange &range) {
   ScopedTask task(getDefaultDomain(), "LayoutPass::buildOrdinalOverrideMap");
   uint64_t index = 0;
   for (const DefinedAtom *ai : range) {
@@ -419,12 +421,12 @@
 std::vector<LayoutPass::SortKey>
 LayoutPass::decorate(SimpleFile::DefinedAtomRange &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;
 }
@@ -433,7 +435,7 @@
                             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
diff --git a/lld/lib/ReaderWriter/MachO/LayoutPass.h b/lld/lib/ReaderWriter/MachO/LayoutPass.h
index d6072b0..d254891 100644
--- a/lld/lib/ReaderWriter/MachO/LayoutPass.h
+++ b/lld/lib/ReaderWriter/MachO/LayoutPass.h
@@ -33,9 +33,10 @@
 class LayoutPass : public Pass {
 public:
   struct SortKey {
-    SortKey(const DefinedAtom *atom, const DefinedAtom *root, uint64_t override)
-        : _atom(atom), _root(root), _override(override) {}
-    const DefinedAtom *_atom;
+    SortKey(OwningAtomPtr<DefinedAtom> &&atom,
+            const DefinedAtom *root, uint64_t override)
+    : _atom(std::move(atom)), _root(root), _override(override) {}
+    OwningAtomPtr<DefinedAtom> _atom;
     const DefinedAtom *_root;
     uint64_t _override;
   };
@@ -53,10 +54,10 @@
 private:
   // Build the followOn atoms chain as specified by the kindLayoutAfter
   // reference type
-  void buildFollowOnTable(SimpleFile::DefinedAtomRange &range);
+  void buildFollowOnTable(const SimpleFile::DefinedAtomRange &range);
 
   // Build a map of Atoms to ordinals for sorting the atoms
-  void buildOrdinalOverrideMap(SimpleFile::DefinedAtomRange &range);
+  void buildOrdinalOverrideMap(const SimpleFile::DefinedAtomRange &range);
 
   const Registry &_registry;
   SortOverride _customSorter;
@@ -85,11 +86,12 @@
   void setChainRoot(const DefinedAtom *targetAtom, const DefinedAtom *root);
 
   std::vector<SortKey> decorate(SimpleFile::DefinedAtomRange &atomRange) const;
+
   void undecorate(SimpleFile::DefinedAtomRange &atomRange,
                   std::vector<SortKey> &keys) const;
 
   // Check if the follow-on graph is a correct structure. For debugging only.
-  void checkFollowonChain(SimpleFile::DefinedAtomRange &range);
+  void checkFollowonChain(const SimpleFile::DefinedAtomRange &range);
 };
 
 } // namespace mach_o
diff --git a/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp b/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp
index 1cc87f0..4431f34 100644
--- a/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachOLinkingContext.cpp
@@ -171,7 +171,19 @@
 
 MachOLinkingContext::MachOLinkingContext() {}
 
-MachOLinkingContext::~MachOLinkingContext() {}
+MachOLinkingContext::~MachOLinkingContext() {
+  // Atoms are allocated on BumpPtrAllocator's on File's.
+  // As we transfer atoms from one file to another, we need to clear all of the
+  // atoms before we remove any of the BumpPtrAllocator's.
+  auto &nodes = getNodes();
+  for (unsigned i = 0, e = nodes.size(); i != e; ++i) {
+    FileNode *node = dyn_cast<FileNode>(nodes[i].get());
+    if (!node)
+      continue;
+    File *file = node->getFile();
+    file->clearAtoms();
+  }
+}
 
 void MachOLinkingContext::configure(HeaderFileType type, Arch arch, OS os,
                                     uint32_t minOSVersion,
diff --git a/lld/lib/ReaderWriter/MachO/ObjCPass.cpp b/lld/lib/ReaderWriter/MachO/ObjCPass.cpp
index 13fa988..f8c7310 100644
--- a/lld/lib/ReaderWriter/MachO/ObjCPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/ObjCPass.cpp
@@ -56,6 +56,8 @@
     Data.info.flags |= (swiftVersion << 8);
   }
 
+  ~ObjCImageInfoAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeObjCImageInfo;
   }
diff --git a/lld/lib/ReaderWriter/MachO/SectCreateFile.h b/lld/lib/ReaderWriter/MachO/SectCreateFile.h
index 5236071..49e65f6 100644
--- a/lld/lib/ReaderWriter/MachO/SectCreateFile.h
+++ b/lld/lib/ReaderWriter/MachO/SectCreateFile.h
@@ -31,6 +31,8 @@
         _combinedName((segName + "/" + sectName).str()),
         _content(std::move(content)) {}
 
+    ~SectCreateAtom() override = default;
+
     uint64_t size() const override { return _content->getBufferSize(); }
 
     Scope scope() const override { return scopeGlobal; }
@@ -67,22 +69,29 @@
       new (allocator()) SectCreateAtom(*this, seg, sect, std::move(content)));
   }
 
-  const AtomVector<DefinedAtom> &defined() const override {
+  const AtomRange<DefinedAtom> defined() const override {
     return _definedAtoms;
   }
 
-  const AtomVector<UndefinedAtom> &undefined() const override {
+  const AtomRange<UndefinedAtom> undefined() const override {
     return _noUndefinedAtoms;
   }
 
-  const AtomVector<SharedLibraryAtom> &sharedLibrary() const override {
+  const AtomRange<SharedLibraryAtom> sharedLibrary() const override {
     return _noSharedLibraryAtoms;
   }
 
-  const AtomVector<AbsoluteAtom> &absolute() const override {
+  const AtomRange<AbsoluteAtom> absolute() const override {
     return _noAbsoluteAtoms;
   }
 
+  void clearAtoms() override {
+    _definedAtoms.clear();
+    _noUndefinedAtoms.clear();
+    _noSharedLibraryAtoms.clear();
+    _noAbsoluteAtoms.clear();
+  }
+
 private:
   AtomVector<DefinedAtom> _definedAtoms;
 };
diff --git a/lld/lib/ReaderWriter/MachO/StubsPass.cpp b/lld/lib/ReaderWriter/MachO/StubsPass.cpp
index e69d669..2539c26 100644
--- a/lld/lib/ReaderWriter/MachO/StubsPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/StubsPass.cpp
@@ -37,6 +37,8 @@
   LazyPointerAtom(const File &file, bool is64)
     : SimpleDefinedAtom(file), _is64(is64) { }
 
+  ~LazyPointerAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeLazyPointer;
   }
@@ -71,6 +73,8 @@
   NonLazyPointerAtom(const File &file, bool is64, ContentType contentType)
     : SimpleDefinedAtom(file), _is64(is64), _contentType(contentType) { }
 
+  ~NonLazyPointerAtom() override = default;
+
   ContentType contentType() const override {
     return _contentType;
   }
@@ -106,6 +110,8 @@
   StubAtom(const File &file, const ArchHandler::StubInfo &stubInfo)
       : SimpleDefinedAtom(file), _stubInfo(stubInfo){ }
 
+  ~StubAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeStub;
   }
@@ -138,6 +144,8 @@
   StubHelperAtom(const File &file, const ArchHandler::StubInfo &stubInfo)
       : SimpleDefinedAtom(file), _stubInfo(stubInfo) { }
 
+  ~StubHelperAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeStubHelper;
   }
@@ -171,6 +179,8 @@
   StubHelperCommonAtom(const File &file, const ArchHandler::StubInfo &stubInfo)
       : SimpleDefinedAtom(file), _stubInfo(stubInfo) { }
 
+  ~StubHelperCommonAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeStubHelper;
   }
diff --git a/lld/lib/ReaderWriter/MachO/TLVPass.cpp b/lld/lib/ReaderWriter/MachO/TLVPass.cpp
index b12f006..41aa223 100644
--- a/lld/lib/ReaderWriter/MachO/TLVPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/TLVPass.cpp
@@ -30,6 +30,8 @@
   TLVPEntryAtom(const File &file, bool is64, StringRef name)
       : SimpleDefinedAtom(file), _is64(is64), _name(name) {}
 
+  ~TLVPEntryAtom() override = default;
+
   ContentType contentType() const override {
     return DefinedAtom::typeTLVInitializerPtr;
   }
diff --git a/lld/lib/ReaderWriter/YAML/ReaderWriterYAML.cpp b/lld/lib/ReaderWriter/YAML/ReaderWriterYAML.cpp
index 091da9c..bec2c68 100644
--- a/lld/lib/ReaderWriter/YAML/ReaderWriterYAML.cpp
+++ b/lld/lib/ReaderWriter/YAML/ReaderWriterYAML.cpp
@@ -195,7 +195,7 @@
 
 /// Mapping of Atoms.
 template <typename T> class AtomList {
-  typedef lld::File::AtomVector<T> Ty;
+  using Ty = std::vector<OwningAtomPtr<T>>;
 
 public:
   typename Ty::iterator begin() { return _atoms.begin(); }
@@ -503,10 +503,20 @@
 // Declare that an AtomList is a yaml sequence.
 template <typename T> struct SequenceTraits<AtomList<T> > {
   static size_t size(IO &io, AtomList<T> &seq) { return seq._atoms.size(); }
-  static const T *&element(IO &io, AtomList<T> &seq, size_t index) {
+  static T *&element(IO &io, AtomList<T> &seq, size_t index) {
     if (index >= seq._atoms.size())
       seq._atoms.resize(index + 1);
-    return seq._atoms[index];
+    return seq._atoms[index].get();
+  }
+};
+
+// Declare that an AtomRange is a yaml sequence.
+template <typename T> struct SequenceTraits<File::AtomRange<T> > {
+  static size_t size(IO &io, File::AtomRange<T> &seq) { return seq.size(); }
+  static T *&element(IO &io, File::AtomRange<T> &seq, size_t index) {
+    assert(io.outputting() && "AtomRange only used when outputting");
+    assert(index < seq.size() && "Out of range access");
+    return seq[index].get();
   }
 };
 
@@ -558,23 +568,29 @@
 
     const lld::File *denormalize(IO &io) { return this; }
 
-    const AtomVector<lld::DefinedAtom> &defined() const override {
+    const AtomRange<lld::DefinedAtom> defined() const override {
       return _noDefinedAtoms;
     }
 
-    const AtomVector<lld::UndefinedAtom> &undefined() const override {
+    const AtomRange<lld::UndefinedAtom> undefined() const override {
       return _noUndefinedAtoms;
     }
 
-    const AtomVector<lld::SharedLibraryAtom> &
-    sharedLibrary() const override {
+    const AtomRange<lld::SharedLibraryAtom> sharedLibrary() const override {
       return _noSharedLibraryAtoms;
     }
 
-    const AtomVector<lld::AbsoluteAtom> &absolute() const override {
+    const AtomRange<lld::AbsoluteAtom> absolute() const override {
       return _noAbsoluteAtoms;
     }
 
+    void clearAtoms() override {
+      _noDefinedAtoms.clear();
+      _noUndefinedAtoms.clear();
+      _noSharedLibraryAtoms.clear();
+      _noAbsoluteAtoms.clear();
+    }
+
     File *find(StringRef name, bool dataSymbolOnly) override {
       for (const ArchMember &member : _members) {
         for (const lld::DefinedAtom *atom : member._content->defined()) {
@@ -606,36 +622,46 @@
   class NormalizedFile : public lld::File {
   public:
     NormalizedFile(IO &io)
-      : File("", kindNormalizedObject), _io(io), _rnb(nullptr) {}
+      : File("", kindNormalizedObject), _io(io), _rnb(nullptr),
+        _definedAtomsRef(_definedAtoms._atoms),
+        _undefinedAtomsRef(_undefinedAtoms._atoms),
+        _sharedLibraryAtomsRef(_sharedLibraryAtoms._atoms),
+        _absoluteAtomsRef(_absoluteAtoms._atoms) {}
     NormalizedFile(IO &io, const lld::File *file)
         : File(file->path(), kindNormalizedObject), _io(io),
-          _rnb(new RefNameBuilder(*file)), _path(file->path()) {
-      for (const lld::DefinedAtom *a : file->defined())
-        _definedAtoms._atoms.push_back(a);
-      for (const lld::UndefinedAtom *a : file->undefined())
-        _undefinedAtoms._atoms.push_back(a);
-      for (const lld::SharedLibraryAtom *a : file->sharedLibrary())
-        _sharedLibraryAtoms._atoms.push_back(a);
-      for (const lld::AbsoluteAtom *a : file->absolute())
-        _absoluteAtoms._atoms.push_back(a);
+          _rnb(new RefNameBuilder(*file)), _path(file->path()),
+        _definedAtomsRef(file->defined()),
+        _undefinedAtomsRef(file->undefined()),
+        _sharedLibraryAtomsRef(file->sharedLibrary()),
+        _absoluteAtomsRef(file->absolute()) {
     }
+
+    ~NormalizedFile() override {
+    }
+
     const lld::File *denormalize(IO &io);
 
-    const AtomVector<lld::DefinedAtom> &defined() const override {
-      return _definedAtoms._atoms;
+    const AtomRange<lld::DefinedAtom> defined() const override {
+      return _definedAtomsRef;
     }
 
-    const AtomVector<lld::UndefinedAtom> &undefined() const override {
-      return _undefinedAtoms._atoms;
+    const AtomRange<lld::UndefinedAtom> undefined() const override {
+      return _undefinedAtomsRef;
     }
 
-    const AtomVector<lld::SharedLibraryAtom> &
-    sharedLibrary() const override {
-      return _sharedLibraryAtoms._atoms;
+    const AtomRange<lld::SharedLibraryAtom> sharedLibrary() const override {
+      return _sharedLibraryAtomsRef;
     }
 
-    const AtomVector<lld::AbsoluteAtom> &absolute() const override {
-      return _absoluteAtoms._atoms;
+    const AtomRange<lld::AbsoluteAtom> absolute() const override {
+      return _absoluteAtomsRef;
+    }
+
+    void clearAtoms() override {
+      _definedAtoms._atoms.clear();
+      _undefinedAtoms._atoms.clear();
+      _sharedLibraryAtoms._atoms.clear();
+      _absoluteAtoms._atoms.clear();
     }
 
     // Allocate a new copy of this string in _storage, so the strings
@@ -653,6 +679,10 @@
     AtomList<lld::UndefinedAtom>         _undefinedAtoms;
     AtomList<lld::SharedLibraryAtom>     _sharedLibraryAtoms;
     AtomList<lld::AbsoluteAtom>          _absoluteAtoms;
+    AtomRange<lld::DefinedAtom>          _definedAtomsRef;
+    AtomRange<lld::UndefinedAtom>        _undefinedAtomsRef;
+    AtomRange<lld::SharedLibraryAtom>    _sharedLibraryAtomsRef;
+    AtomRange<lld::AbsoluteAtom>         _absoluteAtomsRef;
     llvm::BumpPtrAllocator               _storage;
   };
 
@@ -676,10 +706,18 @@
     info->_file = keys.operator->();
 
     io.mapOptional("path",                 keys->_path);
-    io.mapOptional("defined-atoms",        keys->_definedAtoms);
-    io.mapOptional("undefined-atoms",      keys->_undefinedAtoms);
-    io.mapOptional("shared-library-atoms", keys->_sharedLibraryAtoms);
-    io.mapOptional("absolute-atoms",       keys->_absoluteAtoms);
+
+    if (io.outputting()) {
+      io.mapOptional("defined-atoms",        keys->_definedAtomsRef);
+      io.mapOptional("undefined-atoms",      keys->_undefinedAtomsRef);
+      io.mapOptional("shared-library-atoms", keys->_sharedLibraryAtomsRef);
+      io.mapOptional("absolute-atoms",       keys->_absoluteAtomsRef);
+    } else {
+      io.mapOptional("defined-atoms",        keys->_definedAtoms);
+      io.mapOptional("undefined-atoms",      keys->_undefinedAtoms);
+      io.mapOptional("shared-library-atoms", keys->_sharedLibraryAtoms);
+      io.mapOptional("absolute-atoms",       keys->_absoluteAtoms);
+    }
   }
 
   static void mappingArchive(IO &io, const lld::File *&file) {
@@ -790,6 +828,9 @@
       for (uint8_t x : cont)
         _content.push_back(x);
     }
+
+    ~NormalizedAtom() override = default;
+    
     const lld::DefinedAtom *denormalize(IO &io) {
       YamlContext *info = reinterpret_cast<YamlContext *>(io.getContext());
       assert(info != nullptr);
@@ -938,6 +979,14 @@
   }
 };
 
+template <> struct MappingTraits<lld::DefinedAtom *> {
+  static void mapping(IO &io, lld::DefinedAtom *&atom) {
+    const lld::DefinedAtom *atomPtr = atom;
+    MappingTraits<const lld::DefinedAtom *>::mapping(io, atomPtr);
+    atom = const_cast<lld::DefinedAtom *>(atomPtr);
+  }
+};
+
 // YAML conversion for const lld::UndefinedAtom*
 template <> struct MappingTraits<const lld::UndefinedAtom *> {
 
@@ -950,6 +999,8 @@
         : _file(fileFromContext(io)), _name(atom->name()),
           _canBeNull(atom->canBeNull()) {}
 
+    ~NormalizedAtom() override = default;
+
     const lld::UndefinedAtom *denormalize(IO &io) {
       YamlContext *info = reinterpret_cast<YamlContext *>(io.getContext());
       assert(info != nullptr);
@@ -993,6 +1044,14 @@
   }
 };
 
+template <> struct MappingTraits<lld::UndefinedAtom *> {
+  static void mapping(IO &io, lld::UndefinedAtom *&atom) {
+    const lld::UndefinedAtom *atomPtr = atom;
+    MappingTraits<const lld::UndefinedAtom *>::mapping(io, atomPtr);
+    atom = const_cast<lld::UndefinedAtom *>(atomPtr);
+  }
+};
+
 // YAML conversion for const lld::SharedLibraryAtom*
 template <> struct MappingTraits<const lld::SharedLibraryAtom *> {
 
@@ -1006,6 +1065,8 @@
           _loadName(atom->loadName()), _canBeNull(atom->canBeNullAtRuntime()),
           _type(atom->type()), _size(atom->size()) {}
 
+    ~NormalizedAtom() override = default;
+
     const lld::SharedLibraryAtom *denormalize(IO &io) {
       YamlContext *info = reinterpret_cast<YamlContext *>(io.getContext());
       assert(info != nullptr);
@@ -1061,6 +1122,14 @@
   }
 };
 
+template <> struct MappingTraits<lld::SharedLibraryAtom *> {
+  static void mapping(IO &io, lld::SharedLibraryAtom *&atom) {
+    const lld::SharedLibraryAtom *atomPtr = atom;
+    MappingTraits<const lld::SharedLibraryAtom *>::mapping(io, atomPtr);
+    atom = const_cast<lld::SharedLibraryAtom *>(atomPtr);
+  }
+};
+
 // YAML conversion for const lld::AbsoluteAtom*
 template <> struct MappingTraits<const lld::AbsoluteAtom *> {
 
@@ -1071,6 +1140,9 @@
     NormalizedAtom(IO &io, const lld::AbsoluteAtom *atom)
         : _file(fileFromContext(io)), _name(atom->name()),
           _scope(atom->scope()), _value(atom->value()) {}
+
+    ~NormalizedAtom() override = default;
+
     const lld::AbsoluteAtom *denormalize(IO &io) {
       YamlContext *info = reinterpret_cast<YamlContext *>(io.getContext());
       assert(info != nullptr);
@@ -1129,6 +1201,14 @@
   }
 };
 
+template <> struct MappingTraits<lld::AbsoluteAtom *> {
+  static void mapping(IO &io, lld::AbsoluteAtom *&atom) {
+    const lld::AbsoluteAtom *atomPtr = atom;
+    MappingTraits<const lld::AbsoluteAtom *>::mapping(io, atomPtr);
+    atom = const_cast<lld::AbsoluteAtom *>(atomPtr);
+  }
+};
+
 } // namespace llvm
 } // namespace yaml