ZJIT: Specialize ObjectAlloc with known class pointer

This has fewer effects (can be elided!) and will eventually get better
codegen, too.

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

3 files changed, 74 insertions, 13 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 2830b623b8..0f167ceec3 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -351,6 +351,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::NewRangeFixnum { low, high, flag, state } => gen_new_range_fixnum(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)),
         Insn::ObjectAlloc { val, state } => gen_object_alloc(jit, asm, opnd!(val), &function.frame_state(*state)),
+        &Insn::ObjectAllocClass { class, state } => gen_object_alloc_class(asm, class, &function.frame_state(state)),
         Insn::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)),
         // concatstrings shouldn't have 0 strings
         // If it happens we abort the compilation for now
@@ -1234,12 +1235,19 @@ fn gen_new_range_fixnum(
 }
 
 fn gen_object_alloc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, state: &FrameState) -> lir::Opnd {
-    // TODO: this is leaf in the vast majority of cases,
-    // Should specialize to avoid `gen_prepare_non_leaf_call` (Shopify#747)
+    // Allocating an object from an unknown class is non-leaf; see doc for `ObjectAlloc`.
     gen_prepare_non_leaf_call(jit, asm, state);
     asm_ccall!(asm, rb_obj_alloc, val)
 }
 
+fn gen_object_alloc_class(asm: &mut Assembler, class: VALUE, state: &FrameState) -> lir::Opnd {
+    // Allocating an object for a known class with default allocator is leaf; see doc for
+    // `ObjectAllocClass`.
+    gen_prepare_leaf_call_with_gc(asm, state);
+    // TODO(max): directly call `class_call_alloc_func`
+    asm_ccall!(asm, rb_obj_alloc, class.into())
+}
+
 /// Compile code that exits from JIT code with a return value
 fn gen_return(asm: &mut Assembler, val: lir::Opnd) {
     // Pop the current frame (ec->cfp++)
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index af30f60fe4..114cfda549 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -560,8 +560,15 @@ pub enum Insn {
 
     HashDup { val: InsnId, state: InsnId },
 
-    /// Allocate an instance of the `val` class without calling `#initialize` on it.
+    /// Allocate an instance of the `val` object without calling `#initialize` on it.
+    /// This can:
+    /// * raise an exception if `val` is not a class
+    /// * run arbitrary code if `val` is a class with a custom allocator
     ObjectAlloc { val: InsnId, state: InsnId },
+    /// Allocate an instance of the `val` class without calling `#initialize` on it.
+    /// This requires that `class` has the default allocator (for example via `IsMethodCfunc`).
+    /// This won't raise or run arbitrary code because `class` has the default allocator.
+    ObjectAllocClass { class: VALUE, state: InsnId },
 
     /// Check if the value is truthy and "return" a C boolean. In reality, we will likely fuse this
     /// with IfTrue/IfFalse in the backend to generate jcc.
@@ -755,6 +762,7 @@ impl Insn {
             Insn::LoadIvarEmbedded { .. } => false,
             Insn::LoadIvarExtended { .. } => false,
             Insn::CCall { elidable, .. } => !elidable,
+            Insn::ObjectAllocClass { .. } => false,
             // TODO: NewRange is effects free if we can prove the two ends to be Fixnum,
             // but we don't have type information here in `impl Insn`. See rb_range_new().
             Insn::NewRange { .. } => true,
@@ -819,6 +827,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::ArrayDup { val, .. } => { write!(f, "ArrayDup {val}") }
             Insn::HashDup { val, .. } => { write!(f, "HashDup {val}") }
             Insn::ObjectAlloc { val, .. } => { write!(f, "ObjectAlloc {val}") }
+            Insn::ObjectAllocClass { class, .. } => { write!(f, "ObjectAllocClass {}", class.print(self.ptr_map)) }
             Insn::StringCopy { val, .. } => { write!(f, "StringCopy {val}") }
             Insn::StringConcat { strings, .. } => {
                 write!(f, "StringConcat")?;
@@ -1411,6 +1420,7 @@ impl Function {
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
             &ObjectAlloc { val, state } => ObjectAlloc { val: find!(val), state },
+            &ObjectAllocClass { class, state } => ObjectAllocClass { class, state: find!(state) },
             &CCall { cfun, ref args, name, return_type, elidable } => CCall { cfun, args: find_vec!(args), name, return_type, elidable },
             &Defined { op_type, obj, pushval, v, state } => Defined { op_type, obj, pushval, v: find!(v), state: find!(state) },
             &DefinedIvar { self_val, pushval, id, state } => DefinedIvar { self_val: find!(self_val), pushval, id, state },
@@ -1497,6 +1507,7 @@ impl Function {
             Insn::NewRange { .. } => types::RangeExact,
             Insn::NewRangeFixnum { .. } => types::RangeExact,
             Insn::ObjectAlloc { .. } => types::HeapObject,
+            Insn::ObjectAllocClass { class, .. } => Type::from_class(*class),
             Insn::CCall { return_type, .. } => *return_type,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
             Insn::GuardTypeNot { .. } => types::BasicObject,
@@ -1895,6 +1906,17 @@ impl Function {
                             self.push_insn_id(block, insn_id);
                         }
                     }
+                    Insn::ObjectAlloc { val, state } => {
+                        let val_type = self.type_of(val);
+                        if val_type.is_subtype(types::Class) && val_type.ruby_object_known() {
+                            let class = val_type.ruby_object().unwrap();
+                            let replacement = self.push_insn(block, Insn::ObjectAllocClass { class, state });
+                            self.insn_types[replacement.0] = self.infer_type(replacement);
+                            self.make_equal_to(insn_id, replacement);
+                        } else {
+                            self.push_insn_id(block, insn_id);
+                        }
+                    }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
             }
@@ -2373,6 +2395,7 @@ impl Function {
             &Insn::GetGlobal { state, .. } |
             &Insn::GetSpecialSymbol { state, .. } |
             &Insn::GetSpecialNumber { state, .. } |
+            &Insn::ObjectAllocClass { state, .. } |
             &Insn::SideExit { state, .. } => worklist.push_back(state),
         }
     }
@@ -8071,10 +8094,10 @@ mod opt_tests {
           v6:NilClass = Const Value(nil)
           v7:CBool = IsMethodCFunc v34, :new
           IfFalse v7, bb1(v0, v6, v34)
-          v10:HeapObject = ObjectAlloc v34
-          v12:BasicObject = SendWithoutBlock v10, :initialize
+          v35:HeapObject[class_exact:C] = ObjectAllocClass VALUE(0x1008)
+          v12:BasicObject = SendWithoutBlock v35, :initialize
           CheckInterrupts
-          Jump bb2(v0, v10, v12)
+          Jump bb2(v0, v35, v12)
         bb1(v16:BasicObject, v17:NilClass, v18:Class[VALUE(0x1008)]):
           v21:BasicObject = SendWithoutBlock v18, :new
           Jump bb2(v16, v21, v17)
@@ -8105,12 +8128,11 @@ mod opt_tests {
           v7:Fixnum[1] = Const Value(1)
           v8:CBool = IsMethodCFunc v36, :new
           IfFalse v8, bb1(v0, v6, v36, v7)
-          v11:HeapObject = ObjectAlloc v36
+          v37:HeapObject[class_exact:C] = ObjectAllocClass VALUE(0x1008)
           PatchPoint MethodRedefined(C@0x1008, initialize@0x1010, cme:0x1018)
-          v38:HeapObject[class_exact:C] = GuardType v11, HeapObject[class_exact:C]
-          v39:BasicObject = SendWithoutBlockDirect v38, :initialize (0x1040), v7
+          v39:BasicObject = SendWithoutBlockDirect v37, :initialize (0x1040), v7
           CheckInterrupts
-          Jump bb2(v0, v11, v39)
+          Jump bb2(v0, v37, v39)
         bb1(v17:BasicObject, v18:NilClass, v19:Class[VALUE(0x1008)], v20:Fixnum[1]):
           v23:BasicObject = SendWithoutBlock v19, :new, v20
           Jump bb2(v17, v23, v18)
diff --git a/zjit/src/hir_type/mod.rs b/zjit/src/hir_type/mod.rs
index 33f6ab223b..f2fb870257 100644
--- a/zjit/src/hir_type/mod.rs
+++ b/zjit/src/hir_type/mod.rs
@@ -257,6 +257,20 @@ impl Type {
         }
     }
 
+    pub fn from_class(class: VALUE) -> Type {
+        if class == unsafe { rb_cArray } { types::ArrayExact }
+        else if class == unsafe { rb_cFalseClass } { types::FalseClass }
+        else if class == unsafe { rb_cHash } { types::HashExact }
+        else if class == unsafe { rb_cInteger } { types::Integer}
+        else if class == unsafe { rb_cNilClass } { types::NilClass }
+        else if class == unsafe { rb_cString } { types::StringExact }
+        else if class == unsafe { rb_cTrueClass } { types::TrueClass }
+        else {
+            // TODO(max): Add more cases for inferring type bits from built-in types
+            Type { bits: bits::HeapObject, spec: Specialization::TypeExact(class) }
+        }
+    }
+
     /// Private. Only for creating type globals.
     const fn from_bits(bits: u64) -> Type {
         Type {
@@ -669,10 +683,27 @@ mod tests {
     }
 
     #[test]
+    fn from_class() {
+        crate::cruby::with_rubyvm(|| {
+            assert_bit_equal(Type::from_class(unsafe { rb_cInteger }), types::Integer);
+            assert_bit_equal(Type::from_class(unsafe { rb_cString }), types::StringExact);
+            assert_bit_equal(Type::from_class(unsafe { rb_cArray }), types::ArrayExact);
+            assert_bit_equal(Type::from_class(unsafe { rb_cHash }), types::HashExact);
+            assert_bit_equal(Type::from_class(unsafe { rb_cNilClass }), types::NilClass);
+            assert_bit_equal(Type::from_class(unsafe { rb_cTrueClass }), types::TrueClass);
+            assert_bit_equal(Type::from_class(unsafe { rb_cFalseClass }), types::FalseClass);
+            let c_class = define_class("C", unsafe { rb_cObject });
+            assert_bit_equal(Type::from_class(c_class), Type { bits: bits::HeapObject, spec: Specialization::TypeExact(c_class) });
+        });
+    }
+
+    #[test]
     fn integer_has_ruby_class() {
-        assert_eq!(Type::fixnum(3).inexact_ruby_class(), Some(unsafe { rb_cInteger }));
-        assert_eq!(types::Fixnum.inexact_ruby_class(), None);
-        assert_eq!(types::Integer.inexact_ruby_class(), None);
+        crate::cruby::with_rubyvm(|| {
+            assert_eq!(Type::fixnum(3).inexact_ruby_class(), Some(unsafe { rb_cInteger }));
+            assert_eq!(types::Fixnum.inexact_ruby_class(), None);
+            assert_eq!(types::Integer.inexact_ruby_class(), None);
+        });
     }
 
     #[test]