[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

Author: JDevlieghere

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 @@ class MCInstPrinter {
   /// 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 @@ class MCInstPrinter {
 
   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 @@ class MCInstPrinter {
   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();
+}