1 files changed, 50 insertions, 6 deletions
diff --git a/lib/ObjectYAML/DWARFEmitter.cpp b/lib/ObjectYAML/DWARFEmitter.cpp
index c49c2249cdfe..f23fa1237600 100644
--- a/lib/ObjectYAML/DWARFEmitter.cpp
+++ b/lib/ObjectYAML/DWARFEmitter.cpp
@@ -8,7 +8,7 @@
 //===----------------------------------------------------------------------===//
 ///
 /// \file
-/// \brief The DWARF component of yaml2obj. Provided as library code for tests.
+/// The DWARF component of yaml2obj. Provided as library code for tests.
 ///
 //===----------------------------------------------------------------------===//
 
@@ -132,7 +132,7 @@ void DWARFYAML::EmitPubSection(raw_ostream &OS,
 }
 
 namespace {
-/// \brief An extension of the DWARFYAML::ConstVisitor which writes compile
+/// An extension of the DWARFYAML::ConstVisitor which writes compile
 /// units and DIEs to a stream.
 class DumpVisitor : public DWARFYAML::ConstVisitor {
   raw_ostream &OS;
@@ -149,7 +149,6 @@ protected:
       writeInteger((uint32_t)CU.AbbrOffset, OS, DebugInfo.IsLittleEndian);
       writeInteger((uint8_t)CU.AddrSize, OS, DebugInfo.IsLittleEndian);
     }
-    
   }
 
   void onStartDIE(const DWARFYAML::Unit &CU,
@@ -308,11 +307,50 @@ EmitDebugSectionImpl(const DWARFYAML::Data &DI, EmitFuncType EmitFunc,
     OutputBuffers[Sec] = MemoryBuffer::getMemBufferCopy(Data);
 }
 
+namespace {
+class DIEFixupVisitor : public DWARFYAML::Visitor {
+  uint64_t Length;
+
+public:
+  DIEFixupVisitor(DWARFYAML::Data &DI) : DWARFYAML::Visitor(DI){};
+
+private:
+  virtual void onStartCompileUnit(DWARFYAML::Unit &CU) { Length = 7; }
+
+  virtual void onEndCompileUnit(DWARFYAML::Unit &CU) {
+    CU.Length.setLength(Length);
+  }
+
+  virtual void onStartDIE(DWARFYAML::Unit &CU, DWARFYAML::Entry &DIE) {
+    Length += getULEB128Size(DIE.AbbrCode);
+  }
+
+  virtual void onValue(const uint8_t U) { Length += 1; }
+  virtual void onValue(const uint16_t U) { Length += 2; }
+  virtual void onValue(const uint32_t U) { Length += 4; }
+  virtual void onValue(const uint64_t U, const bool LEB = false) {
+    if (LEB)
+      Length += getULEB128Size(U);
+    else
+      Length += 8;
+  }
+  virtual void onValue(const int64_t S, const bool LEB = false) {
+    if (LEB)
+      Length += getSLEB128Size(S);
+    else
+      Length += 8;
+  }
+  virtual void onValue(const StringRef String) { Length += String.size() + 1; }
+
+  virtual void onValue(const MemoryBufferRef MBR) {
+    Length += MBR.getBufferSize();
+  }
+};
+} // namespace
+
 Expected<StringMap<std::unique_ptr<MemoryBuffer>>>
-DWARFYAML::EmitDebugSections(StringRef YAMLString,
+DWARFYAML::EmitDebugSections(StringRef YAMLString, bool ApplyFixups,
                              bool IsLittleEndian) {
-  StringMap<std::unique_ptr<MemoryBuffer>> DebugSections;
-
   yaml::Input YIn(YAMLString);
 
   DWARFYAML::Data DI;
@@ -321,6 +359,12 @@ DWARFYAML::EmitDebugSections(StringRef YAMLString,
   if (YIn.error())
     return errorCodeToError(YIn.error());
 
+  if (ApplyFixups) {
+    DIEFixupVisitor DIFixer(DI);
+    DIFixer.traverseDebugInfo();
+  }
+
+  StringMap<std::unique_ptr<MemoryBuffer>> DebugSections;
   EmitDebugSectionImpl(DI, &DWARFYAML::EmitDebugInfo, "debug_info",
                        DebugSections);
   EmitDebugSectionImpl(DI, &DWARFYAML::EmitDebugLine, "debug_line",