Index: lib/DebugInfo/DWARFCompileUnit.h
===================================================================
--- lib/DebugInfo/DWARFCompileUnit.h
+++ lib/DebugInfo/DWARFCompileUnit.h
@@ -18,8 +18,9 @@
 public:
   DWARFCompileUnit(DWARFContext& Context, const DWARFDebugAbbrev *DA,
                    StringRef IS, StringRef RS, StringRef SS, StringRef SOS,
-                   StringRef AOS, const RelocAddrMap *M, bool LE)
-    : DWARFUnit(Context, DA, IS, RS, SS, SOS, AOS, M, LE) {}
+                   StringRef AOS, const RelocAddrMap *M, bool LE,
+                   const DWARFUnitSectionBase &UnitSection)
+    : DWARFUnit(Context, DA, IS, RS, SS, SOS, AOS, M, LE, UnitSection) {}
   void dump(raw_ostream &OS);
   // VTable anchor.
   ~DWARFCompileUnit() override;
Index: lib/DebugInfo/DWARFContext.h
===================================================================
--- lib/DebugInfo/DWARFContext.h
+++ lib/DebugInfo/DWARFContext.h
@@ -28,19 +28,17 @@
 /// information parsing. The actual data is supplied through pure virtual
 /// methods that a concrete implementation provides.
 class DWARFContext : public DIContext {
-  typedef SmallVector<std::unique_ptr<DWARFCompileUnit>, 1> CUVector;
-  typedef SmallVector<std::unique_ptr<DWARFTypeUnit>, 1> TUVector;
 
-  CUVector CUs;
-  TUVector TUs;
+  DWARFUnitSection<DWARFCompileUnit> CUs;
+  DWARFUnitSection<DWARFTypeUnit> TUs;
   std::unique_ptr<DWARFDebugAbbrev> Abbrev;
   std::unique_ptr<DWARFDebugLoc> Loc;
   std::unique_ptr<DWARFDebugAranges> Aranges;
   std::unique_ptr<DWARFDebugLine> Line;
   std::unique_ptr<DWARFDebugFrame> DebugFrame;
 
-  CUVector DWOCUs;
-  TUVector DWOTUs;
+  DWARFUnitSection<DWARFCompileUnit> DWOCUs;
+  DWARFUnitSection<DWARFTypeUnit> DWOTUs;
   std::unique_ptr<DWARFDebugAbbrev> AbbrevDWO;
   std::unique_ptr<DWARFDebugLocDWO> LocDWO;
 
@@ -77,8 +75,8 @@
 
   void dump(raw_ostream &OS, DIDumpType DumpType = DIDT_All) override;
 
-  typedef iterator_range<CUVector::iterator> cu_iterator_range;
-  typedef iterator_range<TUVector::iterator> tu_iterator_range;
+  typedef DWARFUnitSection<DWARFCompileUnit>::iterator_range cu_iterator_range;
+  typedef DWARFUnitSection<DWARFTypeUnit>::iterator_range tu_iterator_range;
 
   /// Get compile units in this context.
   cu_iterator_range compile_units() {
Index: lib/DebugInfo/DWARFContext.cpp
===================================================================
--- lib/DebugInfo/DWARFContext.cpp
+++ lib/DebugInfo/DWARFContext.cpp
@@ -321,7 +321,7 @@
     std::unique_ptr<DWARFCompileUnit> CU(new DWARFCompileUnit(*this,
         getDebugAbbrev(), getInfoSection().Data, getRangeSection(),
         getStringSection(), StringRef(), getAddrSection(),
-        &getInfoSection().Relocs, isLittleEndian()));
+        &getInfoSection().Relocs, isLittleEndian(), CUs));
     if (!CU->extract(DIData, &offset)) {
       break;
     }
@@ -341,7 +341,7 @@
       std::unique_ptr<DWARFTypeUnit> TU(new DWARFTypeUnit(*this,
            getDebugAbbrev(), I.second.Data, getRangeSection(),
            getStringSection(), StringRef(), getAddrSection(),
-           &I.second.Relocs, isLittleEndian()));
+           &I.second.Relocs, isLittleEndian(), TUs));
       if (!TU->extract(DIData, &offset))
         break;
       TUs.push_back(std::move(TU));
