[llvm] 95062d7 - [llvm-objdump] Enable disassembly color highlighting

[Jonas Devlieghere via llvm-commits]

Author: Jonas Devlieghere
Date: 2023-09-01T09:32:34-07:00
New Revision: 95062d74676b98a85e4d3db6628d239b15dbeb35

URL: https://github.com/llvm/llvm-project/commit/95062d74676b98a85e4d3db6628d239b15dbeb35
DIFF: https://github.com/llvm/llvm-project/commit/95062d74676b98a85e4d3db6628d239b15dbeb35.diff

LOG: [llvm-objdump] Enable disassembly color highlighting

Enable color highlighting of disassembly in llvm-objdump. This patch
introduces a new flag --disassembler-color=<mode> that enables or
disables highlighting disassembly with ANSI escape codes. The default
mode is to enable color highlighting if outputting to a color-enabled
terminal.

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

Added: 
    llvm/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
    llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
    llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s

Modified: 
    llvm/docs/CommandGuide/llvm-objdump.rst
    llvm/tools/llvm-objdump/ObjdumpOpts.td
    llvm/tools/llvm-objdump/llvm-objdump.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/docs/CommandGuide/llvm-objdump.rst b/llvm/docs/CommandGuide/llvm-objdump.rst
index cd3d14b4114347..b156b212e461f2 100644
--- a/llvm/docs/CommandGuide/llvm-objdump.rst
+++ b/llvm/docs/CommandGuide/llvm-objdump.rst
@@ -179,6 +179,15 @@ OPTIONS
   * ``att``: x86 only (default). Print in the AT&T syntax.
   * ``intel``: x86 only. Print in the intel syntax.
 
+
+.. option::  --disassembler-color=<mode>
+
+  Enable or disable disassembler color output.
+
+  * ``off``: Disable disassembler color output.
+  * ``on``: Enable disassembler color output.
+  * ``terminal``: Enable disassembler color output if the terminal supports it (default).
+
 .. option:: --mcpu=<cpu-name>
 
   Target a specific CPU type for disassembly. Specify ``--mcpu=help`` to display

