ZJIT: Only print bold text in TTY (#15864)

This is really irritating when piping to a file or pager.

4 files changed, 47 insertions, 10 deletions

diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index f4ed3d7cb7..d8d82a09ca 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -1901,18 +1901,19 @@ impl Assembler
     }
 }
 
-const BOLD_BEGIN: &str = "\x1b[1m";
-const BOLD_END: &str = "\x1b[22m";
-
 /// Return a result of fmt::Display for Assembler without escape sequence
 pub fn lir_string(asm: &Assembler) -> String {
-    format!("{asm}").replace(BOLD_BEGIN, "").replace(BOLD_END, "")
+    use crate::ttycolors::TTY_TERMINAL_COLOR;
+    format!("{asm}").replace(TTY_TERMINAL_COLOR.bold_begin, "").replace(TTY_TERMINAL_COLOR.bold_end, "")
 }
 
 impl fmt::Display for Assembler {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         // Count the number of duplicated label names to disambiguate them if needed
         let mut label_counts: HashMap<&String, usize> = HashMap::new();
+        let colors = crate::ttycolors::get_colors();
+        let bold_begin = colors.bold_begin;
+        let bold_end = colors.bold_end;
         for label_name in self.label_names.iter() {
             let counter = label_counts.entry(label_name).or_insert(0);
             *counter += 1;
@@ -1932,7 +1933,7 @@ impl fmt::Display for Assembler {
         for insn in self.insns.iter() {
             match insn {
                 Insn::Comment(comment) => {
-                    writeln!(f, "    {BOLD_BEGIN}# {comment}{BOLD_END}")?;
+                    writeln!(f, "    {bold_begin}# {comment}{bold_end}")?;
                 }
                 Insn::Label(target) => {
                     let &Target::Label(Label(label_idx)) = target else {
@@ -2506,6 +2507,9 @@ impl AssemblerPanicHook {
 
     /// Dump Assembler, highlighting the insn_idx line
     fn dump_asm(asm: &Assembler, insn_idx: usize) {
+        let colors = crate::ttycolors::get_colors();
+        let bold_begin = colors.bold_begin;
+        let bold_end = colors.bold_end;
         let lir_string = lir_string(asm);
         let lines: Vec<&str> = lir_string.split('\n').collect();
 
@@ -2520,7 +2524,7 @@ impl AssemblerPanicHook {
         eprintln!("Failed to compile LIR at insn_idx={insn_idx}:");
         for (idx, line) in lines.iter().enumerate().filter(|(idx, _)| (min_idx..=max_idx).contains(idx)) {
             if idx == insn_idx && line.starts_with("  ") {
-                eprintln!("{BOLD_BEGIN}=>{}{BOLD_END}", &line["  ".len()..]);
+                eprintln!("{bold_begin}=>{}{bold_end}", &line["  ".len()..]);
             } else {
                 eprintln!("{line}");
             }
diff --git a/zjit/src/disasm.rs b/zjit/src/disasm.rs
index 70a56d780d..d4f6591594 100644
--- a/zjit/src/disasm.rs
+++ b/zjit/src/disasm.rs
@@ -1,8 +1,5 @@
 use crate::asm::CodeBlock;
 
-pub const BOLD_BEGIN: &str = "\x1b[1m";
-pub const BOLD_END: &str = "\x1b[22m";
-
 pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) -> String {
     use std::fmt::Write;
 
@@ -36,12 +33,16 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     let start_addr = 0;
     let insns = cs.disasm_all(code_slice, start_addr as u64).unwrap();
 
+    let colors = crate::ttycolors::get_colors();
+    let bold_begin = colors.bold_begin;
+    let bold_end = colors.bold_end;
+
     // For each instruction in this block
     for insn in insns.as_ref() {
         // Comments for this block
         if let Some(comment_list) = cb.comments_at(insn.address() as usize) {
             for comment in comment_list {
-                writeln!(&mut out, "  {BOLD_BEGIN}# {comment}{BOLD_END}").unwrap();
+                writeln!(&mut out, "  {bold_begin}# {comment}{bold_end}").unwrap();
             }
         }
         writeln!(&mut out, "  {}", format!("{insn}").trim()).unwrap();
diff --git a/zjit/src/lib.rs b/zjit/src/lib.rs
index 88e7af02b7..e7420c16dd 100644
--- a/zjit/src/lib.rs
+++ b/zjit/src/lib.rs
@@ -30,3 +30,4 @@ mod bitset;
 mod gc;
 mod payload;
 mod json;
+mod ttycolors;
diff --git a/zjit/src/ttycolors.rs b/zjit/src/ttycolors.rs
new file mode 100644
index 0000000000..f325772431
--- /dev/null
+++ b/zjit/src/ttycolors.rs
@@ -0,0 +1,31 @@
+use std::io::IsTerminal;
+
+pub fn stdout_supports_colors() -> bool {
+    std::io::stdout().is_terminal()
+}
+
+#[cfg_attr(not(feature = "disasm"), allow(dead_code))]
+#[derive(Copy, Clone, Debug)]
+pub struct TerminalColor {
+    pub bold_begin: &'static str,
+    pub bold_end: &'static str,
+}
+
+pub static TTY_TERMINAL_COLOR: TerminalColor = TerminalColor {
+    bold_begin: "\x1b[1m",
+    bold_end: "\x1b[22m",
+};
+
+pub static NON_TTY_TERMINAL_COLOR: TerminalColor = TerminalColor {
+    bold_begin: "",
+    bold_end: "",
+};
+
+/// Terminal escape codes for colors, font weight, etc. Only enabled if stdout is a TTY.
+pub fn get_colors() -> &'static TerminalColor {
+    if stdout_supports_colors() {
+        &TTY_TERMINAL_COLOR
+    } else {
+        &NON_TTY_TERMINAL_COLOR
+    }
+}