@@ -360,7 +360,7 @@
     std::unique_ptr<DWARFCompileUnit> DWOCU(new DWARFCompileUnit(*this,
         getDebugAbbrevDWO(), getInfoDWOSection().Data, getRangeDWOSection(),
         getStringDWOSection(), getStringOffsetDWOSection(), getAddrSection(),
-        &getInfoDWOSection().Relocs, isLittleEndian()));
+        &getInfoDWOSection().Relocs, isLittleEndian(), DWOCUs));
     if (!DWOCU->extract(DIData, &offset)) {
       break;
     }
@@ -380,7 +380,7 @@
       std::unique_ptr<DWARFTypeUnit> TU(new DWARFTypeUnit(*this,
           getDebugAbbrevDWO(), I.second.Data, getRangeDWOSection(),
           getStringDWOSection(), getStringOffsetDWOSection(), getAddrSection(),
-          &I.second.Relocs, isLittleEndian()));
+          &I.second.Relocs, isLittleEndian(), DWOTUs));
       if (!TU->extract(DIData, &offset))
         break;
       DWOTUs.push_back(std::move(TU));
@@ -389,33 +389,9 @@
   }
 }
 
-namespace {
-  struct OffsetComparator {
-
-    bool operator()(const std::unique_ptr<DWARFCompileUnit> &LHS,
-                    const std::unique_ptr<DWARFCompileUnit> &RHS) const {
-      return LHS->getOffset() < RHS->getOffset();
-    }
-    bool operator()(const std::unique_ptr<DWARFCompileUnit> &LHS,
-                    uint32_t RHS) const {
-      return LHS->getOffset() < RHS;
-    }
-    bool operator()(uint32_t LHS,
-                    const std::unique_ptr<DWARFCompileUnit> &RHS) const {
-      return LHS < RHS->getOffset();
-    }
-  };
-}
-
 DWARFCompileUnit *DWARFContext::getCompileUnitForOffset(uint32_t Offset) {
   parseCompileUnits();
-
-  std::unique_ptr<DWARFCompileUnit> *CU =
-      std::lower_bound(CUs.begin(), CUs.end(), Offset, OffsetComparator());
-  if (CU != CUs.end()) {
-    return CU->get();
-  }
-  return nullptr;
+  return CUs.getUnitForOffset(Offset);
 }
 
 DWARFCompileUnit *DWARFContext::getCompileUnitForAddress(uint64_t Address) {
Index: lib/DebugInfo/DWARFTypeUnit.h
===================================================================
--- lib/DebugInfo/DWARFTypeUnit.h
+++ lib/DebugInfo/DWARFTypeUnit.h
@@ -21,8 +21,9 @@
 public:
   DWARFTypeUnit(DWARFContext &Context, const DWARFDebugAbbrev *DA,
                 StringRef IS, StringRef RS, StringRef SS, StringRef SOS,
-                StringRef AOS, const RelocAddrMap *M, bool LE)
-    : DWARFUnit(Context, DA, IS, RS, SS, SOS, AOS, M, LE) {}
+                StringRef AOS, const RelocAddrMap *M, bool LE,
+                const DWARFUnitSectionBase &UnitSection)
+    : DWARFUnit(Context, DA, IS, RS, SS, SOS, AOS, M, LE, UnitSection) {}
   uint32_t getHeaderSize() const override {
     return DWARFUnit::getHeaderSize() + 12;
   }
Index: lib/DebugInfo/DWARFUnit.h
===================================================================
--- lib/DebugInfo/DWARFUnit.h
+++ lib/DebugInfo/DWARFUnit.h
@@ -26,6 +26,50 @@
 class DWARFDebugAbbrev;
 class StringRef;
 class raw_ostream;
