[llvm] Add assembly color highlighting

Add support for syntax highlighting assembly. The patch introduces new
RAII helper called WithMarkup that takes care of both emitting colors
and markup annotations. It makes adding markup easier and ensures colors
and annotations remain consistent.

This patch adopts the new helper in the AArch64 backend. If your backend
already uses markup annotations, adoption is as easy as using the new
MCInstPrinter::markup overload.

Differential revision: https://reviews.llvm.org/D159162

llvm/docs/CommandGuide/llvm-mc.rst

@@ -75,6 +75,10 @@ End-user Options
 
  Marked up disassembly of string of hex bytes.
 
+.. option:: --cdis
+
+ Colored disassembly of string of hex bytes.
+
 .. option:: --filetype=[asm,null,obj]
 
  Sets the output filetype. Setting this flag to `asm` will make the tool

llvm/include/llvm/MC/MCInstPrinter.h

@@ -55,6 +55,9 @@ protected:
   /// True if we are printing marked up assembly.
   bool UseMarkup = false;
 
+  /// True if we are printing colored assembly.
+  bool UseColor = false;
+
   /// True if we prefer aliases (e.g. nop) to raw mnemonics.
   bool PrintAliases = true;
 
@@ -85,6 +88,35 @@ public:
 
   virtual ~MCInstPrinter();
 
+  enum class Markup {
+    Immediate,
+    Register,
+    Target,
+    Memory,
+  };
+
+  class WithMarkup {
+  public:
+    [[nodiscard]] WithMarkup(raw_ostream &OS, Markup M, bool EnableMarkup,
+                             bool EnableColor);
+    ~WithMarkup();
+
+    template <typename T> WithMarkup &operator<<(T &O) {
+      OS << O;
+      return *this;
+    }
+
+    template <typename T> WithMarkup &operator<<(const T &O) {
+      OS << O;
+      return *this;
+    }
+
+  private:
+    raw_ostream &OS;
+    bool EnableMarkup;
+    bool EnableColor;
+  };
+
   /// Customize the printer according to a command line option.
   /// @return true if the option is recognized and applied.
   virtual bool applyTargetSpecificCLOption(StringRef Opt) { return false; }
@@ -116,6 +148,11 @@ public:
   bool getUseMarkup() const { return UseMarkup; }
   void setUseMarkup(bool Value) { UseMarkup = Value; }
 
+  bool getUseColor() const { return UseColor; }
+  void setUseColor(bool Value) { UseColor = Value; }
+
+  WithMarkup markup(raw_ostream &OS, Markup M) const;
+
   /// Utility functions to make adding mark ups simpler.
   StringRef markup(StringRef s) const;
 

llvm/lib/MC/MCInstPrinter.cpp

@@ -231,3 +231,53 @@ format_object<uint64_t> MCInstPrinter::formatHex(uint64_t Value) const {
   }
   llvm_unreachable("unsupported print style");
 }
+
+MCInstPrinter::WithMarkup MCInstPrinter::markup(raw_ostream &OS,
+                                                Markup S) const {
+  return WithMarkup(OS, S, getUseMarkup(), getUseColor());
+}
+
+MCInstPrinter::WithMarkup::WithMarkup(raw_ostream &OS, Markup M,
+                                      bool EnableMarkup, bool EnableColor)
+    : OS(OS), EnableMarkup(EnableMarkup), EnableColor(EnableColor) {
+  if (EnableColor) {
+    switch (M) {
+    case Markup::Immediate:
+      OS.changeColor(raw_ostream::RED);
+      break;
+    case Markup::Register:
+      OS.changeColor(raw_ostream::CYAN);
+      break;
+    case Markup::Target:
+      OS.changeColor(raw_ostream::YELLOW);
+      break;
+    case Markup::Memory:
+      OS.changeColor(raw_ostream::GREEN);
+      break;
+    }
+  }
+
+  if (EnableMarkup) {
+    switch (M) {
+    case Markup::Immediate:
+      OS << "<imm:";
+      break;
+    case Markup::Register:
+      OS << "<reg:";
+      break;
+    case Markup::Target:
+      OS << "<target:";
+      break;
+    case Markup::Memory:
+      OS << "<mem:";
+      break;
+    }
+  }
+}
+
+MCInstPrinter::WithMarkup::~WithMarkup() {
+  if (EnableMarkup)
+    OS << '>';
+  if (EnableColor)
+    OS.resetColor();
+}

