ZJIT: Improve builtin function annotation collection

Author: st0012

zjit/src/cruby_methods.rs

@@ -15,11 +15,7 @@ use crate::hir_type::{types, Type};
 
 pub struct Annotations {
     cfuncs: HashMap<*mut c_void, FnProperties>,
-    // Builtin annotations by function pointer for fast runtime lookup
-    // Uses Option to cache both hits and misses
-    builtin_ptrs: HashMap<*mut c_void, Option<FnProperties>>,
-    // Temporary name-based annotations used during initialization
-    builtin_names: HashMap<&'static str, FnProperties>,
+    builtin_funcs: HashMap<*mut c_void, FnProperties>,
 }
 
 /// Runtime behaviors of C functions that implement a Ruby method
@@ -48,22 +44,9 @@ impl Annotations {
     }
 
     /// Query about properties of a builtin function by its pointer
-    /// If not found by pointer, checks by name and caches the result (including misses)
-    pub fn get_builtin_properties(&mut self, bf: *const rb_builtin_function) -> Option<FnProperties> {
+    pub fn get_builtin_properties(&self, bf: *const rb_builtin_function) -> Option<FnProperties> {
         let func_ptr = unsafe { (*bf).func_ptr as *mut c_void };
-
-        // First check if we already have it cached by pointer
-        if let Some(&cached_result) = self.builtin_ptrs.get(&func_ptr) {
-            return cached_result;
-        }
-
-        // If not found, check by name and cache the result
-        let name = unsafe { std::ffi::CStr::from_ptr((*bf).name).to_str().ok()? };
-        let result = self.builtin_names.get(name).copied();
-
-        // Cache the result (both hits and misses)
-        self.builtin_ptrs.insert(func_ptr, result);
-        result
+        self.builtin_funcs.get(&func_ptr).copied()
     }
 }
 
@@ -83,11 +66,67 @@ fn annotate_c_method(props_map: &mut HashMap<*mut c_void, FnProperties>, class:
     props_map.insert(fn_ptr, props);
 }
 
+/// Look up a method and find its builtin function pointer by parsing its ISEQ
+/// We currently only support methods with exactly one invokebuiltin instruction
+fn annotate_builtin_method(props_map: &mut HashMap<*mut c_void, FnProperties>, class: VALUE, method_name: &'static str, props: FnProperties) {
+    unsafe {
+        let method_id = rb_intern2(method_name.as_ptr().cast(), method_name.len().try_into().unwrap());
+        let method = rb_method_entry_at(class, method_id);
+        if method.is_null() {
+            panic!("Method {}#{} not found", std::ffi::CStr::from_ptr(rb_class2name(class)).to_str().unwrap_or("?"), method_name);
+        }
+
+        // Cast ME to CME - they have identical layout
+        let cme = method.cast::<rb_callable_method_entry_t>();
+        let def_type = get_cme_def_type(cme);
+
+        if def_type != VM_METHOD_TYPE_ISEQ {
+            panic!("Method {}#{} is not an ISEQ method (type: {})",
+                std::ffi::CStr::from_ptr(rb_class2name(class)).to_str().unwrap_or("?"),
+                method_name, def_type);
+        }
+
+        // Get the ISEQ from the method definition
+        let iseq = get_def_iseq_ptr((*cme).def);
+        if iseq.is_null() {
+            panic!("Failed to get ISEQ for {}#{}",
+                std::ffi::CStr::from_ptr(rb_class2name(class)).to_str().unwrap_or("?"),
+                method_name);
+        }
+
+        // Get the size of the ISEQ in instruction units
+        let encoded_size = rb_iseq_encoded_size(iseq);
+
+        // Scan through the ISEQ to find invokebuiltin instructions
+        let mut insn_idx: u32 = 0;
+        while insn_idx < encoded_size {
+            // Get the PC for this instruction index
+            let pc = rb_iseq_pc_at_idx(iseq, insn_idx);
+
+            // Get the opcode using the proper decoder
+            let opcode = rb_iseq_opcode_at_pc(iseq, pc);
+
+            if opcode == YARVINSN_invokebuiltin as i32 ||
+               opcode == YARVINSN_opt_invokebuiltin_delegate as i32 ||
+               opcode == YARVINSN_opt_invokebuiltin_delegate_leave as i32 {
+                // The first operand is the builtin function pointer
+                let bf_value = *pc.add(1);
+                let bf_ptr = bf_value.as_ptr() as *const rb_builtin_function;
+                let func_ptr = (*bf_ptr).func_ptr as *mut c_void;
+                props_map.insert(func_ptr, props);
+            }
+
+            // Move to the next instruction using the proper length
+            insn_idx = insn_idx.saturating_add(rb_insn_len(VALUE(opcode as usize)).try_into().unwrap());
+        }
+    }
+}
+
 /// Gather annotations. Run this right after boot since the annotations
 /// are about the stock versions of methods.
 pub fn init() -> Annotations {
     let cfuncs = &mut HashMap::new();
-    let builtin_names = &mut HashMap::new();
+    let builtin_funcs = &mut HashMap::new();
 
     macro_rules! annotate {
         ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
@@ -100,18 +139,18 @@ pub fn init() -> Annotations {
     }
 
     macro_rules! annotate_builtin {
-        ($name:literal, $return_type:expr) => {
-            annotate_builtin!($name, $return_type, no_gc, leaf, elidable)
+        ($module:ident, $method_name:literal, $return_type:expr) => {
+            annotate_builtin!($module, $method_name, $return_type, no_gc, leaf, elidable)
         };
-        ($name:literal, $return_type:expr, $($properties:ident),+) => {
+        ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
             let mut props = FnProperties {
                 no_gc: false,
                 leaf: false,
                 elidable: false,
                 return_type: $return_type
             };
             $(props.$properties = true;)+
-            builtin_names.insert($name, props);
+            annotate_builtin_method(builtin_funcs, unsafe { $module }, $method_name, props);
         }
     }
 
@@ -124,13 +163,11 @@ pub fn init() -> Annotations {
     annotate!(rb_cNilClass, "nil?", types::TrueClass, no_gc, leaf, elidable);
     annotate!(rb_mKernel, "nil?", types::FalseClass, no_gc, leaf, elidable);
 
-    // Annotate builtin functions
-    annotate_builtin!("rb_f_float", types::Flonum);
-    annotate_builtin!("rb_f_float1", types::Flonum);
+    annotate_builtin!(rb_mKernel, "Float", types::Flonum);
+    annotate_builtin!(rb_mKernel, "Integer", types::Integer);
 
     Annotations {
         cfuncs: std::mem::take(cfuncs),
-        builtin_ptrs: HashMap::new(), // Cache queried built-in functions by pointer
-        builtin_names: std::mem::take(builtin_names)
+        builtin_funcs: std::mem::take(builtin_funcs),
     }
 }

zjit/src/state.rs

@@ -117,8 +117,8 @@ impl ZJITState {
         &mut ZJITState::get_instance().invariants
     }
 
-    pub fn get_method_annotations() -> &'static mut cruby_methods::Annotations {
-        &mut ZJITState::get_instance().method_annotations
+    pub fn get_method_annotations() -> &'static cruby_methods::Annotations {
+        &ZJITState::get_instance().method_annotations
     }
 
     /// Return true if successful compilation should be asserted