+class DWARFUnit;
+
+/// Base class for all DWARFUnitSection classes. This provides the
+/// functionality common to all unit types.
+class DWARFUnitSectionBase {
+public:
+  /// Returns the Unit that contains the given section offset in the
+  /// same section this Unit originated from.
+  virtual DWARFUnit *getUnitForOffset(uint32_t Offset) const = 0;
+};
+
+/// Concrete instance of DWARFUnitSection, specialized for one Unit type.
+template<typename UnitType>
+class DWARFUnitSection : public SmallVector<std::unique_ptr<UnitType>, 1>,
+                         public DWARFUnitSectionBase {
+
+  struct UnitOffsetComparator {
+    bool operator()(const std::unique_ptr<UnitType> &LHS,
+                    const std::unique_ptr<UnitType> &RHS) const {
+      return LHS->getOffset() < RHS->getOffset();
+    }
+    bool operator()(const std::unique_ptr<UnitType> &LHS,
+                    uint32_t RHS) const {
+      return LHS->getOffset() < RHS;
+    }
+    bool operator()(uint32_t LHS,
+                    const std::unique_ptr<UnitType> &RHS) const {
+      return LHS < RHS->getOffset();
+    }
+  };
+
+public:
+  typedef SmallVectorImpl<std::unique_ptr<UnitType>> UnitVector;
+  typedef typename UnitVector::iterator iterator;
+  typedef iterator_range<typename UnitVector::iterator> iterator_range;
+
+  UnitType *getUnitForOffset(uint32_t Offset) const {
+    auto *CU = std::lower_bound(this->begin(), this->end(), Offset,
+                                UnitOffsetComparator());
+    if (CU != this->end())
+      return CU->get();
+    return nullptr;
+  }
+};
 
 class DWARFUnit {
   DWARFContext &Context;
@@ -40,6 +84,7 @@
   uint32_t AddrOffsetSectionBase;
   const RelocAddrMap *RelocMap;
   bool isLittleEndian;
+  const DWARFUnitSectionBase &UnitSection;
 
   uint32_t Offset;
   uint32_t Length;
@@ -68,7 +113,7 @@
 public:
   DWARFUnit(DWARFContext& Context, const DWARFDebugAbbrev *DA, StringRef IS,
             StringRef RS, StringRef SS, StringRef SOS, StringRef AOS,
-            const RelocAddrMap *M, bool LE);
+            const RelocAddrMap *M, bool LE, const DWARFUnitSectionBase &UnitSection);
 
   virtual ~DWARFUnit();
 
@@ -136,6 +181,9 @@
   /// chain is valid as long as parsed compile unit DIEs are not cleared.
   DWARFDebugInfoEntryInlinedChain getInlinedChainForAddress(uint64_t Address);
 
+  /// getUnitSection - Return the DWARFUnitSection containing this unit.
+  const DWARFUnitSectionBase &getUnitSection() const { return UnitSection; }
+
 private:
   /// Size in bytes of the .debug_info data associated with this compile unit.
   size_t getDebugInfoSize() const { return Length + 4 - getHeaderSize(); }
Index: lib/DebugInfo/DWARFUnit.cpp
===================================================================
--- lib/DebugInfo/DWARFUnit.cpp
+++ lib/DebugInfo/DWARFUnit.cpp
@@ -19,10 +19,11 @@
 
 DWARFUnit::DWARFUnit(DWARFContext &DC, const DWARFDebugAbbrev *DA,
                      StringRef IS, StringRef RS, StringRef SS, StringRef SOS,
-                     StringRef AOS, const RelocAddrMap *M, bool LE)
+                     StringRef AOS, const RelocAddrMap *M, bool LE,
+                     const DWARFUnitSectionBase& UnitSection)
   : Context(DC), Abbrev(DA), InfoSection(IS), RangeSection(RS),
     StringSection(SS), StringOffsetSection(SOS), AddrOffsetSection(AOS),
-    RelocMap(M), isLittleEndian(LE) {
+    RelocMap(M), isLittleEndian(LE), UnitSection(UnitSection)  {
   clear();
 }