3 files changed, 29 insertions, 3 deletions
diff --git a/llvm/include/llvm/DebugInfo/DWARF/DWARFUnitIndex.h b/llvm/include/llvm/DebugInfo/DWARF/DWARFUnitIndex.h
index 8e2ce023695..8a5e9af84fa 100644
--- a/llvm/include/llvm/DebugInfo/DWARF/DWARFUnitIndex.h
+++ b/llvm/include/llvm/DebugInfo/DWARF/DWARFUnitIndex.h
@@ -83,9 +83,12 @@ public:
   DWARFUnitIndex(DWARFSectionKind InfoColumnKind)
       : InfoColumnKind(InfoColumnKind) {}
 
+  explicit operator bool() const { return Header.NumBuckets; }
+
   bool parse(DataExtractor IndexData);
   void dump(raw_ostream &OS) const;
   const Entry *getFromOffset(uint32_t Offset) const;
+  const Entry *getFromHash(uint64_t Offset) const;
 
   ArrayRef<DWARFSectionKind> getColumnKinds() const {
     return makeArrayRef(ColumnKinds.get(), Header.NumColumns);
diff --git a/llvm/lib/DebugInfo/DWARF/DWARFContext.cpp b/llvm/lib/DebugInfo/DWARF/DWARFContext.cpp
index aeb1dea2bca..d8121912a30 100644
--- a/llvm/lib/DebugInfo/DWARF/DWARFContext.cpp
+++ b/llvm/lib/DebugInfo/DWARF/DWARFContext.cpp
@@ -400,8 +400,17 @@ void DWARFContext::dump(raw_ostream &OS, DIDumpOptions DumpOpts) {
 }
 
 DWARFCompileUnit *DWARFContext::getDWOCompileUnitForHash(uint64_t Hash) {
-  // FIXME: Improve this for the case where this DWO file is really a DWP file
-  // with an index - use the index for lookup instead of a linear search.
+  if (const auto &CUI = getCUIndex()) {
+    if (const auto *R = CUI.getFromHash(Hash))
+      if (auto CUOff = R->getOffset(DW_SECT_INFO))
+        return CUs.getUnitForOffset(CUOff->Offset);
+    return nullptr;
+  }
+
+  // If there's no index, just search through the CUs in the DWO - there's
+  // probably only one unless this is something like LTO - though an in-process
+  // built/cached lookup table could be used in that case to improve repeated
+  // lookups of different CUs in the DWO.
   for (const auto &DWOCU : dwo_compile_units())
     if (DWOCU->getDWOId() == Hash)
       return DWOCU.get();
diff --git a/llvm/lib/DebugInfo/DWARF/DWARFUnitIndex.cpp b/llvm/lib/DebugInfo/DWARF/DWARFUnitIndex.cpp
index 59b3d0ca55a..6e8cbcad491 100644
--- a/llvm/lib/DebugInfo/DWARF/DWARFUnitIndex.cpp
+++ b/llvm/lib/DebugInfo/DWARF/DWARFUnitIndex.cpp
@@ -123,7 +123,7 @@ StringRef DWARFUnitIndex::getColumnHeader(DWARFSectionKind DS) {
 }
 
 void DWARFUnitIndex::dump(raw_ostream &OS) const {
-  if (!Header.NumBuckets)
+  if (!*this)
     return;
 
   Header.dump(OS);
@@ -170,3 +170,17 @@ DWARFUnitIndex::getFromOffset(uint32_t Offset) const {
         return &Rows[i];
   return nullptr;
 }
+
+const DWARFUnitIndex::Entry *DWARFUnitIndex::getFromHash(uint64_t S) const {
+  uint64_t Mask = Header.NumBuckets - 1;
+
+  auto H = S & Mask;
+  auto HP = ((S >> 32) & Mask) | 1;
+  while (Rows[H].getSignature() != S && Rows[H].getSignature() != 0)
+    H = (H + HP) & Mask;
+
+  if (Rows[H].getSignature() != S)
+    return nullptr;
+
+  return &Rows[H];
+}