ZJIT: Print local names in FrameState (#14571)

4 files changed, 49 insertions, 3 deletions

diff --git a/zjit.c b/zjit.c
index 565a362a8a..6bbe508f24 100644
--- a/zjit.c
+++ b/zjit.c
@@ -164,6 +164,12 @@ rb_zjit_insn_leaf(int insn, const VALUE *opes)
     return insn_leaf(insn, opes);
 }
 
+ID
+rb_zjit_local_id(const rb_iseq_t *iseq, unsigned idx)
+{
+    return ISEQ_BODY(iseq)->local_table[idx];
+}
+
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key);
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index 7d66bf0ecf..c6f02be415 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -271,6 +271,7 @@ fn main() {
         .allowlist_function("rb_zjit_iseq_builtin_attrs")
         .allowlist_function("rb_zjit_iseq_inspect")
         .allowlist_function("rb_zjit_iseq_insn_set")
+        .allowlist_function("rb_zjit_local_id")
         .allowlist_function("rb_set_cfp_(pc|sp)")
         .allowlist_function("rb_c_method_tracing_currently_enabled")
         .allowlist_function("rb_full_cfunc_return")
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index 9ee4b1bb74..4bb1c3dffd 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -936,6 +936,7 @@ unsafe extern "C" {
     pub fn rb_zjit_singleton_class_p(klass: VALUE) -> bool;
     pub fn rb_zjit_defined_ivar(obj: VALUE, id: ID, pushval: VALUE) -> VALUE;
     pub fn rb_zjit_insn_leaf(insn: ::std::os::raw::c_int, opes: *const VALUE) -> bool;
+    pub fn rb_zjit_local_id(iseq: *const rb_iseq_t, idx: ::std::os::raw::c_uint) -> ID;
     pub fn rb_iseq_encoded_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_iseq_pc_at_idx(iseq: *const rb_iseq_t, insn_idx: u32) -> *mut VALUE;
     pub fn rb_iseq_opcode_at_pc(iseq: *const rb_iseq_t, pc: *const VALUE) -> ::std::os::raw::c_int;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 1a268e62da..033531fd56 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -2903,9 +2903,14 @@ impl Display for FrameStatePrinter<'_> {
         let inner = self.inner;
         write!(f, "FrameState {{ pc: {:?}, stack: ", self.ptr_map.map_ptr(inner.pc))?;
         write_vec(f, &inner.stack)?;
-        write!(f, ", locals: ")?;
-        write_vec(f, &inner.locals)?;
-        write!(f, " }}")
+        write!(f, ", locals: [")?;
+        for (idx, local) in inner.locals.iter().enumerate() {
+            let name: ID = unsafe { rb_zjit_local_id(inner.iseq, idx.try_into().unwrap()) };
+            let name = name.contents_lossy();
+            if idx > 0 { write!(f, ", ")?; }
+            write!(f, "{name}={local}")?;
+        }
+        write!(f, "] }}")
     }
 }
 
@@ -4229,6 +4234,39 @@ mod infer_tests {
 }
 
 #[cfg(test)]
+mod snapshot_tests {
+    use super::*;
+    use insta::assert_snapshot;
+
+    #[track_caller]
+    fn hir_string(method: &str) -> String {
+        let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", method));
+        unsafe { crate::cruby::rb_zjit_profile_disable(iseq) };
+        let function = iseq_to_hir(iseq).unwrap();
+        format!("{}", FunctionPrinter::with_snapshot(&function))
+    }
+
+    #[test]
+    fn test_new_array_with_elements() {
+        eval("def test(a, b) = [a, b]");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:1:
+        bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+          v3:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [a=v1, b=v2] }
+          v4:Any = Snapshot FrameState { pc: 0x1008, stack: [], locals: [a=v1, b=v2] }
+          PatchPoint NoTracePoint
+          v6:Any = Snapshot FrameState { pc: 0x1010, stack: [v1, v2], locals: [a=v1, b=v2] }
+          v7:ArrayExact = NewArray v1, v2
+          v8:Any = Snapshot FrameState { pc: 0x1018, stack: [v7], locals: [a=v1, b=v2] }
+          PatchPoint NoTracePoint
+          v10:Any = Snapshot FrameState { pc: 0x1018, stack: [v7], locals: [a=v1, b=v2] }
+          CheckInterrupts
+          Return v7
+        ");
+    }
+}
+
+#[cfg(test)]
 mod tests {
     use super::*;
     use insta::assert_snapshot;