ZJIT: Annotate builtin functions

Author: st0012

zjit/src/codegen.rs

@@ -337,6 +337,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args))?,
         Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), opnds!(args), &function.frame_state(*state))?,
         Insn::InvokeBuiltin { bf, args, state } => gen_invokebuiltin(asm, &function.frame_state(*state), bf, opnds!(args))?,
+        Insn::BuiltinCall { bf, args, state, return_type: _ } => gen_invokebuiltin(asm, &function.frame_state(*state), bf, opnds!(args))?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
         Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,

zjit/src/cruby_methods.rs

@@ -15,6 +15,7 @@ use crate::hir_type::{types, Type};
 
 pub struct Annotations {
     cfuncs: HashMap<*mut c_void, FnProperties>,
+    builtins: HashMap<&'static str, FnProperties>,
 }
 
 /// Runtime behaviors of C functions that implement a Ruby method
@@ -41,6 +42,11 @@ impl Annotations {
         };
         self.cfuncs.get(&fn_ptr).copied()
     }
+
+    /// Query about properties of a builtin function
+    pub fn get_builtin_properties(&self, name: &str) -> Option<FnProperties> {
+        self.builtins.get(&name).copied()
+    }
 }
 
 fn annotate_c_method(props_map: &mut HashMap<*mut c_void, FnProperties>, class: VALUE, method_name: &'static str, props: FnProperties) {
@@ -63,6 +69,7 @@ fn annotate_c_method(props_map: &mut HashMap<*mut c_void, FnProperties>, class:
 /// are about the stock versions of methods.
 pub fn init() -> Annotations {
     let cfuncs = &mut HashMap::new();
+    let builtins = &mut HashMap::new();
 
     macro_rules! annotate {
         ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
@@ -74,6 +81,22 @@ pub fn init() -> Annotations {
         }
     }
 
+    macro_rules! annotate_builtin {
+        ($name:literal, $return_type:expr) => {
+            annotate_builtin!($name, $return_type, no_gc, leaf, elidable)
+        };
+        ($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;)+
+            builtins.insert($name, props);
+        }
+    }
+
     annotate!(rb_mKernel, "itself", types::BasicObject, no_gc, leaf, elidable);
     annotate!(rb_cString, "bytesize", types::Fixnum, no_gc, leaf);
     annotate!(rb_cModule, "name", types::StringExact.union(types::NilClass), no_gc, leaf, elidable);
@@ -83,7 +106,12 @@ 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);
+
     Annotations {
-        cfuncs: std::mem::take(cfuncs)
+        cfuncs: std::mem::take(cfuncs),
+        builtins: std::mem::take(builtins)
     }
 }

zjit/src/hir.rs

