[llvm] [llvm-objdump] Fix coloring with nested WithMarkup (PR #113834)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mc

@llvm/pr-subscribers-llvm-binary-utilities

Author: Fangrui Song (MaskRay)

<details>
<summary>Changes</summary>

WithMarkup objects may nest, resulting in the `)` in `leaq
(%rdx,%rax), %rbx` to be green instead of the default color,
mismatching the color of `(`.

```
% llvm-mc -triple=x86_64 -mdis <<< '0x48 0x8d 0x1c 0x02'
        .text
        leaq    <mem:(<reg:%rdx>,<reg:%rax>)>, <reg:%rbx>
```

To ensure that `(` and `)` get the same color, maintain a color stack
within MCInstPrinter.


---
Full diff: https://github.com/llvm/llvm-project/pull/113834.diff


3 Files Affected:

- (modified) llvm/include/llvm/MC/MCInstPrinter.h (+7-3) 
- (modified) llvm/lib/MC/MCInstPrinter.cpp (+17-10) 
- (added) llvm/test/tools/llvm-objdump/X86/disassemble-color.s (+21) 


``````````diff
diff --git a/llvm/include/llvm/MC/MCInstPrinter.h b/llvm/include/llvm/MC/MCInstPrinter.h
index 60a901e3d0deae..c296a6fe1c045d 100644
--- a/llvm/include/llvm/MC/MCInstPrinter.h
+++ b/llvm/include/llvm/MC/MCInstPrinter.h
@@ -9,8 +9,10 @@
 #ifndef LLVM_MC_MCINSTPRINTER_H
 #define LLVM_MC_MCINSTPRINTER_H
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Format.h"
+#include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 
 namespace llvm {
@@ -24,7 +26,6 @@ class MCRegister;
 class MCRegisterInfo;
 class MCSubtargetInfo;
 class StringRef;
-class raw_ostream;
 
 /// Convert `Bytes' to a hex string and output to `OS'
 void dumpBytes(ArrayRef<uint8_t> Bytes, raw_ostream &OS);
@@ -76,6 +77,8 @@ class MCInstPrinter {
   /// If true, symbolize branch target and memory reference operands.
   bool SymbolizeOperands = false;
 
+  SmallVector<raw_ostream::Colors, 4> ColorStack{raw_ostream::Colors::RESET};
+
   /// Utility function for printing annotations.
   void printAnnotation(raw_ostream &OS, StringRef Annot);
 
@@ -98,8 +101,8 @@ class MCInstPrinter {
 
   class WithMarkup {
   public:
-    LLVM_CTOR_NODISCARD WithMarkup(raw_ostream &OS, Markup M, bool EnableMarkup,
-                                   bool EnableColor);
+    LLVM_CTOR_NODISCARD WithMarkup(MCInstPrinter &IP, raw_ostream &OS, Markup M,
+                                   bool EnableMarkup, bool EnableColor);
     ~WithMarkup();
 
     template <typename T> WithMarkup &operator<<(T &O) {
@@ -113,6 +116,7 @@ class MCInstPrinter {
     }
 
   private:
+    MCInstPrinter &IP;
     raw_ostream &OS;
     bool EnableMarkup;
     bool EnableColor;
diff --git a/llvm/lib/MC/MCInstPrinter.cpp b/llvm/lib/MC/MCInstPrinter.cpp
index e4faeba04a8fd7..a57f4e19c81c6b 100644
--- a/llvm/lib/MC/MCInstPrinter.cpp
+++ b/llvm/lib/MC/MCInstPrinter.cpp
@@ -226,27 +226,32 @@ format_object<uint64_t> MCInstPrinter::formatHex(uint64_t Value) const {
 
 MCInstPrinter::WithMarkup MCInstPrinter::markup(raw_ostream &OS,
                                                 Markup S) const {
-  return WithMarkup(OS, S, getUseMarkup(), getUseColor());
+  return WithMarkup(const_cast<MCInstPrinter &>(*this), OS, S, getUseMarkup(),
+                    getUseColor());
 }
 
-MCInstPrinter::WithMarkup::WithMarkup(raw_ostream &OS, Markup M,
-                                      bool EnableMarkup, bool EnableColor)
-    : OS(OS), EnableMarkup(EnableMarkup), EnableColor(EnableColor) {
+MCInstPrinter::WithMarkup::WithMarkup(MCInstPrinter &IP, raw_ostream &OS,
+                                      Markup M, bool EnableMarkup,
+                                      bool EnableColor)
+    : IP(IP), OS(OS), EnableMarkup(EnableMarkup), EnableColor(EnableColor) {
   if (EnableColor) {
+    raw_ostream::Colors Color = raw_ostream::Colors::RESET;
     switch (M) {
     case Markup::Immediate:
-      OS.changeColor(raw_ostream::RED);
+      Color = raw_ostream::RED;
       break;
     case Markup::Register:
-      OS.changeColor(raw_ostream::CYAN);
+      Color = raw_ostream::CYAN;
       break;
     case Markup::Target:
-      OS.changeColor(raw_ostream::YELLOW);
+      Color = raw_ostream::YELLOW;
       break;
     case Markup::Memory:
-      OS.changeColor(raw_ostream::GREEN);
+      Color = raw_ostream::GREEN;
       break;
     }
+    IP.ColorStack.push_back(Color);
+    OS.changeColor(Color);
   }
 
   if (EnableMarkup) {
@@ -270,6 +275,8 @@ MCInstPrinter::WithMarkup::WithMarkup(raw_ostream &OS, Markup M,
 MCInstPrinter::WithMarkup::~WithMarkup() {
   if (EnableMarkup)
     OS << '>';
-  if (EnableColor)
-    OS.resetColor();
+  if (!EnableColor)
+    return;
+  IP.ColorStack.pop_back();
+  OS << IP.ColorStack.back();
 }
diff --git a/llvm/test/tools/llvm-objdump/X86/disassemble-color.s b/llvm/test/tools/llvm-objdump/X86/disassemble-color.s
new file mode 100644
index 00000000000000..4e1d82562fb546
--- /dev/null
+++ b/llvm/test/tools/llvm-objdump/X86/disassemble-color.s
@@ -0,0 +1,21 @@
+# UNSUPPORTED: system-windows
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o %t
+# RUN: llvm-objdump -d --no-show-raw-insn --disassembler-color=on %t | FileCheck %s --check-prefix=ATT
+# RUN: llvm-objdump -d --no-show-raw-insn --disassembler-color=on -M intel %t | FileCheck %s --check-prefix=INTEL
+
+# ATT:      <.text>:
+# ATT-NEXT:  leaq	(%rdx,%rax,4), %rbx
+# ATT-NEXT:  movq	(,%rax), %rbx
+# ATT-NEXT:  leaq	0x3(%rdx,%rax), %rbx
+# ATT-NEXT:  movq	$0x3, %rax
+
+# INTEL:      <.text>:
+# INTEL-NEXT:  lea	rbx, [rdx + 4*rax]
+# INTEL-NEXT:  mov	rbx, qword ptr [1*rax]
+# INTEL-NEXT:  lea	rbx, [rdx + rax + 0x3]
+# INTEL-NEXT:  mov	rax, 0x3
+
+leaq (%rdx,%rax,4), %rbx
+movq (,%rax), %rbx
+leaq 3(%rdx,%rax), %rbx
+movq $3, %rax

``````````

</details>


https://github.com/llvm/llvm-project/pull/113834