llvm/lib/Target/AArch64/MCTargetDesc/AArch64InstPrinter.cpp

@@ -60,12 +60,12 @@ bool AArch64InstPrinter::applyTargetSpecificCLOption(StringRef Opt) {
 }
 
 void AArch64InstPrinter::printRegName(raw_ostream &OS, MCRegister Reg) const {
-  OS << markup("<reg:") << getRegisterName(Reg) << markup(">");
+  markup(OS, Markup::Register) << getRegisterName(Reg);
 }
 
 void AArch64InstPrinter::printRegName(raw_ostream &OS, MCRegister Reg,
                                       unsigned AltIdx) const {
-  OS << markup("<reg:") << getRegisterName(Reg, AltIdx) << markup(">");
+  markup(OS, Markup::Register) << getRegisterName(Reg, AltIdx);
 }
 
 StringRef AArch64InstPrinter::getRegName(MCRegister Reg) const {
@@ -175,7 +175,8 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
         printRegName(O, Op0.getReg());
         O << ", ";
         printRegName(O, Op1.getReg());
-        O << ", " << markup("<imm:") << "#" << shift << markup(">");
+        O << ", ";
+        markup(O, Markup::Immediate) << "#" << shift;
         printAnnotation(O, Annot);
         return;
       }
@@ -187,9 +188,10 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
       printRegName(O, Op0.getReg());
       O << ", ";
       printRegName(O, Op1.getReg());
-      O << ", " << markup("<imm:") << "#" << (Is64Bit ? 64 : 32) - Op2.getImm()
-        << markup(">") << ", " << markup("<imm:") << "#" << Op3.getImm() + 1
-        << markup(">");
+      O << ", ";
+      markup(O, Markup::Immediate) << "#" << (Is64Bit ? 64 : 32) - Op2.getImm();
+      O << ", ";
+      markup(O, Markup::Immediate) << "#" << Op3.getImm() + 1;
       printAnnotation(O, Annot);
       return;
     }
@@ -199,9 +201,10 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
     printRegName(O, Op0.getReg());
     O << ", ";
     printRegName(O, Op1.getReg());
-    O << ", " << markup("<imm:") << "#" << Op2.getImm() << markup(">") << ", "
-      << markup("<imm:") << "#" << Op3.getImm() - Op2.getImm() + 1
-      << markup(">");
+    O << ", ";
+    markup(O, Markup::Immediate) << "#" << Op2.getImm();
+    O << ", ";
+    markup(O, Markup::Immediate) << "#" << Op3.getImm() - Op2.getImm() + 1;
     printAnnotation(O, Annot);
     return;
   }
@@ -221,8 +224,10 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
 
       O << "\tbfc\t";
       printRegName(O, Op0.getReg());
-      O << ", " << markup("<imm:") << "#" << LSB << markup(">") << ", "
-        << markup("<imm:") << "#" << Width << markup(">");
+      O << ", ";
+      markup(O, Markup::Immediate) << "#" << LSB;
+      O << ", ";
+      markup(O, Markup::Immediate) << "#" << Width;
       printAnnotation(O, Annot);
       return;
     } else if (ImmS < ImmR) {
@@ -235,8 +240,10 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
       printRegName(O, Op0.getReg());
       O << ", ";
       printRegName(O, Op2.getReg());
-      O << ", " << markup("<imm:") << "#" << LSB << markup(">") << ", "
-        << markup("<imm:") << "#" << Width << markup(">");
+      O << ", ";
+      markup(O, Markup::Immediate) << "#" << LSB;
+      O << ", ";
+      markup(O, Markup::Immediate) << "#" << Width;
       printAnnotation(O, Annot);
       return;
     }
