Add return_type to FnProperties and CCall

Author: tekknolagi

zjit/src/codegen.rs

@@ -203,7 +203,7 @@ fn gen_insn(jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_i
         Insn::GuardType { val, guard_type, state } => gen_guard_type(asm, opnd!(val), *guard_type, &function.frame_state(*state))?,
         Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(asm, opnd!(val), *expected, &function.frame_state(*state))?,
         Insn::PatchPoint(_) => return Some(()), // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
-        Insn::CCall { cfun, args, name: _ } => gen_ccall(jit, asm, *cfun, args)?,
+        Insn::CCall { cfun, args, name: _, return_type: _ } => gen_ccall(jit, asm, *cfun, args)?,
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
             return None;

zjit/src/cruby_methods.rs

@@ -11,18 +11,21 @@
 use crate::cruby::*;
 use std::collections::HashMap;
 use std::ffi::c_void;
+use crate::hir_type::{types, Type};
 
 pub struct Annotations {
     cfuncs: HashMap<*mut c_void, FnProperties>,
 }
 
 /// Runtime behaviors of C functions that implement a Ruby method
-#[derive(Clone, Copy, Default)]
+#[derive(Clone, Copy)]
 pub struct FnProperties {
     /// Whether it's possible for the function to yield to the GC
     pub no_gc: bool,
     /// Whether it's possible for the function to make a ruby call
     pub leaf: bool,
+    /// What Type the C function returns
+    pub return_type: Type,
 }
 
 impl Annotations {
@@ -60,19 +63,19 @@ pub fn init() -> Annotations {
     let cfuncs = &mut HashMap::new();
 
     macro_rules! annotate {
-        ($module:ident, $method_name:literal, $($properties:ident),+) => {
-            let mut props = FnProperties::default();
+        ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
+            let mut props = FnProperties { no_gc: false, leaf: false, return_type: $return_type };
             $(
                 props.$properties = true;
             )+
             annotate_c_method(cfuncs, unsafe { $module }, $method_name, props);
         }
     }
 
-    annotate!(rb_mKernel, "itself", no_gc, leaf);
-    annotate!(rb_cString, "bytesize", no_gc, leaf);
-    annotate!(rb_cModule, "name", no_gc, leaf);
-    annotate!(rb_cModule, "===", no_gc, leaf);
+    annotate!(rb_mKernel, "itself", types::BasicObject, no_gc, leaf);
+    annotate!(rb_cString, "bytesize", types::Fixnum, no_gc, leaf);
+    annotate!(rb_cModule, "name", types::StringExact.union(types::NilClassExact), no_gc, leaf);
+    annotate!(rb_cModule, "===", types::BoolExact, no_gc, leaf);
 
     Annotations {
         cfuncs: std::mem::take(cfuncs)

zjit/src/hir.rs

@@ -326,7 +326,7 @@ pub enum Insn {
 
     // Call a C function
     // `name` is for printing purposes only
-    CCall { cfun: *const u8, args: Vec<InsnId>, name: ID },
+    CCall { cfun: *const u8, args: Vec<InsnId>, name: ID, return_type: Type },
 
     // Send without block with dynamic dispatch
     // Ignoring keyword arguments etc for now
@@ -475,11 +475,11 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::FixnumLe   { left, right, .. } => { write!(f, "FixnumLe {left}, {right}") },
             Insn::FixnumGt   { left, right, .. } => { write!(f, "FixnumGt {left}, {right}") },
             Insn::FixnumGe   { left, right, .. } => { write!(f, "FixnumGe {left}, {right}") },
-            Insn::GuardType { val, guard_type, .. } => { write!(f, "GuardType {val}, {guard_type}") },
+            Insn::GuardType { val, guard_type, .. } => { write!(f, "GuardType {val}, {}", guard_type.print(self.ptr_map)) },
             Insn::GuardBitEquals { val, expected, .. } => { write!(f, "GuardBitEquals {val}, {}", expected.print(self.ptr_map)) },
             Insn::PatchPoint(invariant) => { write!(f, "PatchPoint {}", invariant.print(self.ptr_map)) },
             Insn::GetConstantPath { ic } => { write!(f, "GetConstantPath {:p}", self.ptr_map.map_ptr(ic)) },
-            Insn::CCall { cfun, args, name } => {
+            Insn::CCall { cfun, args, name, return_type: _ } => {
                 write!(f, "CCall {}@{:p}", name.contents_lossy(), self.ptr_map.map_ptr(cfun))?;
                 for arg in args {
                     write!(f, ", {arg}")?;
@@ -790,7 +790,7 @@ impl Function {
             },
             ArraySet { array, idx, val } => ArraySet { array: find!(*array), idx: *idx, val: find!(*val) },
             ArrayDup { val } => ArrayDup { val: find!(*val) },
-            CCall { cfun, args, name } => CCall { cfun: *cfun, args: args.iter().map(|arg| find!(*arg)).collect(), name: *name },
+            CCall { cfun, args, name, return_type } => CCall { cfun: *cfun, args: args.iter().map(|arg| find!(*arg)).collect(), name: *name, return_type: *return_type },
             Defined { .. } => todo!("find(Defined)"),
         }
     }
@@ -838,7 +838,7 @@ impl Function {
             Insn::StringIntern { .. } => types::StringExact,
             Insn::NewArray { .. } => types::ArrayExact,
             Insn::ArrayDup { .. } => types::ArrayExact,
-            Insn::CCall { .. } => types::Any,
+            Insn::CCall { return_type, .. } => *return_type,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
             Insn::GuardBitEquals { val, expected, .. } => self.type_of(*val).intersection(Type::from_value(*expected)),
             Insn::FixnumAdd  { .. } => types::Fixnum,
@@ -1024,7 +1024,7 @@ impl Function {
 
                     // Filter for a leaf and GC free function
                     use crate::cruby_methods::FnProperties;
-                    let Some(FnProperties { leaf: true, no_gc: true }) =
+                    let Some(FnProperties { leaf: true, no_gc: true, return_type }) =
                         ZJITState::get_method_annotations().get_cfunc_properties(method)
                     else {
                         return Err(());
@@ -1040,7 +1040,7 @@ impl Function {
                         let cfun = unsafe { get_mct_func(cfunc) }.cast();
                         let mut cfunc_args = vec![self_val];
                         cfunc_args.append(&mut args);
-                        let ccall = fun.push_insn(block, Insn::CCall { cfun, args: cfunc_args, name: method_id });
+                        let ccall = fun.push_insn(block, Insn::CCall { cfun, args: cfunc_args, name: method_id, return_type });
                         fun.make_equal_to(send_insn_id, ccall);
                         return Ok(());
                     }
@@ -3306,7 +3306,7 @@ mod opt_tests {
             bb0():
               v1:Fixnum[1] = Const Value(1)
               PatchPoint MethodRedefined(Integer@0x1000, itself@0x1008)
-              v8:Any = CCall itself@0x1010, v1
+              v8:BasicObject = CCall itself@0x1010, v1
               PatchPoint CalleeModifiedLocals(v8)
               Return v8
         "#]]);
@@ -3330,4 +3330,24 @@ mod opt_tests {
               Return v4
         "#]]);
     }
+
+    #[test]
+    fn string_bytesize_simple() {
+        eval("
+            def test = 'abc'.bytesize
+            test
+            test
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0():
+              v1:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v2:StringExact = StringCopy v1
+              PatchPoint MethodRedefined(String@0x1008, bytesize@0x1010)
+              v8:StringExact[VALUE(0x1018)] = GuardType v2, StringExact[VALUE(0x1018)]
+              v9:Fixnum = CCall bytesize@0x1020, v2
+              PatchPoint CalleeModifiedLocals(v9)
+              Return v9
+        "#]]);
+    }
 }