diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 55041945d4..f2f990afde 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -843,6 +843,22 @@ impl<'a> FunctionPrinter<'a> {
}
}
+/// Pretty printer for [`Function`].
+pub struct FunctionGraphvizPrinter<'a> {
+ fun: &'a Function,
+ ptr_map: PtrPrintMap,
+}
+
+impl<'a> FunctionGraphvizPrinter<'a> {
+ pub fn new(fun: &'a Function) -> Self {
+ let mut ptr_map = PtrPrintMap::identity();
+ if cfg!(test) {
+ ptr_map.map_ptrs = true;
+ }
+ Self { fun, ptr_map }
+ }
+}
+
/// Union-Find (Disjoint-Set) is a data structure for managing disjoint sets that has an interface
/// of two operations:
///
@@ -2115,6 +2131,10 @@ impl Function {
Some(DumpHIR::Debug) => println!("Optimized HIR:\n{:#?}", &self),
None => {},
}
+
+ if get_option!(dump_hir_graphviz) {
+ println!("{}", FunctionGraphvizPrinter::new(&self));
+ }
}
@@ -2293,6 +2313,87 @@ impl<'a> std::fmt::Display for FunctionPrinter<'a> {
}
}
+struct HtmlEncoder<'a, 'b> {
+ formatter: &'a mut std::fmt::Formatter<'b>,
+}
+
+impl<'a, 'b> std::fmt::Write for HtmlEncoder<'a, 'b> {
+ fn write_str(&mut self, s: &str) -> std::fmt::Result {
+ for ch in s.chars() {
+ match ch {
+ '<' => self.formatter.write_str("<")?,
+ '>' => self.formatter.write_str(">")?,
+ '&' => self.formatter.write_str("&")?,
+ '"' => self.formatter.write_str(""")?,
+ '\'' => self.formatter.write_str("'")?,
+ _ => self.formatter.write_char(ch)?,
+ }
+ }
+ Ok(())
+ }
+}
+
+impl<'a> std::fmt::Display for FunctionGraphvizPrinter<'a> {
+ fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ macro_rules! write_encoded {
+ ($f:ident, $($arg:tt)*) => {
+ HtmlEncoder { formatter: $f }.write_fmt(format_args!($($arg)*))
+ };
+ }
+ use std::fmt::Write;
+ let fun = &self.fun;
+ let iseq_name = iseq_get_location(fun.iseq, 0);
+ write!(f, "digraph G {{ # ")?;
+ write_encoded!(f, "{iseq_name}")?;
+ write!(f, "\n")?;
+ writeln!(f, "node [shape=plaintext];")?;
+ writeln!(f, "mode=hier; overlap=false; splines=true;")?;
+ for block_id in fun.rpo() {
+ writeln!(f, r#" {block_id} [label=<
"#)?;
+ write!(f, r#"{block_id}("#)?;
+ if !fun.blocks[block_id.0].params.is_empty() {
+ let mut sep = "";
+ for param in &fun.blocks[block_id.0].params {
+ write_encoded!(f, "{sep}{param}")?;
+ let insn_type = fun.type_of(*param);
+ if !insn_type.is_subtype(types::Empty) {
+ write_encoded!(f, ":{}", insn_type.print(&self.ptr_map))?;
+ }
+ sep = ", ";
+ }
+ }
+ let mut edges = vec![];
+ writeln!(f, ") |
")?;
+ for insn_id in &fun.blocks[block_id.0].insns {
+ let insn_id = fun.union_find.borrow().find_const(*insn_id);
+ let insn = fun.find(insn_id);
+ if matches!(insn, Insn::Snapshot {..}) {
+ continue;
+ }
+ write!(f, r#""#)?;
+ if insn.has_output() {
+ let insn_type = fun.type_of(insn_id);
+ if insn_type.is_subtype(types::Empty) {
+ write_encoded!(f, "{insn_id} = ")?;
+ } else {
+ write_encoded!(f, "{insn_id}:{} = ", insn_type.print(&self.ptr_map))?;
+ }
+ }
+ if let Insn::Jump(ref target) | Insn::IfTrue { ref target, .. } | Insn::IfFalse { ref target, .. } = insn {
+ edges.push((insn_id, target.target));
+ }
+ write_encoded!(f, "{}", insn.print(&self.ptr_map))?;
+ writeln!(f, " |
")?;
+ }
+ writeln!(f, "
>];")?;
+ for (src, dst) in edges {
+ writeln!(f, " {block_id}:{src} -> {dst}:params;")?;
+ }
+ }
+ writeln!(f, "}}")
+ }
+}
+
#[derive(Debug, Clone, PartialEq)]
pub struct FrameState {
iseq: IseqPtr,
@@ -5145,6 +5246,81 @@ mod tests {
}
}
+#[cfg(test)]
+mod graphviz_tests {
+ use super::*;
+ use expect_test::{expect, Expect};
+
+ #[track_caller]
+ fn assert_optimized_graphviz(method: &str, expected: Expect) {
+ let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", method));
+ unsafe { crate::cruby::rb_zjit_profile_disable(iseq) };
+ let mut function = iseq_to_hir(iseq).unwrap();
+ function.optimize();
+ function.validate().unwrap();
+ let actual = format!("{}", FunctionGraphvizPrinter::new(&function));
+ expected.assert_eq(&actual);
+ }
+
+ #[test]
+ fn test_guard_fixnum_or_fixnum() {
+ eval(r#"
+ def test(x, y) = x | y
+
+ test(1, 2)
+ "#);
+ assert_optimized_graphviz("test", expect![[r#"
+ digraph G { # test@<compiled>:2
+ node [shape=plaintext];
+ mode=hier; overlap=false; splines=true;
+ bb0 [label=<
+ bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject) |
+ PatchPoint BOPRedefined(INTEGER_REDEFINED_OP_FLAG, 29) |
+ v8:Fixnum = GuardType v1, Fixnum |
+ v9:Fixnum = GuardType v2, Fixnum |
+ v10:Fixnum = FixnumOr v8, v9 |
+ Return v10 |
+
>];
+ }
+ "#]]);
+ }
+
+ #[test]
+ fn test_multiple_blocks() {
+ eval(r#"
+ def test(c)
+ if c
+ 3
+ else
+ 4
+ end
+ end
+
+ test(1)
+ test("x")
+ "#);
+ assert_optimized_graphviz("test", expect![[r#"
+ digraph G { # test@<compiled>:3
+ node [shape=plaintext];
+ mode=hier; overlap=false; splines=true;
+ bb0 [label=<
+ bb0(v0:BasicObject, v1:BasicObject) |
+ v3:CBool = Test v1 |
+ IfFalse v3, bb1(v0, v1) |
+ v5:Fixnum[3] = Const Value(3) |
+ Return v5 |
+
>];
+ bb0:v4 -> bb1:params;
+ bb1 [label=<
+ bb1(v7:BasicObject, v8:BasicObject) |
+ v10:Fixnum[4] = Const Value(4) |
+ Return v10 |
+
>];
+ }
+ "#]]);
+ }
+}
+
#[cfg(test)]
mod opt_tests {
use super::*;
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
index 340812f089..92f56b8916 100644
--- a/zjit/src/options.rs
+++ b/zjit/src/options.rs
@@ -37,6 +37,8 @@ pub struct Options {
/// Dump High-level IR after optimization, right before codegen.
pub dump_hir_opt: Option,
+ pub dump_hir_graphviz: bool,
+
/// Dump low-level IR
pub dump_lir: bool,
@@ -61,6 +63,7 @@ impl Default for Options {
debug: false,
dump_hir_init: None,
dump_hir_opt: None,
+ dump_hir_graphviz: false,
dump_lir: false,
dump_disasm: false,
perf: false,
@@ -186,6 +189,7 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
("dump-hir" | "dump-hir-opt", "") => options.dump_hir_opt = Some(DumpHIR::WithoutSnapshot),
("dump-hir" | "dump-hir-opt", "all") => options.dump_hir_opt = Some(DumpHIR::All),
("dump-hir" | "dump-hir-opt", "debug") => options.dump_hir_opt = Some(DumpHIR::Debug),
+ ("dump-hir-graphviz", "") => options.dump_hir_graphviz = true,
("dump-hir-init", "") => options.dump_hir_init = Some(DumpHIR::WithoutSnapshot),
("dump-hir-init", "all") => options.dump_hir_init = Some(DumpHIR::All),