@@ -248,8 +255,10 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
     printRegName(O, Op0.getReg());
     O << ", ";
     printRegName(O, Op2.getReg());
-    O << ", " << markup("<imm:") << "#" << LSB << markup(">") << ", "
-      << markup("<imm:") << "#" << Width << markup(">");
+    O << ", ";
+    markup(O, Markup::Immediate) << "#" << LSB;
+    O << ", ";
+    markup(O, Markup::Immediate) << "#" << Width;
     printAnnotation(O, Annot);
     return;
   }
@@ -266,9 +275,12 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
       O << "\tmovn\t";
 
     printRegName(O, MI->getOperand(0).getReg());
-    O << ", " << markup("<imm:") << "#";
-    MI->getOperand(1).getExpr()->print(O, &MAI);
-    O << markup(">");
+    O << ", ";
+    {
+      WithMarkup M = markup(O, Markup::Immediate);
+      O << "#";
+      MI->getOperand(1).getExpr()->print(O, &MAI);
+    }
     return;
   }
 
@@ -276,9 +288,12 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
       MI->getOperand(2).isExpr()) {
     O << "\tmovk\t";
     printRegName(O, MI->getOperand(0).getReg());
-    O << ", " << markup("<imm:") << "#";
-    MI->getOperand(2).getExpr()->print(O, &MAI);
-    O << markup(">");
+    O << ", ";
+    {
+      WithMarkup M = markup(O, Markup::Immediate);
+      O << "#";
+      MI->getOperand(2).getExpr()->print(O, &MAI);
+    }
     return;
   }
 
