ZJIT: Inline struct aref

5 files changed, 167 insertions, 2 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 6e8038335f..5faa1bbcc3 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -448,6 +448,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::GuardShape { val, shape, state } => gen_guard_shape(jit, asm, opnd!(val), shape, &function.frame_state(state)),
         Insn::LoadPC => gen_load_pc(asm),
         Insn::LoadSelf => gen_load_self(),
+        &Insn::LoadField { recv, id, offset, return_type: _ } => gen_load_field(asm, opnd!(recv), id, offset),
         &Insn::LoadIvarEmbedded { self_val, id, index } => gen_load_ivar_embedded(asm, opnd!(self_val), id, index),
         &Insn::LoadIvarExtended { self_val, id, index } => gen_load_ivar_extended(asm, opnd!(self_val), id, index),
         &Insn::IsBlockGiven => gen_is_block_given(jit, asm),
@@ -981,6 +982,12 @@ fn gen_load_self() -> Opnd {
     Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF)
 }
 
+fn gen_load_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32) -> Opnd {
+    asm_comment!(asm, "Load field id={} offset={}", id.contents_lossy(), offset);
+    let recv = asm.load(recv);
+    asm.load(Opnd::mem(64, recv, offset))
+}
+
 fn gen_load_ivar_embedded(asm: &mut Assembler, self_val: Opnd, id: ID, index: u16) -> Opnd {
     // See ROBJECT_FIELDS() from include/ruby/internal/core/robject.h
 
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index d4e4079b5c..89488fd255 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -584,6 +584,13 @@ impl VALUE {
         }
     }
 
