[lldb] Add syntax color highlighting for disassembly

JDevlieghere · JDevlieghere · commit 0029013c5575 · 2023-09-14T16:04:58.000-07:00
Add support for syntax color highlighting disassembly in LLDB. This patch relies on 77d1032, which introduces support for syntax highlighting in MC. Currently only AArch64 and X86 have color support, but other interested backends can adopt WithColor in their respective MCInstPrinter. Differential revision: https://reviews.llvm.org/D159164 (cherry picked from commit a69f78b)
diff --git a/lldb/include/lldb/Core/Disassembler.h b/lldb/include/lldb/Core/Disassembler.h
@@ -64,14 +64,16 @@ class Instruction {
 
   const Address &GetAddress() const { return m_address; }
 
-  const char *GetMnemonic(const ExecutionContext *exe_ctx) {
+  const char *GetMnemonic(const ExecutionContext *exe_ctx,
+                          bool markup = false) {
     CalculateMnemonicOperandsAndCommentIfNeeded(exe_ctx);
-    return m_opcode_name.c_str();
+    return markup ? m_markup_opcode_name.c_str() : m_opcode_name.c_str();
   }
 
-  const char *GetOperands(const ExecutionContext *exe_ctx) {
+  const char *GetOperands(const ExecutionContext *exe_ctx,
+                          bool markup = false) {
     CalculateMnemonicOperandsAndCommentIfNeeded(exe_ctx);
-    return m_mnemonics.c_str();
+    return markup ? m_markup_mnemonics.c_str() : m_mnemonics.c_str();
   }
 
   const char *GetComment(const ExecutionContext *exe_ctx) {
@@ -244,7 +246,9 @@ class Instruction {
 protected:
   Opcode m_opcode; // The opcode for this instruction
   std::string m_opcode_name;
+  std::string m_markup_opcode_name;
   std::string m_mnemonics;
+  std::string m_markup_mnemonics;
   std::string m_comment;
   bool m_calculated_strings;
 
diff --git a/lldb/source/Core/Disassembler.cpp b/lldb/source/Core/Disassembler.cpp
@@ -645,18 +645,29 @@ void Instruction::Dump(lldb_private::Stream *s, uint32_t max_opcode_byte_size,
                            instruction_control_flow_kind));
   }
 
+  bool show_color = false;
+  if (exe_ctx) {
+    if (TargetSP target_sp = exe_ctx->GetTargetSP()) {
+      show_color = target_sp->GetDebugger().GetUseColor();
+    }
+  }
   const size_t opcode_pos = ss.GetSizeOfLastLine();
+  const std::string &opcode_name =
+      show_color ? m_markup_opcode_name : m_opcode_name;
+  const std::string &mnemonics = show_color ? m_markup_mnemonics : m_mnemonics;
 
   // The default opcode size of 7 characters is plenty for most architectures
   // but some like arm can pull out the occasional vqrshrun.s16.  We won't get
-  // consistent column spacing in these cases, unfortunately.
+  // consistent column spacing in these cases, unfortunately. Also note that we
+  // need to directly use m_opcode_name here (instead of opcode_name) so we
+  // don't include color codes as characters.
   if (m_opcode_name.length() >= opcode_column_width) {
     opcode_column_width = m_opcode_name.length() + 1;
   }
 
-  ss.PutCString(m_opcode_name);
+  ss.PutCString(opcode_name);
   ss.FillLastLineToColumn(opcode_pos + opcode_column_width, ' ');
-  ss.PutCString(m_mnemonics);
+  ss.PutCString(mnemonics);
 
   if (!m_comment.empty()) {
     ss.FillLastLineToColumn(
diff --git a/lldb/source/Plugins/Disassembler/LLVMC/DisassemblerLLVMC.cpp b/lldb/source/Plugins/Disassembler/LLVMC/DisassemblerLLVMC.cpp
@@ -63,6 +63,8 @@ class DisassemblerLLVMC::MCDisasmInstance {
   void PrintMCInst(llvm::MCInst &mc_inst, lldb::addr_t pc,
                    std::string &inst_string, std::string &comments_string);
   void SetStyle(bool use_hex_immed, HexImmediateStyle hex_style);
+  void SetUseColor(bool use_color);
+  bool GetUseColor() const;
   bool CanBranch(llvm::MCInst &mc_inst) const;
   bool HasDelaySlot(llvm::MCInst &mc_inst) const;
   bool IsCall(llvm::MCInst &mc_inst) const;
@@ -565,7 +567,9 @@ class InstructionLLVMC : public lldb_private::Instruction {
 
     if (m_opcode.GetData(data)) {
       std::string out_string;
+      std::string markup_out_string;
       std::string comment_string;
+      std::string markup_comment_string;
 
       DisassemblerScope disasm(*this, exe_ctx);
       if (disasm) {
@@ -607,7 +611,14 @@ class InstructionLLVMC : public lldb_private::Instruction {
 
         if (inst_size > 0) {
           mc_disasm_ptr->SetStyle(use_hex_immediates, hex_style);
+
+          const bool saved_use_color = mc_disasm_ptr->GetUseColor();
+          mc_disasm_ptr->SetUseColor(false);
           mc_disasm_ptr->PrintMCInst(inst, pc, out_string, comment_string);
+          mc_disasm_ptr->SetUseColor(true);
+          mc_disasm_ptr->PrintMCInst(inst, pc, markup_out_string,
+                                     markup_comment_string);
+          mc_disasm_ptr->SetUseColor(saved_use_color);
 
           if (!comment_string.empty()) {
             AppendComment(comment_string);
@@ -672,6 +683,11 @@ class InstructionLLVMC : public lldb_private::Instruction {
           m_opcode_name = matches[1].str();
           m_mnemonics = matches[2].str();
         }
+        matches.clear();
+        if (s_regex.Execute(markup_out_string, &matches)) {
+          m_markup_opcode_name = matches[1].str();
+          m_markup_mnemonics = matches[2].str();
+        }
       }
     }
   }
@@ -1344,10 +1360,12 @@ void DisassemblerLLVMC::MCDisasmInstance::PrintMCInst(
   llvm::raw_string_ostream inst_stream(inst_string);
   llvm::raw_string_ostream comments_stream(comments_string);
 
+  inst_stream.enable_colors(m_instr_printer_up->getUseColor());
   m_instr_printer_up->setCommentStream(comments_stream);
   m_instr_printer_up->printInst(&mc_inst, pc, llvm::StringRef(),
                                 *m_subtarget_info_up, inst_stream);
   m_instr_printer_up->setCommentStream(llvm::nulls());
+
   comments_stream.flush();
 
   static std::string g_newlines("\r\n");
@@ -1374,6 +1392,14 @@ void DisassemblerLLVMC::MCDisasmInstance::SetStyle(
   }
 }
 
+void DisassemblerLLVMC::MCDisasmInstance::SetUseColor(bool use_color) {
+  m_instr_printer_up->setUseColor(use_color);
+}
+
+bool DisassemblerLLVMC::MCDisasmInstance::GetUseColor() const {
+  return m_instr_printer_up->getUseColor();
+}
+
 bool DisassemblerLLVMC::MCDisasmInstance::CanBranch(
     llvm::MCInst &mc_inst) const {
   if (m_instr_analysis_up)
diff --git a/lldb/test/API/python_api/disassemble-raw-data/TestDisassembleRawData.py b/lldb/test/API/python_api/disassemble-raw-data/TestDisassembleRawData.py
@@ -59,9 +59,19 @@ def test_disassemble_raw_data(self):
         elif arch in ("aarch64", "arm64"):
             self.assertEqual(inst.GetMnemonic(target), "mov")
             self.assertEqual(inst.GetOperands(target), "w0, #0x63")
+            self.assertEqual(inst.GetComment(target), "=99 ")
             self.assertEqual(
                 inst.GetControlFlowKind(target), lldb.eInstructionControlFlowKindUnknown
             )
+            # Make sure that using colors doesn't affect the output here.
+            res = lldb.SBCommandReturnObject()
+            ci = self.dbg.GetCommandInterpreter()
+            ci.HandleCommand("settings set use-color true", res)
+            self.assertEqual(inst.GetOperands(target), "w0, #0x63")
+            self.assertEqual(inst.GetMnemonic(target), "mov")
+            self.assertEqual(inst.GetComment(target), "=99 ")
+            ci.HandleCommand("settings set use-color false", res)
+
         elif arch == "arm":
             self.assertEqual(inst.GetMnemonic(target), "mov")
             self.assertEqual(inst.GetOperands(target), "r3, #99")
diff --git a/lldb/test/Shell/Commands/command-disassemble-aarch64-color.s b/lldb/test/Shell/Commands/command-disassemble-aarch64-color.s
@@ -0,0 +1,33 @@
+# UNSUPPORTED: system-windows
+# REQUIRES: aarch64
+
+# This checks that lldb's disassembler colors AArch64 disassembly.
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-linux-gnueabihf %s -o %t --mattr=+all
+# RUN: %lldb %t -o "settings set use-color true" -o "disassemble -n fn" -o exit 2>&1 | FileCheck %s
+
+.globl  fn
+.type   fn, @function
+fn:
+  // These are in alphabetical order by extension name
+  aesd v0.16b, v0.16b                   // AEK_AES
+  bfadd z23.h, p3/m, z23.h, z13.h       // AEK_B16B16
+  bfdot v2.2s, v3.4h, v4.4h             // AEK_BF16
+  brb iall                              // AEK_BRBE
+  crc32b w0, w0, w0                     // AEK_CRC
+  // AEK_CRYPTO enables a combination of other features
+  smin x0, x0, #0                       // AEK_CSSC
+  sysp	#0, c2, c0, #0, x0, x1          // AEK_D128
+  sdot v0.2s, v1.8b, v2.8b              // AEK_DOTPROD
+  fmmla z0.s, z1.s, z2.s                // AEK_F32MM
+
+# CHECK: `fn:
+# CHECK-NEXT: [0x0] <+0>:    aesd   [0;36mv0[0m.16b, [0;36mv0[0m.16b
+# CHECK-NEXT: [0x4] <+4>:    bfadd  [0;36mz23[0m.h, [0;36mp3[0m/m, [0;36mz23[0m.h, [0;36mz13[0m.h
+# CHECK-NEXT: [0x8] <+8>:    bfdot  [0;36mv2[0m.2s, [0;36mv3[0m.4h, [0;36mv4[0m.4h
+# CHECK-NEXT: [0xc] <+12>:   brb    iall
+# CHECK-NEXT: [0x10] <+16>:  crc32b [0;36mw0[0m, [0;36mw0[0m, [0;36mw0[0m
+# CHECK-NEXT: [0x14] <+20>:  smin   [0;36mx0[0m, [0;36mx0[0m, [0;31m#0x0[0m
+# CHECK-NEXT: [0x18] <+24>:  sysp   [0;31m#0x0[0m, c2, c0, [0;31m#0x0[0m, [0;36mx0[0m, [0;36mx1[0m
+# CHECK-NEXT: [0x1c] <+28>:  sdot   [0;36mv0[0m.2s, [0;36mv1[0m.8b, [0;36mv2[0m.8b
+# CHECK-NEXT: [0x20] <+32>:  fmmla  [0;36mz0[0m.s, [0;36mz1[0m.s, [0;36mz2[0m.s