@@ -286,8 +301,8 @@ void AArch64InstPrinter::printInst(const MCInst *MI, uint64_t Address,
     int64_t SExtVal = SignExtend64(Value, RegWidth);
     O << "\tmov\t";
     printRegName(O, MI->getOperand(0).getReg());
-    O << ", " << markup("<imm:") << "#"
-      << formatImm(SExtVal) << markup(">");
+    O << ", ";
+    markup(O, Markup::Immediate) << "#" << formatImm(SExtVal);
     if (CommentStream) {
       // Do the opposite to that used for instruction operands.
       if (getPrintImmHex())
@@ -813,8 +828,8 @@ void AArch64AppleInstPrinter::printInst(const MCInst *MI, uint64_t Address,
         printRegName(O, Reg);
       } else {
         assert(LdStDesc->NaturalOffset && "no offset on post-inc instruction?");
-        O << ", " << markup("<imm:") << "#" << LdStDesc->NaturalOffset
-          << markup(">");
+        O << ", ";
+        markup(O, Markup::Immediate) << "#" << LdStDesc->NaturalOffset;
       }
     }
 
@@ -1142,14 +1157,14 @@ void AArch64InstPrinter::printImm(const MCInst *MI, unsigned OpNo,
                                      const MCSubtargetInfo &STI,
                                      raw_ostream &O) {
   const MCOperand &Op = MI->getOperand(OpNo);
-  O << markup("<imm:") << "#" << formatImm(Op.getImm()) << markup(">");
+  markup(O, Markup::Immediate) << "#" << formatImm(Op.getImm());
 }
 
 void AArch64InstPrinter::printImmHex(const MCInst *MI, unsigned OpNo,
                                      const MCSubtargetInfo &STI,
                                      raw_ostream &O) {
   const MCOperand &Op = MI->getOperand(OpNo);
-  O << markup("<imm:") << format("#%#llx", Op.getImm()) << markup(">");
+  markup(O, Markup::Immediate) << format("#%#llx", Op.getImm());
 }
 
 template<int Size>
@@ -1158,13 +1173,11 @@ void AArch64InstPrinter::printSImm(const MCInst *MI, unsigned OpNo,
                                   raw_ostream &O) {
   const MCOperand &Op = MI->getOperand(OpNo);
   if (Size == 8)
-    O << markup("<imm:") << "#" << formatImm((signed char)Op.getImm())
-      << markup(">");
+    markup(O, Markup::Immediate) << "#" << formatImm((signed char)Op.getImm());
   else if (Size == 16)
-    O << markup("<imm:") << "#" << formatImm((signed short)Op.getImm())
-      << markup(">");
+    markup(O, Markup::Immediate) << "#" << formatImm((signed short)Op.getImm());
   else
-    O << markup("<imm:") << "#" << formatImm(Op.getImm()) << markup(">");
+    markup(O, Markup::Immediate) << "#" << formatImm(Op.getImm());
 }
 
 void AArch64InstPrinter::printPostIncOperand(const MCInst *MI, unsigned OpNo,
@@ -1173,7 +1186,7 @@ void AArch64InstPrinter::printPostIncOperand(const MCInst *MI, unsigned OpNo,
   if (Op.isReg()) {
     unsigned Reg = Op.getReg();
     if (Reg == AArch64::XZR)
-      O << markup("<imm:") << "#" << Imm << markup(">");
+      markup(O, Markup::Immediate) << "#" << Imm;
     else
       printRegName(O, Reg);
   } else
@@ -1206,7 +1219,7 @@ void AArch64InstPrinter::printAddSubImm(const MCInst *MI, unsigned OpNum,
     assert(Val == MO.getImm() && "Add/sub immediate out of range!");
     unsigned Shift =
         AArch64_AM::getShiftValue(MI->getOperand(OpNum + 1).getImm());
-    O << markup("<imm:") << '#' << formatImm(Val) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatImm(Val);
     if (Shift != 0) {
       printShifter(MI, OpNum + 1, STI, O);
       if (CommentStream)
@@ -1224,9 +1237,9 @@ void AArch64InstPrinter::printLogicalImm(const MCInst *MI, unsigned OpNum,
                                          const MCSubtargetInfo &STI,
                                          raw_ostream &O) {
   uint64_t Val = MI->getOperand(OpNum).getImm();
-  O << markup("<imm:") << "#0x";
+  WithMarkup M = markup(O, Markup::Immediate);
+  O << "#0x";
   O.write_hex(AArch64_AM::decodeLogicalImmediate(Val, 8 * sizeof(T)));
-  O << markup(">");
 }
 
 void AArch64InstPrinter::printShifter(const MCInst *MI, unsigned OpNum,
@@ -1238,8 +1251,8 @@ void AArch64InstPrinter::printShifter(const MCInst *MI, unsigned OpNum,
       AArch64_AM::getShiftValue(Val) == 0)
     return;
   O << ", " << AArch64_AM::getShiftExtendName(AArch64_AM::getShiftType(Val))
-    << " " << markup("<imm:") << "#" << AArch64_AM::getShiftValue(Val)
-    << markup(">");
+    << " ";
+  markup(O, Markup::Immediate) << "#" << AArch64_AM::getShiftValue(Val);
 }
 
 void AArch64InstPrinter::printShiftedRegister(const MCInst *MI, unsigned OpNum,
@@ -1273,19 +1286,23 @@ void AArch64InstPrinter::printArithExtend(const MCInst *MI, unsigned OpNum,
           ExtType == AArch64_AM::UXTX) ||
          ((Dest == AArch64::WSP || Src1 == AArch64::WSP) &&
           ExtType == AArch64_AM::UXTW) ) {
-      if (ShiftVal != 0)
-        O << ", lsl " << markup("<imm:") << "#" << ShiftVal << markup(">");
+      if (ShiftVal != 0) {
+        O << ", lsl ";
+        markup(O, Markup::Immediate) << "#" << ShiftVal;
+      }
       return;
     }
   }
   O << ", " << AArch64_AM::getShiftExtendName(ExtType);
-  if (ShiftVal != 0)
-    O << " " << markup("<imm:") << "#" << ShiftVal << markup(">");
+  if (ShiftVal != 0) {
+    O << " ";
+    markup(O, Markup::Immediate) << "#" << ShiftVal;
+  }
 }
 
-static void printMemExtendImpl(bool SignExtend, bool DoShift, unsigned Width,
-                               char SrcRegKind, raw_ostream &O,
-                               bool UseMarkup) {
+void AArch64InstPrinter::printMemExtendImpl(bool SignExtend, bool DoShift,
+                                            unsigned Width, char SrcRegKind,
+                                            raw_ostream &O) {
   // sxtw, sxtx, uxtw or lsl (== uxtx)
   bool IsLSL = !SignExtend && SrcRegKind == 'x';
   if (IsLSL)
@@ -1295,11 +1312,7 @@ static void printMemExtendImpl(bool SignExtend, bool DoShift, unsigned Width,
 
   if (DoShift || IsLSL) {
     O << " ";
-    if (UseMarkup)
-      O << "<imm:";
-    O << "#" << Log2_32(Width / 8);
-    if (UseMarkup)
-      O << ">";
+    markup(O, Markup::Immediate) << "#" << Log2_32(Width / 8);
   }
 }
 
@@ -1308,7 +1321,7 @@ void AArch64InstPrinter::printMemExtend(const MCInst *MI, unsigned OpNum,
                                         unsigned Width) {
   bool SignExtend = MI->getOperand(OpNum).getImm();
   bool DoShift = MI->getOperand(OpNum + 1).getImm();
-  printMemExtendImpl(SignExtend, DoShift, Width, SrcRegKind, O, UseMarkup);
+  printMemExtendImpl(SignExtend, DoShift, Width, SrcRegKind, O);
 }
 
 template <bool SignExtend, int ExtWidth, char SrcRegKind, char Suffix>
@@ -1325,7 +1338,7 @@ void AArch64InstPrinter::printRegWithShiftExtend(const MCInst *MI,
   bool DoShift = ExtWidth != 8;
   if (SignExtend || DoShift || SrcRegKind == 'w') {
     O << ", ";
-    printMemExtendImpl(SignExtend, DoShift, ExtWidth, SrcRegKind, O, UseMarkup);
+    printMemExtendImpl(SignExtend, DoShift, ExtWidth, SrcRegKind, O);
   }
 }
 
@@ -1384,8 +1397,8 @@ template <int Scale>
 void AArch64InstPrinter::printImmScale(const MCInst *MI, unsigned OpNum,
                                        const MCSubtargetInfo &STI,
                                        raw_ostream &O) {
-  O << markup("<imm:") << '#'
-    << formatImm(Scale * MI->getOperand(OpNum).getImm()) << markup(">");
+  markup(O, Markup::Immediate)
+      << '#' << formatImm(Scale * MI->getOperand(OpNum).getImm());
 }
 
 template <int Scale, int Offset>
@@ -1401,8 +1414,7 @@ void AArch64InstPrinter::printUImm12Offset(const MCInst *MI, unsigned OpNum,
                                            unsigned Scale, raw_ostream &O) {
   const MCOperand MO = MI->getOperand(OpNum);
   if (MO.isImm()) {
-    O << markup("<imm:") << '#' << formatImm(MO.getImm() * Scale)
-      << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatImm(MO.getImm() * Scale);
   } else {
     assert(MO.isExpr() && "Unexpected operand type!");
     MO.getExpr()->print(O, &MAI);
@@ -1415,8 +1427,8 @@ void AArch64InstPrinter::printAMIndexedWB(const MCInst *MI, unsigned OpNum,
   O << '[';
   printRegName(O, MI->getOperand(OpNum).getReg());
   if (MO1.isImm()) {
-    O << ", " << markup("<imm:") << "#" << formatImm(MO1.getImm() * Scale)
-      << markup(">");
+    O << ", ";
+    markup(O, Markup::Immediate) << "#" << formatImm(MO1.getImm() * Scale);
   } else {
     assert(MO1.isExpr() && "Unexpected operand type!");
     O << ", ";
@@ -1455,7 +1467,7 @@ void AArch64InstPrinter::printPrefetchOp(const MCInst *MI, unsigned OpNum,
     }
   }
 
-  O << markup("<imm:") << '#' << formatImm(prfop) << markup(">");
+  markup(O, Markup::Immediate) << '#' << formatImm(prfop);
 }
 
 void AArch64InstPrinter::printPSBHintOp(const MCInst *MI, unsigned OpNum,
@@ -1466,7 +1478,7 @@ void AArch64InstPrinter::printPSBHintOp(const MCInst *MI, unsigned OpNum,
   if (PSB)
     O << PSB->Name;
   else
-    O << markup("<imm:") << '#' << formatImm(psbhintop) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatImm(psbhintop);
 }
 
 void AArch64InstPrinter::printBTIHintOp(const MCInst *MI, unsigned OpNum,
@@ -1477,7 +1489,7 @@ void AArch64InstPrinter::printBTIHintOp(const MCInst *MI, unsigned OpNum,
   if (BTI)
     O << BTI->Name;
   else
-    O << markup("<imm:") << '#' << formatImm(btihintop) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatImm(btihintop);
 }
 
 void AArch64InstPrinter::printFPImmOperand(const MCInst *MI, unsigned OpNum,
@@ -1488,7 +1500,7 @@ void AArch64InstPrinter::printFPImmOperand(const MCInst *MI, unsigned OpNum,
                               : AArch64_AM::getFPImmFloat(MO.getImm());
 
   // 8 decimal places are enough to perfectly represent permitted floats.
-  O << markup("<imm:") << format("#%.8f", FPImm) << markup(">");
+  markup(O, Markup::Immediate) << format("#%.8f", FPImm);
 }
 
 static unsigned getNextVectorRegister(unsigned Reg, unsigned Stride = 1) {
@@ -1758,13 +1770,11 @@ void AArch64InstPrinter::printAlignedLabel(const MCInst *MI, uint64_t Address,
   // If the label has already been resolved to an immediate offset (say, when
   // we're running the disassembler), just print the immediate.
   if (Op.isImm()) {
-    O << markup("<imm:");
     int64_t Offset = Op.getImm() * 4;
     if (PrintBranchImmAsAddress)
-      O << formatHex(Address + Offset);
+      markup(O, Markup::Target) << formatHex(Address + Offset);
     else
-      O << "#" << formatImm(Offset);
-    O << markup(">");
+      markup(O, Markup::Immediate) << "#" << formatImm(Offset);
     return;
   }
 
@@ -1773,7 +1783,7 @@ void AArch64InstPrinter::printAlignedLabel(const MCInst *MI, uint64_t Address,
       dyn_cast<MCConstantExpr>(MI->getOperand(OpNum).getExpr());
   int64_t TargetAddress;
   if (BranchTarget && BranchTarget->evaluateAsAbsolute(TargetAddress)) {
-    O << formatHex((uint64_t)TargetAddress);
+    markup(O, Markup::Target) << formatHex((uint64_t)TargetAddress);
   } else {
     // Otherwise, just print the expression.
     MI->getOperand(OpNum).getExpr()->print(O, &MAI);
@@ -1794,12 +1804,11 @@ void AArch64InstPrinter::printAdrAdrpLabel(const MCInst *MI, uint64_t Address,
       Offset = Offset * 4096;
       Address = Address & -4096;
     }
-    O << markup("<imm:");
+    WithMarkup M = markup(O, Markup::Immediate);
     if (PrintBranchImmAsAddress)
-      O << formatHex(Address + Offset);
+      markup(O, Markup::Target) << formatHex(Address + Offset);
     else
-      O << "#" << Offset;
-    O << markup(">");
+      markup(O, Markup::Immediate) << "#" << Offset;
     return;
   }
 
@@ -1827,7 +1836,7 @@ void AArch64InstPrinter::printBarrierOption(const MCInst *MI, unsigned OpNo,
   if (!Name.empty())
     O << Name;
   else
-    O << markup("<imm:") << "#" << Val << markup(">");
+    markup(O, Markup::Immediate) << "#" << Val;
 }
 
 void AArch64InstPrinter::printBarriernXSOption(const MCInst *MI, unsigned OpNo,
@@ -1843,7 +1852,7 @@ void AArch64InstPrinter::printBarriernXSOption(const MCInst *MI, unsigned OpNo,
   if (!Name.empty())
     O << Name;
   else
-    O << markup("<imm:") << "#" << Val << markup(">");
+    markup(O, Markup::Immediate) << "#" << Val;
 }
 
 static bool isValidSysReg(const AArch64SysReg::SysReg *Reg, bool Read,
@@ -1942,7 +1951,7 @@ void AArch64InstPrinter::printSIMDType10Operand(const MCInst *MI, unsigned OpNo,
                                                 raw_ostream &O) {
   unsigned RawVal = MI->getOperand(OpNo).getImm();
   uint64_t Val = AArch64_AM::decodeAdvSIMDModImmType10(RawVal);
-  O << markup("<imm:") << format("#%#016llx", Val) << markup(">");
+  markup(O, Markup::Immediate) << format("#%#016llx", Val);
 }
 
 template<int64_t Angle, int64_t Remainder>
@@ -1950,7 +1959,7 @@ void AArch64InstPrinter::printComplexRotationOp(const MCInst *MI, unsigned OpNo,
                                                 const MCSubtargetInfo &STI,
                                                 raw_ostream &O) {
   unsigned Val = MI->getOperand(OpNo).getImm();
-  O << markup("<imm:") << "#" << (Val * Angle) + Remainder << markup(">");
+  markup(O, Markup::Immediate) << "#" << (Val * Angle) + Remainder;
 }
 
 void AArch64InstPrinter::printSVEPattern(const MCInst *MI, unsigned OpNum,
@@ -1960,7 +1969,7 @@ void AArch64InstPrinter::printSVEPattern(const MCInst *MI, unsigned OpNum,
   if (auto Pat = AArch64SVEPredPattern::lookupSVEPREDPATByEncoding(Val))
     O << Pat->Name;
   else
-    O << markup("<imm:") << '#' << formatImm(Val) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatImm(Val);
 }
 
 void AArch64InstPrinter::printSVEVecLenSpecifier(const MCInst *MI,
@@ -2004,9 +2013,9 @@ void AArch64InstPrinter::printImmSVE(T Value, raw_ostream &O) {
   std::make_unsigned_t<T> HexValue = Value;
 
   if (getPrintImmHex())
-    O << markup("<imm:") << '#' << formatHex((uint64_t)HexValue) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatHex((uint64_t)HexValue);
   else
-    O << markup("<imm:") << '#' << formatDec(Value) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatDec(Value);
 
   if (CommentStream) {
     // Do the opposite to that used for instruction operands.
@@ -2028,7 +2037,7 @@ void AArch64InstPrinter::printImm8OptLsl(const MCInst *MI, unsigned OpNum,
 
   // #0 lsl #8 is never pretty printed
   if ((UnscaledVal == 0) && (AArch64_AM::getShiftValue(Shift) != 0)) {
-    O << markup("<imm:") << '#' << formatImm(UnscaledVal) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatImm(UnscaledVal);
     printShifter(MI, OpNum + 1, STI, O);
     return;
   }
@@ -2058,7 +2067,7 @@ void AArch64InstPrinter::printSVELogicalImm(const MCInst *MI, unsigned OpNum,
   else if ((uint16_t)PrintVal == PrintVal)
     printImmSVE(PrintVal, O);
   else
-    O << markup("<imm:") << '#' << formatHex((uint64_t)PrintVal) << markup(">");
+    markup(O, Markup::Immediate) << '#' << formatHex((uint64_t)PrintVal);
 }
 
 template <int Width>
@@ -2086,8 +2095,8 @@ void AArch64InstPrinter::printExactFPImm(const MCInst *MI, unsigned OpNum,
   auto *Imm0Desc = AArch64ExactFPImm::lookupExactFPImmByEnum(ImmIs0);
   auto *Imm1Desc = AArch64ExactFPImm::lookupExactFPImmByEnum(ImmIs1);
   unsigned Val = MI->getOperand(OpNum).getImm();
-  O << markup("<imm:") << "#" << (Val ? Imm1Desc->Repr : Imm0Desc->Repr)
-    << markup(">");
+  markup(O, Markup::Immediate)
+      << "#" << (Val ? Imm1Desc->Repr : Imm0Desc->Repr);
 }
 
 void AArch64InstPrinter::printGPR64as32(const MCInst *MI, unsigned OpNum,

llvm/lib/Target/AArch64/MCTargetDesc/AArch64InstPrinter.h

@@ -91,7 +91,8 @@ protected:
                              const MCSubtargetInfo &STI, raw_ostream &O);
   void printArithExtend(const MCInst *MI, unsigned OpNum,
                         const MCSubtargetInfo &STI, raw_ostream &O);
-
+  void printMemExtendImpl(bool SignExtend, bool DoShift, unsigned Width,
+                          char SrcRegKind, raw_ostream &O);
   void printMemExtend(const MCInst *MI, unsigned OpNum, raw_ostream &O,
                       char SrcRegKind, unsigned Width);
   template <char SrcRegKind, unsigned Width>

llvm/test/MC/Disassembler/AArch64/colored.txt

@@ -0,0 +1,7 @@
+# UNSUPPORTED: system-windows
+# RUN: llvm-mc -triple=aarch64 -mattr=+all --cdis %s | FileCheck %s --strict-whitespace --match-full-lines
+
+0xa1 0x00 0x00 0x54
+# CHECK:	b.ne	#20
+0x00 0x7c 0x00 0x13
+# CHECK-NEXT:	asr	w0, w0, #0

llvm/tools/llvm-mc/llvm-mc.cpp

@@ -216,6 +216,7 @@ enum ActionType {
   AC_Assemble,
   AC_Disassemble,
   AC_MDisassemble,
+  AC_CDisassemble,
 };
 
 static cl::opt<ActionType> Action(
@@ -226,7 +227,9 @@ static cl::opt<ActionType> Action(
                clEnumValN(AC_Disassemble, "disassemble",
                           "Disassemble strings of hex bytes"),
                clEnumValN(AC_MDisassemble, "mdis",
-                          "Marked up disassembly of strings of hex bytes")),
+                          "Marked up disassembly of strings of hex bytes"),
+               clEnumValN(AC_CDisassemble, "cdis",
+                          "Colored disassembly of strings of hex bytes")),
     cl::cat(MCCategory));
 
 static const Target *GetTarget(const char *ProgName) {
@@ -544,6 +547,11 @@ int main(int argc, char **argv) {
     std::unique_ptr<MCAsmBackend> MAB(
         TheTarget->createMCAsmBackend(*STI, *MRI, MCOptions));
     auto FOut = std::make_unique<formatted_raw_ostream>(*OS);
+    // FIXME: Workaround for bug in formatted_raw_ostream. Color escape codes
+    // are (incorrectly) written directly to the unbuffered raw_ostream wrapped
+    // by the formatted_raw_ostream.
+    if (Action == AC_CDisassemble)
+      FOut->SetUnbuffered();
     Str.reset(
         TheTarget->createAsmStreamer(Ctx, std::move(FOut), /*asmverbose*/ true,
                                      /*useDwarfDirectory*/ true, IP,
@@ -586,10 +594,13 @@ int main(int argc, char **argv) {
                         *MCII, MCOptions);
     break;
   case AC_MDisassemble:
-    assert(IP && "Expected assembly output");
     IP->setUseMarkup(true);
     disassemble = true;
     break;
+  case AC_CDisassemble:
+    IP->setUseColor(true);
+    disassemble = true;
+    break;
   case AC_Disassemble:
     disassemble = true;
     break;