2 files changed, 46 insertions, 9 deletions

diff --git a/llvm/lib/ObjectYAML/DWARFEmitter.cpp b/llvm/lib/ObjectYAML/DWARFEmitter.cpp
index 682969a..ad4b316 100644
--- a/llvm/lib/ObjectYAML/DWARFEmitter.cpp
+++ b/llvm/lib/ObjectYAML/DWARFEmitter.cpp
@@ -246,7 +246,8 @@ Error DWARFYAML::emitDebugGNUPubtypes(raw_ostream &OS, const Data &DI) {
                         /*IsGNUStyle=*/true);
 }
 
-static Expected<uint64_t> writeDIE(ArrayRef<DWARFYAML::AbbrevTable> AbbrevTable,
+static Expected<uint64_t> writeDIE(const DWARFYAML::Data &DI, uint64_t CUIndex,
+                                   uint64_t AbbrevTableID,
                                    const dwarf::FormParams &Params,
                                    const DWARFYAML::Entry &Entry,
                                    raw_ostream &OS, bool IsLittleEndian) {
@@ -256,12 +257,16 @@ static Expected<uint64_t> writeDIE(ArrayRef<DWARFYAML::AbbrevTable> AbbrevTable,
   if (AbbrCode == 0 || Entry.Values.empty())
     return OS.tell() - EntryBegin;
 
-  if (AbbrevTable.empty())
-    return createStringError(
-        errc::invalid_argument,
-        "non-empty compilation unit should have an associated abbrev table");
+  Expected<uint64_t> AbbrevTableIndexOrErr =
+      DI.getAbbrevTableIndexByID(AbbrevTableID);
+  if (!AbbrevTableIndexOrErr)
+    return createStringError(errc::invalid_argument,
+                             toString(AbbrevTableIndexOrErr.takeError()) +
+                                 " for compilation unit with index " +
+                                 utostr(CUIndex));
 
-  ArrayRef<DWARFYAML::Abbrev> AbbrevDecls(AbbrevTable[0].Table);
+  ArrayRef<DWARFYAML::Abbrev> AbbrevDecls(
+      DI.DebugAbbrev[*AbbrevTableIndexOrErr].Table);
 
   if (AbbrCode > AbbrevDecls.size())
     return createStringError(
@@ -384,7 +389,8 @@ static Expected<uint64_t> writeDIE(ArrayRef<DWARFYAML::AbbrevTable> AbbrevTable,
 }
 
 Error DWARFYAML::emitDebugInfo(raw_ostream &OS, const DWARFYAML::Data &DI) {
-  for (const DWARFYAML::Unit &Unit : DI.CompileUnits) {
+  for (uint64_t I = 0; I < DI.CompileUnits.size(); ++I) {
+    const DWARFYAML::Unit &Unit = DI.CompileUnits[I];
     uint8_t AddrSize;
     if (Unit.AddrSize)
       AddrSize = *Unit.AddrSize;
@@ -402,9 +408,11 @@ Error DWARFYAML::emitDebugInfo(raw_ostream &OS, const DWARFYAML::Data &DI) {
     std::string EntryBuffer;
     raw_string_ostream EntryBufferOS(EntryBuffer);
 
+    uint64_t AbbrevTableID = Unit.AbbrevTableID.getValueOr(I);
     for (const DWARFYAML::Entry &Entry : Unit.Entries) {
-      if (Expected<uint64_t> EntryLength = writeDIE(
-              DI.DebugAbbrev, Params, Entry, EntryBufferOS, DI.IsLittleEndian))
+      if (Expected<uint64_t> EntryLength =
+              writeDIE(DI, I, AbbrevTableID, Params, Entry, EntryBufferOS,
+                       DI.IsLittleEndian))
         Length += *EntryLength;
       else
         return EntryLength.takeError();
diff --git a/llvm/lib/ObjectYAML/DWARFYAML.cpp b/llvm/lib/ObjectYAML/DWARFYAML.cpp
index 5c4fd53..c55ae9e1 100644
--- a/llvm/lib/ObjectYAML/DWARFYAML.cpp
+++ b/llvm/lib/ObjectYAML/DWARFYAML.cpp
@@ -13,6 +13,8 @@
 
 #include "llvm/ObjectYAML/DWARFYAML.h"
 #include "llvm/BinaryFormat/Dwarf.h"
+#include "llvm/Support/Errc.h"
+#include "llvm/Support/Error.h"
 
 namespace llvm {
 
@@ -53,6 +55,31 @@ SetVector<StringRef> DWARFYAML::Data::getNonEmptySectionNames() const {
   return SecNames;
 }
 
+Expected<uint64_t> DWARFYAML::Data::getAbbrevTableIndexByID(uint64_t ID) const {
+  if (AbbrevTableID2Index.empty()) {
+    for (auto &AbbrevTable : enumerate(DebugAbbrev)) {
+      // If the abbrev table's ID isn't specified, we use the index as its ID.
+      uint64_t AbbrevTableID =
+          AbbrevTable.value().ID.getValueOr(AbbrevTable.index());
+      auto It =
+          AbbrevTableID2Index.insert({AbbrevTableID, AbbrevTable.index()});
+      if (!It.second)
+        return createStringError(
+            errc::invalid_argument,
+            "the ID (%" PRIu64 ") of abbrev table with index %zu has been used "
+            "by abbrev table with index %" PRIu64,
+            AbbrevTableID, AbbrevTable.index(), It.first->second);
+    }
+  }
+
+  auto It = AbbrevTableID2Index.find(ID);
+  if (It == AbbrevTableID2Index.end())
+    return createStringError(errc::invalid_argument,
+                             "cannot find abbrev table whose ID is %" PRIu64,
+                             ID);
+  return It->second;
+}
+
 namespace yaml {
 
 void MappingTraits<DWARFYAML::Data>::mapping(IO &IO, DWARFYAML::Data &DWARF) {
@@ -80,6 +107,7 @@ void MappingTraits<DWARFYAML::Data>::mapping(IO &IO, DWARFYAML::Data &DWARF) {
 
 void MappingTraits<DWARFYAML::AbbrevTable>::mapping(
     IO &IO, DWARFYAML::AbbrevTable &AbbrevTable) {
+  IO.mapOptional("ID", AbbrevTable.ID);
   IO.mapOptional("Table", AbbrevTable.Table);
 }
 
@@ -153,6 +181,7 @@ void MappingTraits<DWARFYAML::Unit>::mapping(IO &IO, DWARFYAML::Unit &Unit) {
   IO.mapRequired("Version", Unit.Version);
   if (Unit.Version >= 5)
     IO.mapRequired("UnitType", Unit.Type);
+  IO.mapOptional("AbbrevTableID", Unit.AbbrevTableID);
   IO.mapRequired("AbbrOffset", Unit.AbbrOffset);
   IO.mapOptional("AddrSize", Unit.AddrSize);
   IO.mapOptional("Entries", Unit.Entries);