From 8010d79bb429ed8dac5244e481ba5ddc108797a2 Mon Sep 17 00:00:00 2001
From: Max Bernstein <max.bernstein@shopify.com>
Date: Mon, 9 Dec 2024 07:36:17 -0800
Subject: [PATCH] YJIT: Only enable disassembly colors for tty (#12283)

* YJIT: Use fully-qualified name for OPTIONS in get_options!

* YJIT: Only enable disassembly colors for tty
---
 yjit/src/disasm.rs  | 43 ++++++++++++++++++++++++++++++++++++++-----
 yjit/src/options.rs |  2 +-
 yjit/src/utils.rs   |  9 +++++++++
 3 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
index f98a3645fb..a505973331 100644
--- a/yjit/src/disasm.rs
+++ b/yjit/src/disasm.rs
@@ -7,6 +7,37 @@ use crate::options::DumpDisasm;
 
 use std::fmt::Write;
 
+#[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,
+}
+
+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: "",
+};
+
+/// Terminal escape codes for colors, font weight, etc. Only enabled if stdout is a TTY.
+pub fn get_colors() -> &'static TerminalColor {
+    if crate::utils::stdout_supports_colors() {
+        &TTY_TERMINAL_COLOR
+    } else {
+        &NON_TTY_TERMINAL_COLOR
+    }
+}
+
 /// Primitive called in yjit.rb
 /// Produce a string representing the disassembly for an ISEQ
 #[no_mangle]
@@ -158,6 +189,7 @@ 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 colors = get_colors();
 
     // For each instruction in this block
     for insn in insns.as_ref() {
@@ -165,17 +197,17 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
         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, "{}", colors.blue_begin).unwrap(); // Make outlined code blue
                 }
-                writeln!(&mut out, "  \x1b[1m# {comment}\x1b[22m").unwrap(); // Make comments bold
+                writeln!(&mut out, "  {}# {comment}{}", colors.bold_begin, colors.bold_end).unwrap(); // Make comments bold
             }
         }
         if cb.outlined {
-            write!(&mut out, "\x1b[34m").unwrap(); // Make outlined code blue
+            write!(&mut out, "{}", colors.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, "{}", colors.blue_end).unwrap(); // Disable blue
         }
     }
 
@@ -188,6 +220,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 colors = get_colors();
 
     for addr in start_addr..end_addr {
         if let Some(comment_list) = cb.comments_at(addr) {
@@ -197,7 +230,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}{}", colors.bold_begin, colors.bold_end).unwrap(); // Make comments bold
             }
         }
         if line_byte_idx == 0 {
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index bd1882dd0b..3fcc40711c 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -173,7 +173,7 @@ macro_rules! get_option {
         {
             // Make this a statement since attributes on expressions are experimental
             #[allow(unused_unsafe)]
-            let ret = unsafe { OPTIONS.$option_name };
+            let ret = unsafe { crate::options::OPTIONS.$option_name };
             ret
         }
     };
diff --git a/yjit/src/utils.rs b/yjit/src/utils.rs
index c4b5fbd2e7..8c4133546d 100644
--- a/yjit/src/utils.rs
+++ b/yjit/src/utils.rs
@@ -3,6 +3,7 @@
 use crate::backend::ir::*;
 use crate::cruby::*;
 use std::slice;
+use std::os::raw::c_int;
 
 /// Trait for casting to [usize] that allows you to say `.as_usize()`.
 /// Implementation conditional on the cast preserving the numeric value on
@@ -239,6 +240,14 @@ pub fn print_str(asm: &mut Assembler, str: &str) {
     asm.cpop_all();
 }
 
+pub fn stdout_supports_colors() -> bool {
+    // TODO(max): Use std::io::IsTerminal after upgrading Rust to 1.70
+    extern "C" { fn isatty(fd: c_int) -> c_int; }
+    let stdout = 1;
+    let is_terminal = unsafe { isatty(stdout) } == 1;
+    is_terminal
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;