+    pub fn struct_embedded_p(self) -> bool {
+        unsafe {
+            RB_TYPE_P(self, RUBY_T_STRUCT) &&
+            FL_TEST_RAW(self, VALUE(RSTRUCT_EMBED_LEN_MASK)) != VALUE(0)
+        }
+    }
+
     pub fn as_fixnum(self) -> i64 {
         assert!(self.fixnum_p());
         (self.0 as i64) >> 1
@@ -1369,6 +1376,7 @@ pub(crate) mod ids {
         name: freeze
         name: minusat            content: b"-@"
         name: aref               content: b"[]"
+        name: _as_heap
     }
 
     /// Get an CRuby `ID` to an interned string, e.g. a particular method name.
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 574f5e40ea..269336fb16 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -444,6 +444,10 @@ impl PtrPrintMap {
         self.map_ptr(id as *const c_void)
     }
 
+    fn map_offset(&self, id: i32) -> *const c_void {
+        self.map_ptr(id as *const c_void)
+    }
+
     /// Map shape ID into a pointer for printing
     fn map_shape(&self, id: ShapeId) -> *const c_void {
         self.map_ptr(id.0 as *const c_void)
@@ -671,6 +675,7 @@ pub enum Insn {
     LoadPC,
     /// Load cfp->self
     LoadSelf,
+    LoadField { recv: InsnId, id: ID, offset: i32, return_type: Type },
     /// Read an instance variable at the given index, embedded in the object
     LoadIvarEmbedded { self_val: InsnId, id: ID, index: u16 },
     /// Read an instance variable at the given index, from the extended table
@@ -909,6 +914,7 @@ impl Insn {
             Insn::IsNil      { .. } => false,
             Insn::LoadPC => false,
             Insn::LoadSelf => false,
+            Insn::LoadField { .. } => false,
             Insn::LoadIvarEmbedded { .. } => false,
             Insn::LoadIvarExtended { .. } => false,
             Insn::CCall { elidable, .. } => !elidable,
@@ -1170,6 +1176,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::GetIvar { self_val, id, .. } => write!(f, "GetIvar {self_val}, :{}", id.contents_lossy()),
             Insn::LoadPC => write!(f, "LoadPC"),
             Insn::LoadSelf => write!(f, "LoadSelf"),
+            &Insn::LoadField { recv, id, offset, return_type: _ } => write!(f, "LoadField {recv}, :{}@{:p}", id.contents_lossy(), self.ptr_map.map_offset(offset)),
             &Insn::LoadIvarEmbedded { self_val, id, index } => write!(f, "LoadIvarEmbedded {self_val}, :{}@{:p}", id.contents_lossy(), self.ptr_map.map_index(index as u64)),
             &Insn::LoadIvarExtended { self_val, id, index } => write!(f, "LoadIvarExtended {self_val}, :{}@{:p}", id.contents_lossy(), self.ptr_map.map_index(index as u64)),
             Insn::SetIvar { self_val, id, val, .. } => write!(f, "SetIvar {self_val}, :{}, {val}", id.contents_lossy()),
@@ -1792,6 +1799,7 @@ impl Function {
             &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
             &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
+            &LoadField { recv, id, offset, return_type } => LoadField { recv: find!(recv), id, offset, return_type },
             &LoadIvarEmbedded { self_val, id, index } => LoadIvarEmbedded { self_val: find!(self_val), id, index },
             &LoadIvarExtended { self_val, id, index } => LoadIvarExtended { self_val: find!(self_val), id, index },
             &SetIvar { self_val, id, val, state } => SetIvar { self_val: find!(self_val), id, val: find!(val), state },
@@ -1929,6 +1937,7 @@ impl Function {
             Insn::GetIvar { .. } => types::BasicObject,
             Insn::LoadPC => types::CPtr,
             Insn::LoadSelf => types::BasicObject,
+            &Insn::LoadField { return_type, .. } => return_type,
             Insn::LoadIvarEmbedded { .. } => types::BasicObject,
             Insn::LoadIvarExtended { .. } => types::BasicObject,
             Insn::GetSpecialSymbol { .. } => types::BasicObject,
@@ -2390,8 +2399,52 @@ impl Function {
                             self.make_equal_to(insn_id, val);
                         } else if def_type == VM_METHOD_TYPE_OPTIMIZED {
                             let opt_type = unsafe { get_cme_def_body_optimized_type(cme) };
-                            self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedOptimizedMethodType(OptimizedMethodType::from(opt_type)));
-                            self.push_insn_id(block, insn_id); continue;
+                            if opt_type == OPTIMIZED_METHOD_TYPE_STRUCT_AREF {
+                                if unsafe { vm_ci_argc(ci) } != 0 {
+                                    self.push_insn_id(block, insn_id); continue;
+                                }
+                                let index: i32 = unsafe { get_cme_def_body_optimized_index(cme) }
+                                                .try_into()
+                                                .unwrap();
+                                // We are going to use an encoding that takes a 4-byte immediate which
+                                // limits the offset to INT32_MAX.
+                                {
+                                    let native_index = (index as i64) * (SIZEOF_VALUE as i64);
+                                    if native_index > (i32::MAX as i64) {
+                                        self.push_insn_id(block, insn_id); continue;
+                                    }
+                                }
+                                // Get the profiled type to check if the fields is embedded or heap allocated.
+                                let Some(is_embedded) = self.profiled_type_of_at(recv, frame_state.insn_idx).map(|t| t.flags().is_struct_embedded()) else {
+                                    // No (monomorphic) profile info
+                                    self.push_insn_id(block, insn_id); continue;
+                                };
+                                self.push_insn(block, Insn::PatchPoint { invariant: Invariant::MethodRedefined { klass, method: mid, cme }, state });
+                                if klass.instance_can_have_singleton_class() {
+                                    self.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoSingletonClass { klass }, state });
+                                }
+                                if let Some(profiled_type) = profiled_type {
+                                    recv = self.push_insn(block, Insn::GuardType { val: recv, guard_type: Type::from_profiled_type(profiled_type), state });
+                                }
+                                // All structs from the same Struct class should have the same
+                                // length. So if our recv is embedded all runtime
+                                // structs of the same class should be as well, and the same is
+                                // true of the converse.
+                                //
+                                // No need for a GuardShape.
+                                let replacement = if is_embedded {
+                                    let offset = RUBY_OFFSET_RSTRUCT_AS_ARY + (SIZEOF_VALUE_I32 * index);
+                                    self.push_insn(block, Insn::LoadField { recv, id: mid, offset, return_type: types::BasicObject })
+                                } else {
+                                    let as_heap = self.push_insn(block, Insn::LoadField { recv, id: ID!(_as_heap), offset: RUBY_OFFSET_RSTRUCT_AS_HEAP_PTR, return_type: types::CPtr });
+                                    let offset = SIZEOF_VALUE_I32 * index;
+                                    self.push_insn(block, Insn::LoadField { recv: as_heap, id: mid, offset, return_type: types::BasicObject })
+                                };
+                                self.make_equal_to(insn_id, replacement);
+                            } else {
+                                self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedOptimizedMethodType(OptimizedMethodType::from(opt_type)));
+                                self.push_insn_id(block, insn_id); continue;
+                            }
                         } else {
                             self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedMethodType(MethodType::from(def_type)));
                             self.push_insn_id(block, insn_id); continue;
@@ -3320,6 +3373,9 @@ impl Function {
                 worklist.push_back(str);
                 worklist.push_back(state);
             }
+            &Insn::LoadField { recv, .. } => {
+                worklist.push_back(recv);
+            }
             &Insn::LoadIvarEmbedded { self_val, .. }
             | &Insn::LoadIvarExtended { self_val, .. } => {
                 worklist.push_back(self_val);
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index f29af66bde..fd93b4a444 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -4936,6 +4936,94 @@ mod hir_opt_tests {
     }
 
     #[test]
+    fn test_inline_struct_aref_embedded() {
+        eval(r#"
+            C = Struct.new(:foo)
+            def test(o) = o.foo
+            test C.new
+            test C.new
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          PatchPoint MethodRedefined(C@0x1000, foo@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v22:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
+          v23:BasicObject = LoadField v22, :foo@0x1038
+          CheckInterrupts
+          Return v23
+        ");
+    }
+
+    #[test]
+    fn test_inline_struct_aref_heap() {
+        eval(r#"
+            C = Struct.new(*(0..1000).map {|i| :"a#{i}"}, :foo)
+            def test(o) = o.foo
+            test C.new
+            test C.new
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          PatchPoint MethodRedefined(C@0x1000, foo@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v22:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
+          v23:CPtr = LoadField v22, :_as_heap@0x1038
+          v24:BasicObject = LoadField v23, :foo@0x1039
+          CheckInterrupts
+          Return v24
+        ");
+    }
+
+    #[test]
+    fn test_elide_struct_aref() {
+        eval(r#"
+            C = Struct.new(*(0..1000).map {|i| :"a#{i}"}, :foo)
+            def test(o)
+              o.foo
+              5
+            end
+            test C.new
+            test C.new
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:4:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          PatchPoint MethodRedefined(C@0x1000, foo@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v25:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
+          v17:Fixnum[5] = Const Value(5)
+          CheckInterrupts
+          Return v17
+        ");
+    }
+
+    #[test]
     fn test_array_reverse_returns_array() {
         eval(r#"
             def test = [].reverse
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
index a6c837df5a..08fdf3eb97 100644
--- a/zjit/src/profile.rs
+++ b/zjit/src/profile.rs
@@ -145,6 +145,8 @@ impl Flags {
     const IS_EMBEDDED: u32 = 1 << 1;
     /// Object is a T_OBJECT
     const IS_T_OBJECT: u32 = 1 << 2;
+    /// Object is a struct with embedded fields
+    const IS_STRUCT_EMBEDDED: u32 = 1 << 3;
 
     pub fn none() -> Self { Self(Self::NONE) }
 
@@ -152,6 +154,7 @@ impl Flags {
     pub fn is_immediate(self) -> bool { (self.0 & Self::IS_IMMEDIATE) != 0 }
     pub fn is_embedded(self) -> bool { (self.0 & Self::IS_EMBEDDED) != 0 }
     pub fn is_t_object(self) -> bool { (self.0 & Self::IS_T_OBJECT) != 0 }
+    pub fn is_struct_embedded(self) -> bool { (self.0 & Self::IS_STRUCT_EMBEDDED) != 0 }
 }
 
 /// opt_send_without_block/opt_plus/... should store:
@@ -214,6 +217,9 @@ impl ProfiledType {
         if obj.embedded_p() {
             flags.0 |= Flags::IS_EMBEDDED;
         }
+        if obj.struct_embedded_p() {
+            flags.0 |= Flags::IS_STRUCT_EMBEDDED;
+        }
         if unsafe { RB_TYPE_P(obj, RUBY_T_OBJECT) } {
             flags.0 |= Flags::IS_T_OBJECT;
         }