@@ -523,6 +523,13 @@ pub enum Insn {
 
     // Invoke a builtin function
     InvokeBuiltin { bf: rb_builtin_function, args: Vec<InsnId>, state: InsnId },
+    /// Call a builtin function with known return type
+    BuiltinCall {
+        bf: rb_builtin_function,
+        args: Vec<InsnId>,
+        state: InsnId,
+        return_type: Type,
+    },
 
     /// Control flow instructions
     Return { val: InsnId },
@@ -705,6 +712,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::BuiltinCall { bf, args, .. } => {
+                write!(f, "BuiltinCall {}", unsafe { CStr::from_ptr(bf.name) }.to_str().unwrap())?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            }
             Insn::Return { val } => { write!(f, "Return {val}") }
             Insn::FixnumAdd  { left, right, .. } => { write!(f, "FixnumAdd {left}, {right}") },
             Insn::FixnumSub  { left, right, .. } => { write!(f, "FixnumSub {left}, {right}") },
@@ -1160,6 +1174,7 @@ impl Function {
                 state,
             },
             &InvokeBuiltin { bf, ref args, state } => InvokeBuiltin { bf: bf, args: find_vec!(args), state },
+            &BuiltinCall { bf, ref args, state, return_type } => BuiltinCall { bf: bf, args: find_vec!(args), state, return_type },
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
             &CCall { cfun, ref args, name, return_type, elidable } => CCall { cfun, args: find_vec!(args), name, return_type, elidable },
@@ -1257,6 +1272,7 @@ impl Function {
             Insn::SendWithoutBlockDirect { .. } => types::BasicObject,
             Insn::Send { .. } => types::BasicObject,
             Insn::InvokeBuiltin { .. } => types::BasicObject,
+            Insn::BuiltinCall { return_type, .. } => *return_type,
             Insn::Defined { .. } => types::BasicObject,
             Insn::DefinedIvar { .. } => types::BasicObject,
             Insn::GetConstantPath { .. } => types::BasicObject,
@@ -1915,6 +1931,10 @@ impl Function {
                 worklist.extend(args);
                 worklist.push_back(state)
             }
+            &Insn::BuiltinCall { ref args, state, .. } => {
+                worklist.extend(args);
+                worklist.push_back(state)
+            }
             &Insn::CCall { ref args, .. } => worklist.extend(args),
             &Insn::GetIvar { self_val, state, .. } | &Insn::DefinedIvar { self_val, state, .. } => {
                 worklist.push_back(self_val);
@@ -3115,7 +3135,21 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     args.reverse();
 
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
-                    let insn_id = fun.push_insn(block, Insn::InvokeBuiltin { bf, args, state: exit_id });
+                    
+                    // Check if this builtin is annotated
+                    let builtin_name = unsafe { CStr::from_ptr(bf.name).to_str().unwrap() };
+                    let insn_id = if let Some(props) = ZJITState::get_method_annotations().get_builtin_properties(builtin_name) {
+                        // Use BuiltinCall with known return type
+                        fun.push_insn(block, Insn::BuiltinCall {
+                            bf,
+                            args,
+                            state: exit_id,
+                            return_type: props.return_type,
+                        })
+                    } else {
+                        // Fall back to InvokeBuiltin for unannotated builtins
+                        fun.push_insn(block, Insn::InvokeBuiltin { bf, args, state: exit_id })
+                    };
                     state.stack_push(insn_id);
                 }
                 YARVINSN_opt_invokebuiltin_delegate |
@@ -3130,7 +3164,21 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     }
 
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
-                    let insn_id = fun.push_insn(block, Insn::InvokeBuiltin { bf, args, state: exit_id });
+                    
+                    // Check if this builtin is annotated
+                    let builtin_name = unsafe { CStr::from_ptr(bf.name).to_str().unwrap() };
+                    let insn_id = if let Some(props) = ZJITState::get_method_annotations().get_builtin_properties(builtin_name) {
+                        // Use BuiltinCall with known return type
+                        fun.push_insn(block, Insn::BuiltinCall {
+                            bf,
+                            args,
+                            state: exit_id,
+                            return_type: props.return_type,
+                        })
+                    } else {
+                        // Fall back to InvokeBuiltin for unannotated builtins
+                        fun.push_insn(block, Insn::InvokeBuiltin { bf, args, state: exit_id })
+                    };
                     state.stack_push(insn_id);
                 }
                 YARVINSN_objtostring => {
@@ -4970,12 +5018,28 @@ mod tests {
 
     #[test]
     fn test_invokebuiltin_delegate_with_args() {
+        // Using an unannotated builtin to test InvokeBuiltin generation
+        let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("GC", "start"));
+        assert!(iseq_contains_opcode(iseq, YARVINSN_invokebuiltin), "iseq GC.start does not contain invokebuiltin");
+        let function = iseq_to_hir(iseq).unwrap();
+        assert_function_hir(function, expect![[r#"
+            fn start@<internal:gc>:36:
+            bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject, v3:BasicObject, v4:BasicObject):
+              v6:FalseClass = Const Value(false)
+              v8:BasicObject = InvokeBuiltin gc_start_internal, v0, v1, v2, v3, v6
+              Return v8
+        "#]]);
+    }
+
+    #[test]
+    fn test_invokebuiltin_delegate_annotated() {
+        // Test that Float now generates BuiltinCall with correct type
         assert_method_hir_with_opcode("Float", YARVINSN_opt_invokebuiltin_delegate_leave, expect![[r#"
             fn Float@<internal:kernel>:197:
             bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject, v3:BasicObject):
-              v6:BasicObject = InvokeBuiltin rb_f_float, v0, v1, v2
+              v6:Flonum = BuiltinCall rb_f_float, v0, v1, v2
               Jump bb1(v0, v1, v2, v3, v6)
-            bb1(v8:BasicObject, v9:BasicObject, v10:BasicObject, v11:BasicObject, v12:BasicObject):
+            bb1(v8:BasicObject, v9:BasicObject, v10:BasicObject, v11:BasicObject, v12:Flonum):
               Return v12
         "#]]);
     }