diff  --git a/llvm/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s b/llvm/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
new file mode 100644
index 00000000000000..336eb11140eeb9
--- /dev/null
+++ b/llvm/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
@@ -0,0 +1,25 @@
+// RUN: llvm-mc -triple arm64-apple-macosx %s -filetype=obj -o %t
+// RUN: llvm-objdump --disassembler-color=on --disassemble %t | FileCheck %s --check-prefix=COLOR
+// RUN: llvm-objdump --disassembler-color=off --disassemble %t | FileCheck %s --check-prefix=NOCOLOR
+// RUN: llvm-objdump --disassembler-color=terminal --disassemble %t | FileCheck %s --check-prefix=NOCOLOR
+
+sub	sp, sp, #16
+str	w0, [sp, #12]
+ldr	w8, [sp, #12]
+ldr	w9, [sp, #12]
+mul	w0, w8, w9
+add	sp, sp, #16
+
+// NOCOLOR: sub	sp, sp, #0x10
+// NOCOLOR: str	w0, [sp, #0xc]
+// NOCOLOR: ldr	w8, [sp, #0xc]
+// NOCOLOR: ldr	w9, [sp, #0xc]
+// NOCOLOR: mul	w0, w8, w9
+// NOCOLOR: add	sp, sp, #0x10
+
+// COLOR: sub	sp, sp, #0x10
+// COLOR: str	w0, [sp, #0xc]
+// COLOR: ldr	w8, [sp, #0xc]
+// COLOR: ldr	w9, [sp, #0xc]
+// COLOR: mul	w0, w8, w9
+// COLOR: add	sp, sp, #0x10

diff  --git a/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s b/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
new file mode 100644
index 00000000000000..336eb11140eeb9
--- /dev/null
+++ b/llvm/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
@@ -0,0 +1,25 @@
+// RUN: llvm-mc -triple arm64-apple-macosx %s -filetype=obj -o %t
+// RUN: llvm-objdump --disassembler-color=on --disassemble %t | FileCheck %s --check-prefix=COLOR
+// RUN: llvm-objdump --disassembler-color=off --disassemble %t | FileCheck %s --check-prefix=NOCOLOR
+// RUN: llvm-objdump --disassembler-color=terminal --disassemble %t | FileCheck %s --check-prefix=NOCOLOR
+
+sub	sp, sp, #16
+str	w0, [sp, #12]
+ldr	w8, [sp, #12]
+ldr	w9, [sp, #12]
+mul	w0, w8, w9
+add	sp, sp, #16
+
+// NOCOLOR: sub	sp, sp, #0x10
+// NOCOLOR: str	w0, [sp, #0xc]
+// NOCOLOR: ldr	w8, [sp, #0xc]
+// NOCOLOR: ldr	w9, [sp, #0xc]
+// NOCOLOR: mul	w0, w8, w9
+// NOCOLOR: add	sp, sp, #0x10
+
+// COLOR: sub	sp, sp, #0x10
+// COLOR: str	w0, [sp, #0xc]
+// COLOR: ldr	w8, [sp, #0xc]
+// COLOR: ldr	w9, [sp, #0xc]
+// COLOR: mul	w0, w8, w9
+// COLOR: add	sp, sp, #0x10

diff  --git a/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s b/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
new file mode 100644
index 00000000000000..336eb11140eeb9
--- /dev/null
+++ b/llvm/test/tools/llvm-objdump/MachO/arm64-disassembly-color.s
@@ -0,0 +1,25 @@
+// RUN: llvm-mc -triple arm64-apple-macosx %s -filetype=obj -o %t
+// RUN: llvm-objdump --disassembler-color=on --disassemble %t | FileCheck %s --check-prefix=COLOR
+// RUN: llvm-objdump --disassembler-color=off --disassemble %t | FileCheck %s --check-prefix=NOCOLOR
+// RUN: llvm-objdump --disassembler-color=terminal --disassemble %t | FileCheck %s --check-prefix=NOCOLOR
+
+sub	sp, sp, #16
+str	w0, [sp, #12]
+ldr	w8, [sp, #12]
+ldr	w9, [sp, #12]
+mul	w0, w8, w9
+add	sp, sp, #16
+
+// NOCOLOR: sub	sp, sp, #0x10
+// NOCOLOR: str	w0, [sp, #0xc]
+// NOCOLOR: ldr	w8, [sp, #0xc]
+// NOCOLOR: ldr	w9, [sp, #0xc]
+// NOCOLOR: mul	w0, w8, w9
+// NOCOLOR: add	sp, sp, #0x10
+
+// COLOR: sub	sp, sp, #0x10
+// COLOR: str	w0, [sp, #0xc]
+// COLOR: ldr	w8, [sp, #0xc]
+// COLOR: ldr	w9, [sp, #0xc]
+// COLOR: mul	w0, w8, w9
+// COLOR: add	sp, sp, #0x10

diff  --git a/llvm/tools/llvm-objdump/ObjdumpOpts.td b/llvm/tools/llvm-objdump/ObjdumpOpts.td
index e3e74762420deb..100a95d3d92542 100644
--- a/llvm/tools/llvm-objdump/ObjdumpOpts.td
+++ b/llvm/tools/llvm-objdump/ObjdumpOpts.td
@@ -87,6 +87,11 @@ def disassembler_options_EQ : Joined<["--"], "disassembler-options=">,
 def : JoinedOrSeparate<["-"], "M">, Alias<disassembler_options_EQ>,
   HelpText<"Alias for --disassembler-options=">;
 
+def disassembler_color_EQ : Joined<["--"], "disassembler-color=">,
+  MetaVarName<"mode">,
+  HelpText<"Enable or disable disassembler color output. "
+           "Valid options are \"on\", \"off\" and \"terminal\" (default)">;
+
 def dynamic_reloc : Flag<["--"], "dynamic-reloc">,
   HelpText<"Display the dynamic relocation entries in the file">;
 def : Flag<["-"], "R">, Alias<dynamic_reloc>,

diff  --git a/llvm/tools/llvm-objdump/llvm-objdump.cpp b/llvm/tools/llvm-objdump/llvm-objdump.cpp
index 1e5175b527710b..96d74d6e2d5e86 100644
--- a/llvm/tools/llvm-objdump/llvm-objdump.cpp
+++ b/llvm/tools/llvm-objdump/llvm-objdump.cpp
@@ -177,6 +177,13 @@ class OtoolOptTable : public CommonOptTable {
 
 #define DEBUG_TYPE "objdump"
 
+enum class ColorOutput {
+  Auto,
+  Enable,
+  Disable,
+  Invalid,
+};
+
 static uint64_t AdjustVMA;
 static bool AllHeaders;
 static std::string ArchName;
@@ -189,6 +196,7 @@ bool objdump::TracebackTable;
 static std::vector<std::string> DisassembleSymbols;
 static bool DisassembleZeroes;
 static std::vector<std::string> DisassemblerOptions;
+static ColorOutput DisassemblyColor;
 DIDumpType objdump::DwarfDumpType;
 static bool DynamicRelocations;
 static bool FaultMapSection;
@@ -900,6 +908,19 @@ DisassemblerTarget::DisassemblerTarget(const Target *TheTarget, ObjectFile &Obj,
   InstPrinter->setPrintBranchImmAsAddress(true);
   InstPrinter->setSymbolizeOperands(SymbolizeOperands);
   InstPrinter->setMCInstrAnalysis(InstrAnalysis.get());
+
+  switch (DisassemblyColor) {
+  case ColorOutput::Enable:
+    InstPrinter->setUseColor(true);
+    break;
+  case ColorOutput::Auto:
+    InstPrinter->setUseColor(outs().has_colors());
+    break;
+  case ColorOutput::Disable:
+  case ColorOutput::Invalid:
+    InstPrinter->setUseColor(false);
+    break;
+  };
 }
 
 DisassemblerTarget::DisassemblerTarget(DisassemblerTarget &Other,
@@ -1900,6 +1921,13 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
 
       formatted_raw_ostream FOS(outs());
 
+      // 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 (DisassemblyColor == ColorOutput::Enable ||
+          DisassemblyColor == ColorOutput::Auto)
+        FOS.SetUnbuffered();
+
       std::unordered_map<uint64_t, std::string> AllLabels;
       std::unordered_map<uint64_t, std::vector<std::string>> BBAddrMapLabels;
       if (SymbolizeOperands) {
@@ -3193,6 +3221,16 @@ static void parseObjdumpOptions(const llvm::opt::InputArgList &InputArgs) {
     if (DbgVariables == DVInvalid)
       invalidArgValue(A);
   }
+  if (const opt::Arg *A = InputArgs.getLastArg(OBJDUMP_disassembler_color_EQ)) {
+    DisassemblyColor = StringSwitch<ColorOutput>(A->getValue())
+                           .Case("on", ColorOutput::Enable)
+                           .Case("off", ColorOutput::Disable)
+                           .Case("terminal", ColorOutput::Auto)
+                           .Default(ColorOutput::Invalid);
+    if (DisassemblyColor == ColorOutput::Invalid)
+      invalidArgValue(A);
+  }
+
   parseIntArg(InputArgs, OBJDUMP_debug_vars_indent_EQ, DbgIndent);
 
   parseMachOOptions(InputArgs);