[llvm] 29bff4a - [llvm-objdump] Fix coloring with nested WithMarkup

[via llvm-commits]

Author: Fangrui Song
Date: 2024-10-29T20:06:56-07:00
New Revision: 29bff4aad8eb7f54f99e0496b735aee193063b04

URL: https://github.com/llvm/llvm-project/commit/29bff4aad8eb7f54f99e0496b735aee193063b04
DIFF: https://github.com/llvm/llvm-project/commit/29bff4aad8eb7f54f99e0496b735aee193063b04.diff

LOG: [llvm-objdump] Fix coloring with nested WithMarkup

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.

Fix #99661

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

Added: 
    llvm/test/tools/llvm-objdump/X86/disassemble-color.s

Modified: 
    llvm/include/llvm/MC/MCInstPrinter.h
    llvm/lib/MC/MCInstPrinter.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/MC/MCInstPrinter.h b/llvm/include/llvm/MC/MCInstPrinter.h
index 0b9c738a7a0a31..e825c04a6dba6f 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 488e34a6d53954..069716a3ecf9b7 100644
--- a/llvm/lib/MC/MCInstPrinter.cpp
+++ b/llvm/lib/MC/MCInstPrinter.cpp
@@ -225,27 +225,31 @@ format_object<uint64_t> MCInstPrinter::formatHex(uint64_t Value) const {
 }
 
 MCInstPrinter::WithMarkup MCInstPrinter::markup(raw_ostream &OS, Markup S) {
-  return WithMarkup(OS, S, getUseMarkup(), getUseColor());
+  return WithMarkup(*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) {
@@ -269,6 +273,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