1 files changed, 22 insertions, 12 deletions
diff --git a/include/llvm/MC/MCMachObjectWriter.h b/include/llvm/MC/MCMachObjectWriter.h
index 12a7f0ee7bb5a..e4681c0a3dc37 100644
--- a/include/llvm/MC/MCMachObjectWriter.h
+++ b/include/llvm/MC/MCMachObjectWriter.h
@@ -14,6 +14,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCObjectWriter.h"
+#include "llvm/MC/StringTableBuilder.h"
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/MachO.h"
 #include <vector>
@@ -67,12 +68,10 @@ public:
   /// @name API
   /// @{
 
-  virtual void RecordRelocation(MachObjectWriter *Writer,
-                                const MCAssembler &Asm,
+  virtual void RecordRelocation(MachObjectWriter *Writer, MCAssembler &Asm,
                                 const MCAsmLayout &Layout,
                                 const MCFragment *Fragment,
-                                const MCFixup &Fixup,
-                                MCValue Target,
+                                const MCFixup &Fixup, MCValue Target,
                                 uint64_t &FixedValue) = 0;
 
   /// @}
@@ -96,15 +95,21 @@ class MachObjectWriter : public MCObjectWriter {
   /// @name Relocation Data
   /// @{
 
-  llvm::DenseMap<const MCSectionData*,
-                 std::vector<MachO::any_relocation_info> > Relocations;
+  struct RelAndSymbol {
+    const MCSymbolData *Sym;
+    MachO::any_relocation_info MRE;
+    RelAndSymbol(const MCSymbolData *Sym, const MachO::any_relocation_info &MRE)
+        : Sym(Sym), MRE(MRE) {}
+  };
+
+  llvm::DenseMap<const MCSectionData *, std::vector<RelAndSymbol>> Relocations;
   llvm::DenseMap<const MCSectionData*, unsigned> IndirectSymBase;
 
   /// @}
   /// @name Symbol Table Data
   /// @{
 
-  SmallString<256> StringTable;
+  StringTableBuilder StringTable;
   std::vector<MachSymbolData> LocalSymbolData;
   std::vector<MachSymbolData> ExternalSymbolData;
   std::vector<MachSymbolData> UndefinedSymbolData;
@@ -212,9 +217,15 @@ public:
   //  - Input errors, where something cannot be correctly encoded. 'as' allows
   //    these through in many cases.
 
-  void addRelocation(const MCSectionData *SD,
+  // Add a relocation to be output in the object file. At the time this is
+  // called, the symbol indexes are not know, so if the relocation refers
+  // to a symbol it should be passed as \p RelSymbol so that it can be updated
+  // afterwards. If the relocation doesn't refer to a symbol, nullptr should be
+  // used.
+  void addRelocation(const MCSymbolData *RelSymbol, const MCSectionData *SD,
                      MachO::any_relocation_info &MRE) {
-    Relocations[SD].push_back(MRE);
+    RelAndSymbol P(RelSymbol, MRE);
+    Relocations[SD].push_back(P);
   }
 
   void RecordScatteredRelocation(const MCAssembler &Asm,
@@ -230,7 +241,7 @@ public:
                             const MCFixup &Fixup, MCValue Target,
                             uint64_t &FixedValue);
 
-  void RecordRelocation(const MCAssembler &Asm, const MCAsmLayout &Layout,
+  void RecordRelocation(MCAssembler &Asm, const MCAsmLayout &Layout,
                         const MCFragment *Fragment, const MCFixup &Fixup,
                         MCValue Target, bool &IsPCRel,
                         uint64_t &FixedValue) override;
@@ -239,8 +250,7 @@ public:
 
   /// ComputeSymbolTable - Compute the symbol table data
   ///
-  /// \param StringTable [out] - The string table data.
-  void ComputeSymbolTable(MCAssembler &Asm, SmallString<256> &StringTable,
+  void ComputeSymbolTable(MCAssembler &Asm,
                           std::vector<MachSymbolData> &LocalSymbolData,
                           std::vector<MachSymbolData> &ExternalSymbolData,
                           std::vector<MachSymbolData> &UndefinedSymbolData);