1 files changed, 28 insertions, 21 deletions

diff --git a/bolt/lib/Passes/RetpolineInsertion.cpp b/bolt/lib/Passes/RetpolineInsertion.cpp
index 4d83fb3903a7..c8a42725c974 100644
--- a/bolt/lib/Passes/RetpolineInsertion.cpp
+++ b/bolt/lib/Passes/RetpolineInsertion.cpp
@@ -22,6 +22,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/RetpolineInsertion.h"
+#include "llvm/MC/MCInstPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "bolt-retpoline"
@@ -173,39 +174,45 @@ BinaryFunction *createNewRetpoline(BinaryContext &BC,
 std::string createRetpolineFunctionTag(BinaryContext &BC,
                                        const IndirectBranchInfo &BrInfo,
                                        bool R11Available) {
-  if (BrInfo.isReg())
-    return "__retpoline_r" + to_string(BrInfo.BranchReg) + "_";
+  std::string Tag;
+  llvm::raw_string_ostream TagOS(Tag);
+  TagOS << "__retpoline_";
+
+  if (BrInfo.isReg()) {
+    BC.InstPrinter->printRegName(TagOS, BrInfo.BranchReg);
+    TagOS << "_";
+    TagOS.flush();
+    return Tag;
+  }
 
   // Memory Branch
   if (R11Available)
     return "__retpoline_r11";
 
-  std::string Tag = "__retpoline_mem_";
-
   const IndirectBranchInfo::MemOpInfo &MemRef = BrInfo.Memory;
 
-  std::string DispExprStr;
-  if (MemRef.DispExpr) {
-    llvm::raw_string_ostream Ostream(DispExprStr);
-    MemRef.DispExpr->print(Ostream, BC.AsmInfo.get());
-    Ostream.flush();
-  }
+  TagOS << "mem_";
 
-  Tag += MemRef.BaseRegNum != BC.MIB->getNoRegister()
-             ? "r" + to_string(MemRef.BaseRegNum)
-             : "";
+  if (MemRef.BaseRegNum != BC.MIB->getNoRegister())
+    BC.InstPrinter->printRegName(TagOS, MemRef.BaseRegNum);
 
-  Tag += MemRef.DispExpr ? "+" + DispExprStr : "+" + to_string(MemRef.DispImm);
+  TagOS << "+";
+  if (MemRef.DispExpr)
+    MemRef.DispExpr->print(TagOS, BC.AsmInfo.get());
+  else
+    TagOS << MemRef.DispImm;
 
-  Tag += MemRef.IndexRegNum != BC.MIB->getNoRegister()
-             ? "+" + to_string(MemRef.ScaleImm) + "*" +
-                   to_string(MemRef.IndexRegNum)
-             : "";
+  if (MemRef.IndexRegNum != BC.MIB->getNoRegister()) {
+    TagOS << "+" << MemRef.ScaleImm << "*";
+    BC.InstPrinter->printRegName(TagOS, MemRef.IndexRegNum);
+  }
 
-  Tag += MemRef.SegRegNum != BC.MIB->getNoRegister()
-             ? "_seg_" + to_string(MemRef.SegRegNum)
-             : "";
+  if (MemRef.SegRegNum != BC.MIB->getNoRegister()) {
+    TagOS << "_seg_";
+    BC.InstPrinter->printRegName(TagOS, MemRef.SegRegNum);
+  }
 
+  TagOS.flush();
   return Tag;
 }