YJIT: Only enable disassembly colors for tty

yjit/src/disasm.rs

@@ -3,7 +3,7 @@ use crate::cruby::*;
 use crate::yjit::yjit_enabled_p;
 use crate::asm::CodeBlock;
 use crate::codegen::CodePtr;
-use crate::options::DumpDisasm;
+use crate::options::{DumpDisasm, get_option};
 
 use std::fmt::Write;
 
@@ -158,24 +158,25 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     #[cfg(test)]
     let start_addr = 0;
     let insns = cs.disasm_all(code_slice, start_addr as u64).unwrap();
+    let color = get_option!(color);
 
     // 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 {
                 if cb.outlined {
-                    write!(&mut out, "\x1b[34m").unwrap(); // Make outlined code blue
+                    write!(&mut out, "{}", color.blue_begin).unwrap(); // Make outlined code blue
                 }
-                writeln!(&mut out, "  \x1b[1m# {comment}\x1b[22m").unwrap(); // Make comments bold
+                writeln!(&mut out, "  {}# {comment}{}", color.bold_begin, color.bold_end).unwrap(); // Make comments bold
             }
         }
         if cb.outlined {
-            write!(&mut out, "\x1b[34m").unwrap(); // Make outlined code blue
+            write!(&mut out, "{}", color.blue_begin).unwrap(); // Make outlined code blue
         }
         writeln!(&mut out, "  {insn}").unwrap();
         if cb.outlined {
-            write!(&mut out, "\x1b[0m").unwrap(); // Disable blue
+            write!(&mut out, "{}", color.blue_end).unwrap(); // Disable blue
         }
     }
 
@@ -188,6 +189,7 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     let mut out = String::new();
     let mut line_byte_idx = 0;
     const MAX_BYTES_PER_LINE: usize = 16;
+    let color = get_option!(color);
 
     for addr in start_addr..end_addr {
         if let Some(comment_list) = cb.comments_at(addr) {
@@ -197,7 +199,7 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
                 line_byte_idx = 0;
             }
             for comment in comment_list {
-                writeln!(&mut out, "  \x1b[1m# {comment}\x1b[22m").unwrap(); // Make comments bold
+                writeln!(&mut out, "  {}# {comment}{}", color.bold_begin, color.bold_end).unwrap(); // Make comments bold
             }
         }
         if line_byte_idx == 0 {

yjit/src/options.rs

@@ -23,6 +23,14 @@ pub static mut rb_yjit_call_threshold: u64 = SMALL_CALL_THRESHOLD;
 #[no_mangle]
 pub static mut rb_yjit_cold_threshold: u64 = 200_000;
 
+#[derive(Copy, Clone, Debug)]
+pub struct TerminalColor {
+    pub blue_begin: &'static str,
+    pub blue_end: &'static str,
+    pub bold_begin: &'static str,
+    pub bold_end: &'static str,
+}
+
 // Command-line options
 #[derive(Debug)]
 #[repr(C)]
@@ -88,8 +96,24 @@ pub struct Options {
 
     // Where to store the log. `None` disables the log.
     pub log: Option<LogOutput>,
+
+    pub color: TerminalColor,
 }
 
+pub static TTY_TERMINAL_COLOR: TerminalColor = TerminalColor {
+    blue_begin: "\x1b[34m",
+    blue_end: "\x1b[0m",
+    bold_begin: "\x1b[1m",
+    bold_end: "\x1b[22m",
+};
+
+pub static NON_TTY_TERMINAL_COLOR: TerminalColor = TerminalColor {
+    blue_begin: "",
+    blue_end: "",
+    bold_begin: "",
+    bold_end: "",
+};
+
 // Initialize the options to default values
 pub static mut OPTIONS: Options = Options {
     mem_size: 128 * 1024 * 1024,
@@ -111,6 +135,7 @@ pub static mut OPTIONS: Options = Options {
     code_gc: false,
     perf_map: None,
     log: None,
+    color: NON_TTY_TERMINAL_COLOR,
 };
 
 /// YJIT option descriptions for `ruby --help`.
@@ -300,7 +325,15 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
 
             match opt_val {
-                "" => unsafe { OPTIONS.dump_disasm = Some(DumpDisasm::Stdout) },
+                "" => {
+                    unsafe { OPTIONS.dump_disasm = Some(DumpDisasm::Stdout) };
+                    extern "C" { fn isatty(fd: c_int) -> c_int; }
+                    let stdout = 1;
+                    let is_terminal = unsafe { isatty(stdout) } != 0;
+                    if is_terminal {
+                        unsafe { OPTIONS.color = TTY_TERMINAL_COLOR };
+                    }
+                },
                 directory => {
                     let path = format!("{directory}/yjit_{}.log", std::process::id());
                     match File::options().create(true).append(true).open(&path) {