ZJIT: Fetch Primitive.attr!(leaf) for InvokeBuiltin

Fix https://github.com/Shopify/ruby/issues/670

9 files changed, 76 insertions, 15 deletions

diff --git a/jit.c b/jit.c
index 3111dcc3e3..43c932e5a0 100644
--- a/jit.c
+++ b/jit.c
@@ -181,6 +181,12 @@ rb_jit_get_proc_ptr(VALUE procv)
     return proc;
 }
 
+unsigned int
+rb_jit_iseq_builtin_attrs(const rb_iseq_t *iseq)
+{
+    return iseq->body->builtin_attrs;
+}
+
 int
 rb_get_mct_argc(const rb_method_cfunc_t *mct)
 {
diff --git a/yjit.c b/yjit.c
index 3793b0f1ac..4b78cfbae2 100644
--- a/yjit.c
+++ b/yjit.c
@@ -244,12 +244,6 @@ rb_optimized_call(VALUE *recv, rb_execution_context_t *ec, int argc, VALUE *argv
     return rb_vm_invoke_proc(ec, proc, argc, argv, kw_splat, block_handler);
 }
 
-unsigned int
-rb_yjit_iseq_builtin_attrs(const rb_iseq_t *iseq)
-{
-    return iseq->body->builtin_attrs;
-}
-
 // If true, the iseq has only opt_invokebuiltin_delegate(_leave) and leave insns.
 static bool
 invokebuiltin_delegate_leave_p(const rb_iseq_t *iseq)
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index df287e1bf8..100abbb33f 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -249,7 +249,7 @@ fn main() {
         .allowlist_function("rb_jit_mark_executable")
         .allowlist_function("rb_jit_mark_unused")
         .allowlist_function("rb_jit_get_page_size")
-        .allowlist_function("rb_yjit_iseq_builtin_attrs")
+        .allowlist_function("rb_jit_iseq_builtin_attrs")
         .allowlist_function("rb_yjit_iseq_inspect")
         .allowlist_function("rb_yjit_builtin_function")
         .allowlist_function("rb_set_cfp_(pc|sp)")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 2eb66a298e..3f6f1bb46e 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -7694,7 +7694,7 @@ fn gen_send_iseq(
     gen_counter_incr(jit, asm, Counter::num_send_iseq);
 
     // Shortcut for special `Primitive.attr! :leaf` builtins
-    let builtin_attrs = unsafe { rb_yjit_iseq_builtin_attrs(iseq) };
+    let builtin_attrs = unsafe { rb_jit_iseq_builtin_attrs(iseq) };
     let builtin_func_raw = unsafe { rb_yjit_builtin_function(iseq) };
     let builtin_func = if builtin_func_raw.is_null() { None } else { Some(builtin_func_raw) };
     let opt_send_call = flags & VM_CALL_OPT_SEND != 0; // .send call is not currently supported for builtins
@@ -9635,7 +9635,7 @@ fn gen_invokeblock_specialized(
 
         // If the current ISEQ is annotated to be inlined but it's not being inlined here,
         // generate a dynamic dispatch to avoid making this yield megamorphic.
-        if unsafe { rb_yjit_iseq_builtin_attrs(jit.iseq) } & BUILTIN_ATTR_INLINE_BLOCK != 0 && !asm.ctx.inline() {
+        if unsafe { rb_jit_iseq_builtin_attrs(jit.iseq) } & BUILTIN_ATTR_INLINE_BLOCK != 0 && !asm.ctx.inline() {
             gen_counter_incr(jit, asm, Counter::invokeblock_iseq_not_inlined);
             return None;
         }
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 6542e5ef09..f6e4e0e22f 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1130,7 +1130,6 @@ extern "C" {
         kw_splat: ::std::os::raw::c_int,
         block_handler: VALUE,
     ) -> VALUE;
-    pub fn rb_yjit_iseq_builtin_attrs(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_yjit_builtin_function(iseq: *const rb_iseq_t) -> *const rb_builtin_function;
     pub fn rb_yjit_str_simple_append(str1: VALUE, str2: VALUE) -> VALUE;
     pub fn rb_vm_base_ptr(cfp: *mut rb_control_frame_struct) -> *mut VALUE;
@@ -1198,6 +1197,7 @@ extern "C" {
     pub fn rb_get_def_original_id(def: *const rb_method_definition_t) -> ID;
     pub fn rb_get_def_bmethod_proc(def: *mut rb_method_definition_t) -> VALUE;
     pub fn rb_jit_get_proc_ptr(procv: VALUE) -> *mut rb_proc_t;
+    pub fn rb_jit_iseq_builtin_attrs(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_get_mct_argc(mct: *const rb_method_cfunc_t) -> ::std::os::raw::c_int;
     pub fn rb_get_mct_func(mct: *const rb_method_cfunc_t) -> *mut ::std::os::raw::c_void;
     pub fn rb_get_def_iseq_ptr(def: *mut rb_method_definition_t) -> *const rb_iseq_t;
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index 92f7a10e56..75dbd46794 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -274,7 +274,7 @@ fn main() {
         .allowlist_function("rb_jit_mark_unused")
         .allowlist_function("rb_jit_get_page_size")
         .allowlist_function("rb_jit_array_len")
-        .allowlist_function("rb_zjit_iseq_builtin_attrs")
+        .allowlist_function("rb_jit_iseq_builtin_attrs")
         .allowlist_function("rb_zjit_iseq_inspect")
         .allowlist_function("rb_zjit_iseq_insn_set")
         .allowlist_function("rb_zjit_local_id")
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index 41e0e847aa..d4e4079b5c 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -1175,6 +1175,11 @@ pub mod test_utils {
         get_proc_iseq(&format!("{}.method(:{})", recv, name))
     }
 
+    /// Get IseqPtr for a specified instance method
+    pub fn get_instance_method_iseq(recv: &str, name: &str) -> *const rb_iseq_t {
+        get_proc_iseq(&format!("{}.instance_method(:{})", recv, name))
+    }
+
     /// Get IseqPtr for a specified Proc object
     pub fn get_proc_iseq(obj: &str) -> *const rb_iseq_t {
         let wrapped_iseq = eval(&format!("RubyVM::InstructionSequence.of({obj})"));
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index d9bd2d33c0..f7e6cdde94 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -1363,6 +1363,7 @@ unsafe extern "C" {
     pub fn rb_get_def_original_id(def: *const rb_method_definition_t) -> ID;
     pub fn rb_get_def_bmethod_proc(def: *mut rb_method_definition_t) -> VALUE;
     pub fn rb_jit_get_proc_ptr(procv: VALUE) -> *mut rb_proc_t;
+    pub fn rb_jit_iseq_builtin_attrs(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_get_mct_argc(mct: *const rb_method_cfunc_t) -> ::std::os::raw::c_int;
     pub fn rb_get_mct_func(mct: *const rb_method_cfunc_t) -> *mut ::std::os::raw::c_void;
     pub fn rb_get_def_iseq_ptr(def: *mut rb_method_definition_t) -> *const rb_iseq_t;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 489ea83a44..834a33d23c 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -745,6 +745,7 @@ pub enum Insn {
         bf: rb_builtin_function,
         args: Vec<InsnId>,
         state: InsnId,
+        leaf: bool,
         return_type: Option<Type>,  // None for unannotated builtins
     },
 
@@ -1039,8 +1040,10 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
-            Insn::InvokeBuiltin { bf, args, .. } => {
-                write!(f, "InvokeBuiltin {}", unsafe { CStr::from_ptr(bf.name) }.to_str().unwrap())?;
+            Insn::InvokeBuiltin { bf, args, leaf, .. } => {
+                write!(f, "InvokeBuiltin{} {}",
+                           if *leaf { " leaf" } else { "" },
+                           unsafe { CStr::from_ptr(bf.name) }.to_str().unwrap())?;
                 for arg in args {
                     write!(f, ", {arg}")?;
                 }
@@ -1678,7 +1681,7 @@ impl Function {
                 state,
                 reason,
             },
-            &InvokeBuiltin { bf, ref args, state, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, return_type },
+            &InvokeBuiltin { bf, ref args, state, leaf, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, leaf, return_type },
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
             &HashAref { hash, key, state } => HashAref { hash: find!(hash), key: find!(key), state },
@@ -4671,10 +4674,14 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         .get_builtin_properties(&bf)
                         .map(|props| props.return_type);
 
+                    let builtin_attrs = unsafe { rb_jit_iseq_builtin_attrs(iseq) };
+                    let leaf = builtin_attrs & BUILTIN_ATTR_LEAF != 0;
+
                     let insn_id = fun.push_insn(block, Insn::InvokeBuiltin {
                         bf,
                         args,
                         state: exit_id,
+                        leaf,
                         return_type,
                     });
                     state.stack_push(insn_id);
@@ -4697,10 +4704,14 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         .get_builtin_properties(&bf)
                         .map(|props| props.return_type);
 
+                    let builtin_attrs = unsafe { rb_jit_iseq_builtin_attrs(iseq) };
+                    let leaf = builtin_attrs & BUILTIN_ATTR_LEAF != 0;
+
                     let insn_id = fun.push_insn(block, Insn::InvokeBuiltin {
                         bf,
                         args,
                         state: exit_id,
+                        leaf,
                         return_type,
                     });
                     state.stack_push(insn_id);
@@ -7928,7 +7939,7 @@ mod tests {
           EntryPoint JIT(0)
           Jump bb2(v4)
         bb2(v6:BasicObject):
-          v11:Class = InvokeBuiltin _bi20, v6
+          v11:Class = InvokeBuiltin leaf _bi20, v6
           Jump bb3(v6, v11)
         bb3(v13:BasicObject, v14:Class):
           CheckInterrupts
@@ -8027,6 +8038,50 @@ mod tests {
     }
 
     #[test]
+    fn test_invoke_leaf_builtin_symbol_name() {
+        let iseq = crate::cruby::with_rubyvm(|| get_instance_method_iseq("Symbol", "name"));
+        let function = iseq_to_hir(iseq).unwrap();
+        assert_snapshot!(hir_string_function(&function), @r"
+        fn name@<internal:symbol>:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v11:BasicObject = InvokeBuiltin leaf _bi28, v6
+          Jump bb3(v6, v11)
+        bb3(v13:BasicObject, v14:BasicObject):
+          CheckInterrupts
+          Return v14
+        ");
+    }
+
+    #[test]
+    fn test_invoke_leaf_builtin_symbol_to_s() {
+        let iseq = crate::cruby::with_rubyvm(|| get_instance_method_iseq("Symbol", "to_s"));
+        let function = iseq_to_hir(iseq).unwrap();
+        assert_snapshot!(hir_string_function(&function), @r"
+        fn to_s@<internal:symbol>:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v11:BasicObject = InvokeBuiltin leaf _bi12, v6
+          Jump bb3(v6, v11)
+        bb3(v13:BasicObject, v14:BasicObject):
+          CheckInterrupts
+          Return v14
+        ");
+    }
+
+    #[test]
     fn dupn() {
         eval("
             def test(x) = (x[0, 1] ||= 2)