ZJIT: Print local names in FrameState (ruby#14571)

Author: tekknolagi

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);

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")

zjit/src/cruby_bindings.inc.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, "] }}")
     }
 }
 
@@ -4228,6 +4233,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::*;