ZJIT: Add FXT trace disassembler and snapshot tests

Author: tekknolagi

zjit/src/stats.rs

@@ -1051,7 +1051,13 @@ impl PerfettoTracer {
     pub fn new() -> Self {
         let pid = std::process::id();
         let path = format!("/tmp/perfetto-{pid}.fxt");
-        let file = std::fs::File::create(&path)
+        let tracer = Self::create(&path, pid);
+        eprintln!("ZJIT: writing trace exits to {path}");
+        tracer
+    }
+
+    fn create(path: &str, pid: u32) -> Self {
+        let file = std::fs::File::create(path)
             .unwrap_or_else(|e| panic!("ZJIT: failed to create {path}: {e}"));
         let mut tracer = PerfettoTracer {
             writer: std::io::BufWriter::new(file),
@@ -1115,7 +1121,6 @@ impl PerfettoTracer {
             let _ = tracer.writer.flush();
         }
 
-        eprintln!("ZJIT: writing trace exits to {path}");
         tracer
     }
 
@@ -1256,6 +1261,171 @@ impl Drop for PerfettoTracer {
     }
 }
 
+/// Decode an FXT binary trace into human-readable text.
+/// Used by tests to verify trace output. Replaces raw nanosecond timestamps
+/// with sequential labels (t0, t1, ...) so clock-dependent values become
+/// deterministic in snapshots (similar to PtrPrintMap for pointers in HIR).
+#[cfg(test)]
+fn disassemble_fxt(data: &[u8]) -> String {
+    use std::fmt::Write;
+    let mut out = String::new();
+    let mut strings: std::collections::HashMap<u16, String> = std::collections::HashMap::new();
+    strings.insert(0, String::new());
+    let mut ts_map: std::collections::HashMap<u64, usize> = std::collections::HashMap::new();
+    let mut next_ts_id: usize = 0;
+
+    let read_word = |pos: usize| -> u64 {
+        u64::from_le_bytes(data[pos..pos+8].try_into().unwrap())
+    };
+
+    let resolve = |strings: &std::collections::HashMap<u16, String>, idx: u16| -> String {
+        strings.get(&idx).cloned().unwrap_or_else(|| format!("<str:{idx}>"))
+    };
+
+    // Map raw timestamps to deterministic labels (t0, t1, ...)
+    // like PtrPrintMap for pointers in HIR tests.
+    let mut map_ts = |raw: u64| -> String {
+        let id = match ts_map.entry(raw) {
+            std::collections::hash_map::Entry::Occupied(e) => *e.get(),
+            std::collections::hash_map::Entry::Vacant(e) => {
+                let id = next_ts_id;
+                next_ts_id += 1;
+                e.insert(id);
+                id
+            }
+        };
+        format!("t{id}")
+    };
+
+    fn parse_args(data: &[u8], pos: usize, n_args: u64, strings: &std::collections::HashMap<u16, String>) -> (Vec<String>, usize) {
+        let mut args = Vec::new();
+        let mut p = pos;
+        for _ in 0..n_args {
+            let arg_header = u64::from_le_bytes(data[p..p+8].try_into().unwrap());
+            let arg_type = arg_header & 0xF;
+            let arg_size = ((arg_header >> 4) & 0xFFF) as usize;
+            let arg_name_ref = ((arg_header >> 16) & 0xFFFF) as u16;
+            let arg_name = strings.get(&arg_name_ref).cloned().unwrap_or_else(|| format!("<str:{arg_name_ref}>"));
+            match arg_type {
+                4 => { // uint64
+                    let value = u64::from_le_bytes(data[p+8..p+16].try_into().unwrap());
+                    args.push(format!("{arg_name}={value}"));
+                    p += arg_size * 8;
+                }
+                6 => { // string
+                    let value_ref = ((arg_header >> 32) & 0xFFFF) as u16;
+                    let value = strings.get(&value_ref).cloned().unwrap_or_else(|| format!("<str:{value_ref}>"));
+                    args.push(format!("{arg_name}=\"{value}\""));
+                    p += 8;
+                }
+                8 => { // koid
+                    let value = u64::from_le_bytes(data[p+8..p+16].try_into().unwrap());
+                    args.push(format!("{arg_name}={value}"));
+                    p += arg_size * 8;
+                }
+                _ => {
+                    args.push(format!("unknown_arg(type={arg_type})"));
+                    p += arg_size * 8;
+                }
+            }
+        }
+        (args, p)
+    }
+
+    let mut pos = 0;
+    while pos + 8 <= data.len() {
+        let header = read_word(pos);
+        let record_type = header & 0xF;
+        let record_size = ((header >> 4) & 0xFFF) as usize;
+
+        match record_type {
+            0 => {
+                // Metadata record
+                let metadata_type = (header >> 16) & 0xFF;
+                if metadata_type == 4 {
+                    let magic = (header >> 24) & 0xFFFFFFFF;
+                    writeln!(out, "Magic(0x{magic:x})").unwrap();
+                } else {
+                    writeln!(out, "Metadata(type={metadata_type})").unwrap();
+                }
+                pos += 8; // metadata records are always 1 word
+            }
+            1 => {
+                // Initialization record
+                let ticks = read_word(pos + 8);
+                writeln!(out, "Init(ticks_per_second={ticks})").unwrap();
+                pos += record_size * 8;
+            }
+            2 => {
+                // String record
+                let idx = ((header >> 16) & 0x7FFF) as u16;
+                let len = ((header >> 32) & 0x7FFF) as usize;
+                let str_start = pos + 8;
+                let s = std::str::from_utf8(&data[str_start..str_start+len])
+                    .unwrap_or("<invalid utf8>").to_string();
+                strings.insert(idx, s.clone());
+                writeln!(out, "String({idx}, \"{s}\")").unwrap();
+                pos += record_size * 8;
+            }
+            3 => {
+                // Thread record
+                let thread_index = (header >> 16) & 0xFF;
+                let process_koid = read_word(pos + 8);
+                let thread_koid = read_word(pos + 16);
+                writeln!(out, "Thread(index={thread_index}, process_koid={process_koid}, thread_koid={thread_koid})").unwrap();
+                pos += record_size * 8;
+            }
+            4 => {
+                // Event record
+                let event_type = (header >> 16) & 0xF;
+                let n_args = (header >> 20) & 0xF;
+                let category_ref = ((header >> 32) & 0xFFFF) as u16;
+                let name_ref = ((header >> 48) & 0xFFFF) as u16;
+                let category = resolve(&strings, category_ref);
+                let name = resolve(&strings, name_ref);
+                let ts = map_ts(read_word(pos + 8));
+
+                let (args, _) = parse_args(data, pos + 16, n_args, &strings);
+                let args_str = if args.is_empty() { String::new() } else { format!(", {}", args.join(", ")) };
+
+                match event_type {
+                    0 => writeln!(out, "Instant(category=\"{category}\", name=\"{name}\", ts={ts}{args_str})").unwrap(),
+                    1 => {
+                        // Counter: has counter_id after args
+                        let counter_id = read_word(pos + 16 + n_args as usize * 8 + 8); // after uint64 arg
+                        writeln!(out, "Counter(category=\"{category}\", name=\"{name}\", ts={ts}{args_str}, counter_id={counter_id})").unwrap();
+                    }
+                    2 => writeln!(out, "DurationBegin(category=\"{category}\", name=\"{name}\", ts={ts}{args_str})").unwrap(),
+                    3 => writeln!(out, "DurationEnd(category=\"{category}\", name=\"{name}\", ts={ts}{args_str})").unwrap(),
+                    _ => writeln!(out, "Event(type={event_type}, category=\"{category}\", name=\"{name}\", ts={ts}{args_str})").unwrap(),
+                }
+                pos += record_size * 8;
+            }
+            7 => {
+                // Kernel object record
+                let obj_type = (header >> 16) & 0xFF;
+                let name_ref = ((header >> 24) & 0xFFFF) as u16;
+                let n_args = (header >> 40) & 0xF;
+                let name = resolve(&strings, name_ref);
+                let koid = read_word(pos + 8);
+
+                let type_name = match obj_type { 1 => "process", 2 => "thread", _ => "unknown" };
+                let (args, _) = parse_args(data, pos + 16, n_args, &strings);
+                let args_str = if args.is_empty() { String::new() } else { format!(", {}", args.join(", ")) };
+
+                writeln!(out, "KernelObject(type={type_name}, name=\"{name}\", koid={koid}{args_str})").unwrap();
+                pos += record_size * 8;
+            }
+            _ => {
+                writeln!(out, "Unknown(type={record_type}, size={record_size})").unwrap();
+                if record_size == 0 { break; } // avoid infinite loop
+                pos += record_size * 8;
+            }
+        }
+    }
+    out
+}
+
 /// Primitive called in zjit.rb
 ///
 /// Check if trace_exits generation is enabled.
@@ -1269,3 +1439,108 @@ pub extern "C" fn rb_zjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self:
     }
 }
 
