[llvm-objdump] Fix coloring with nested WithMarkup

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;

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();
 }

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