+#[cfg(test)]
+mod fxt_tests {
+    use super::*;
+    use insta::assert_snapshot;
+
+    fn create_test_tracer() -> (PerfettoTracer, std::path::PathBuf) {
+        let dir = std::env::temp_dir();
+        let path = dir.join(format!("zjit_fxt_test_{}.fxt", std::process::id()));
+        let tracer = PerfettoTracer::create(path.to_str().unwrap(), 42);
+        (tracer, path)
+    }
+
+    fn read_and_disassemble(tracer: PerfettoTracer, path: &std::path::Path) -> String {
+        drop(tracer); // flush via Drop
+        let data = std::fs::read(path).unwrap();
+        let _ = std::fs::remove_file(path);
+        disassemble_fxt(&data)
+    }
+
+    /// Strip the pre-interned strings from header output to keep snapshots focused
+    fn strip_header(output: &str) -> String {
+        output.lines()
+            .skip_while(|l| !l.starts_with("Duration") && !l.starts_with("Instant") && !l.starts_with("Counter"))
+            .collect::<Vec<_>>()
+            .join("\n")
+    }
+
+    #[test]
+    fn test_fxt_header() {
+        let (tracer, path) = create_test_tracer();
+        assert_snapshot!(read_and_disassemble(tracer, &path), @r#"
+        Magic(0x16547846)
+        Init(ticks_per_second=1000000000)
+        Thread(index=1, process_koid=42, thread_koid=1)
+        String(1, "ruby")
+        KernelObject(type=process, name="ruby", koid=42)
+        String(2, "main")
+        String(3, "process")
+        KernelObject(type=thread, name="main", koid=1, process=42)
+        String(4, "side_exit")
+        String(5, "compile")
+        String(6, "invalidation")
+        String(7, "0")
+        String(8, "1")
+        String(9, "2")
+        String(10, "3")
+        String(11, "4")
+        String(12, "5")
+        String(13, "6")
+        String(14, "7")
+        String(15, "8")
+        String(16, "9")
+        String(17, "10")
+        String(18, "11")
+        String(19, "12")
+        String(20, "13")
+        String(21, "14")
+        "#);
+    }
+
+    #[test]
+    fn test_fxt_duration_events() {
+        let (mut tracer, path) = create_test_tracer();
+        tracer.write_duration_begin("compile", "test_fn", 1000, &[]);
+        tracer.write_duration_end("compile", "test_fn", 2000);
+        assert_snapshot!(strip_header(&read_and_disassemble(tracer, &path)), @r#"
+        DurationBegin(category="compile", name="test_fn", ts=t0)
+        DurationEnd(category="compile", name="test_fn", ts=t1)
+        "#);
+    }
+
+    #[test]
+    fn test_fxt_duration_with_frames() {
+        let (mut tracer, path) = create_test_tracer();
+        tracer.write_duration_begin("invalidation", "bop", 500, &["frame_a".into(), "frame_b".into()]);
+        tracer.write_duration_end("invalidation", "bop", 600);
+        assert_snapshot!(strip_header(&read_and_disassemble(tracer, &path)), @r#"
+        DurationBegin(category="invalidation", name="bop", ts=t0, 0="frame_a", 1="frame_b")
+        DurationEnd(category="invalidation", name="bop", ts=t1)
+        "#);
+    }
+
+    #[test]
+    fn test_fxt_instant_event() {
+        let (mut tracer, path) = create_test_tracer();
+        tracer.write_event("side_exit", "guard_failed", &["caller_fn".into()]);
+        assert_snapshot!(strip_header(&read_and_disassemble(tracer, &path)), @r#"Instant(category="side_exit", name="guard_failed", ts=t0, 0="caller_fn")"#);
+    }
+
+    #[test]
+    fn test_fxt_string_dedup() {
+        let (mut tracer, path) = create_test_tracer();
+        tracer.write_duration_begin("compile", "same_fn", 100, &[]);
+        tracer.write_duration_end("compile", "same_fn", 200);
+        tracer.write_duration_begin("compile", "same_fn", 300, &[]);
+        tracer.write_duration_end("compile", "same_fn", 400);
+        assert_snapshot!(strip_header(&read_and_disassemble(tracer, &path)), @r#"
+        DurationBegin(category="compile", name="same_fn", ts=t0)
+        DurationEnd(category="compile", name="same_fn", ts=t1)
+        DurationBegin(category="compile", name="same_fn", ts=t2)
+        DurationEnd(category="compile", name="same_fn", ts=t3)
+        "#);
+    }
+}
+