This commit implements the Object Shapes technique in CRuby.

Object Shapes is used for accessing instance variables and representing the
"frozenness" of objects.  Object instances have a "shape" and the shape
represents some attributes of the object (currently which instance variables are
set and the "frozenness").  Shapes form a tree data structure, and when a new
instance variable is set on an object, that object "transitions" to a new shape
in the shape tree.  Each shape has an ID that is used for caching. The shape
structure is independent of class, so objects of different types can have the
same shape.

For example:

```ruby
class Foo
  def initialize
    # Starts with shape id 0
    @a = 1 # transitions to shape id 1
    @b = 1 # transitions to shape id 2
  end
end

class Bar
  def initialize
    # Starts with shape id 0
    @a = 1 # transitions to shape id 1
    @b = 1 # transitions to shape id 2
  end
end

foo = Foo.new # `foo` has shape id 2
bar = Bar.new # `bar` has shape id 2
```

Both `foo` and `bar` instances have the same shape because they both set
instance variables of the same name in the same order.

This technique can help to improve inline cache hits as well as generate more
efficient machine code in JIT compilers.

This commit also adds some methods for debugging shapes on objects.  See
`RubyVM::Shape` for more details.

For more context on Object Shapes, see [Feature: #18776]

Co-Authored-By: Aaron Patterson <tenderlove@ruby-lang.org>
Co-Authored-By: Eileen M. Uchitelle <eileencodes@gmail.com>
Co-Authored-By: John Hawthorn <john@hawthorn.email>

5 files changed, 107 insertions, 88 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index c3d4a39a2b..4b50d888de 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -40,6 +40,7 @@ fn main() {
         .header("internal.h")
         .header("internal/re.h")
         .header("include/ruby/ruby.h")
+        .header("shape.h")
         .header("vm_core.h")
         .header("vm_callinfo.h")
 
@@ -81,6 +82,12 @@ fn main() {
         // This function prints info about a value and is useful for debugging
         .allowlist_function("rb_obj_info_dump")
 
+        // From shape.h
+        .allowlist_function("rb_shape_get_shape_id")
+        .allowlist_function("rb_shape_get_shape_by_id")
+        .allowlist_function("rb_shape_flags_mask")
+        .allowlist_function("rb_shape_get_iv_index")
+
         // From ruby/internal/intern/object.h
         .allowlist_function("rb_obj_is_kind_of")
 
diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index d310e3bf12..42d97b7e80 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/yjit/src/asm/x86_64/mod.rs
@@ -617,7 +617,7 @@ fn write_rm_multi(cb: &mut CodeBlock, op_mem_reg8: u8, op_mem_reg_pref: u8, op_r
                 write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, op_ext_imm, &[op_mem_imm_lrg]);
                 cb.write_int(uimm.value, if opnd_size > 32 { 32 } else { opnd_size.into() });
             } else {
-                panic!("immediate value too large");
+                panic!("immediate value too large (num_bits={})", num_bits);
             }
         },
         _ => unreachable!()
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index c246c7b48f..4018a314fc 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1938,14 +1938,12 @@ fn gen_set_ivar(
     let val_opnd = ctx.stack_pop(1);
     let recv_opnd = ctx.stack_pop(1);
 
-    let ivar_index: u32 = unsafe { rb_obj_ensure_iv_index_mapping(recv, ivar_name) };
-
-    // Call rb_vm_set_ivar_idx with the receiver, the index of the ivar, and the value
+    // Call rb_vm_set_ivar_id with the receiver, the ivar name, and the value
     let val = asm.ccall(
-        rb_vm_set_ivar_idx as *const u8,
+        rb_vm_set_ivar_id as *const u8,
         vec![
             recv_opnd,
-            Opnd::Imm(ivar_index.into()),
+            Opnd::UImm(ivar_name.into()),
             val_opnd,
         ],
     );
@@ -2023,81 +2021,82 @@ fn gen_get_ivar(
         return EndBlock;
     }
 
-    // FIXME: Mapping the index could fail when there is too many ivar names. If we're
-    // compiling for a branch stub that can cause the exception to be thrown from the
-    // wrong PC.
-    let ivar_index =
-        unsafe { rb_obj_ensure_iv_index_mapping(comptime_receiver, ivar_name) }.as_usize();
+    let ivar_index = unsafe {
+        let shape_id = comptime_receiver.shape_of();
+        let shape = rb_shape_get_shape_by_id(shape_id);
+        let mut ivar_index: u32 = 0;
+        if rb_shape_get_iv_index(shape, ivar_name, &mut ivar_index) {
+            Some(ivar_index as usize)
+        } else {
+            None
+        }
+    };
+
+    // must be before stack_pop
+    let recv_type = ctx.get_opnd_type(recv_opnd);
+
+    // Upgrade type
+    if !recv_type.is_heap() {
+        ctx.upgrade_opnd_type(recv_opnd, Type::UnknownHeap);
+    }
 
     // Pop receiver if it's on the temp stack
     if recv_opnd != SelfOpnd {
         ctx.stack_pop(1);
     }
 
-    if USE_RVARGC != 0 {
-        // Check that the ivar table is big enough
-        // Check that the slot is inside the ivar table (num_slots > index)
-        let num_slots = Opnd::mem(32, recv, ROBJECT_OFFSET_NUMIV);
-        asm.cmp(num_slots, Opnd::UImm(ivar_index as u64));
-        asm.jbe(counted_exit!(ocb, side_exit, getivar_idx_out_of_range).into());
+    // Guard heap object
+    if !recv_type.is_heap() {
+        guard_object_is_heap(asm, recv, side_exit);
     }
 
     // Compile time self is embedded and the ivar index lands within the object
-    let test_result = unsafe { FL_TEST_RAW(comptime_receiver, VALUE(ROBJECT_EMBED.as_usize())) != VALUE(0) };
-    if test_result {
-        // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
+    let embed_test_result = unsafe { FL_TEST_RAW(comptime_receiver, VALUE(ROBJECT_EMBED.as_usize())) != VALUE(0) };
 
-        // Guard that self is embedded
-        // TODO: BT and JC is shorter
-        asm.comment("guard embedded getivar");
-        let flags_opnd = Opnd::mem(64, recv, RUBY_OFFSET_RBASIC_FLAGS);
-        asm.test(flags_opnd, Opnd::UImm(ROBJECT_EMBED as u64));
-        let side_exit = counted_exit!(ocb, side_exit, getivar_megamorphic);
-        jit_chain_guard(
-            JCC_JZ,
-            jit,
-            &starting_context,
-            asm,
-            ocb,
-            max_chain_depth,
-            side_exit,
-        );
+    let flags_mask: usize = unsafe { rb_shape_flags_mask() }.as_usize();
+    let expected_flags_mask: usize = (RUBY_T_MASK as usize) | !flags_mask | (ROBJECT_EMBED as usize);
+    let expected_flags = comptime_receiver.builtin_flags() & expected_flags_mask;
+
+    // Combined guard for all flags: shape, embeddedness, and T_OBJECT
+    let flags_opnd = Opnd::mem(64, recv, RUBY_OFFSET_RBASIC_FLAGS);
+
+    asm.comment("guard shape, embedded, and T_OBJECT");
+    let flags_opnd = asm.and(flags_opnd, Opnd::UImm(expected_flags_mask as u64));
+    asm.cmp(flags_opnd, Opnd::UImm(expected_flags as u64));
+    jit_chain_guard(
+        JCC_JNE,
+        jit,
+        &starting_context,
+        asm,
+        ocb,
+        max_chain_depth,
+        side_exit,
+    );
+
+    // If there is no IVAR index, then the ivar was undefined
+    // when we entered the compiler.  That means we can just return
+    // nil for this shape + iv name
+    if ivar_index.is_none() {
+        let out_opnd = ctx.stack_push(Type::Nil);
+        asm.mov(out_opnd, Qnil.into());
+    } else if embed_test_result {
+        // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
 
         // Load the variable
-        let offs = ROBJECT_OFFSET_AS_ARY + (ivar_index * SIZEOF_VALUE) as i32;
+        let offs = ROBJECT_OFFSET_AS_ARY + (ivar_index.unwrap() * SIZEOF_VALUE) as i32;
         let ivar_opnd = Opnd::mem(64, recv, offs);
 
-        // Guard that the variable is not Qundef
-        asm.cmp(ivar_opnd, Qundef.into());
-        let out_val = asm.csel_e(Qnil.into(), ivar_opnd);
-
         // Push the ivar on the stack
         let out_opnd = ctx.stack_push(Type::Unknown);
-        asm.mov(out_opnd, out_val);
+        asm.mov(out_opnd, ivar_opnd);
     } else {
         // Compile time value is *not* embedded.
 
-        // Guard that value is *not* embedded
-        // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
-        asm.comment("guard extended getivar");
-        let flags_opnd = Opnd::mem(64, recv, RUBY_OFFSET_RBASIC_FLAGS);
-        asm.test(flags_opnd, Opnd::UImm(ROBJECT_EMBED as u64));
-        let megamorphic_side_exit = counted_exit!(ocb, side_exit, getivar_megamorphic);
-        jit_chain_guard(
-            JCC_JNZ,
-            jit,
-            &starting_context,
-            asm,
-            ocb,
-            max_chain_depth,
-            megamorphic_side_exit,
-        );
-
         if USE_RVARGC == 0 {
             // Check that the extended table is big enough
             // Check that the slot is inside the extended table (num_slots > index)
             let num_slots = Opnd::mem(32, recv, ROBJECT_OFFSET_NUMIV);
-            asm.cmp(num_slots, Opnd::UImm(ivar_index as u64));
+            asm.cmp(num_slots, Opnd::UImm(ivar_index.unwrap() as u64));
             asm.jbe(counted_exit!(ocb, side_exit, getivar_idx_out_of_range).into());
         }
 
@@ -2105,15 +2104,10 @@ fn gen_get_ivar(
         let tbl_opnd = asm.load(Opnd::mem(64, recv, ROBJECT_OFFSET_AS_HEAP_IVPTR));
 
         // Read the ivar from the extended table
-        let ivar_opnd = Opnd::mem(64, tbl_opnd, (SIZEOF_VALUE * ivar_index) as i32);
-
-        // Check that the ivar is not Qundef
-        asm.cmp(ivar_opnd, Qundef.into());
-        let out_val = asm.csel_ne(ivar_opnd, Qnil.into());
+        let ivar_opnd = Opnd::mem(64, tbl_opnd, (SIZEOF_VALUE * ivar_index.unwrap()) as i32);
 
-        // Push the ivar on the stack
         let out_opnd = ctx.stack_push(Type::Unknown);
-        asm.mov(out_opnd, out_val);
+        asm.mov(out_opnd, ivar_opnd);
     }
 
     // Jump to next instruction. This allows guard chains to share the same successor.
@@ -2136,25 +2130,12 @@ fn gen_getinstancevariable(
     let ivar_name = jit_get_arg(jit, 0).as_u64();
 
     let comptime_val = jit_peek_at_self(jit);
-    let comptime_val_klass = comptime_val.class_of();
 
     // Generate a side exit
     let side_exit = get_side_exit(jit, ocb, ctx);
 
     // Guard that the receiver has the same class as the one from compile time.
     let self_asm_opnd = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF);
-    jit_guard_known_klass(
-        jit,
-        ctx,
-        asm,
-        ocb,
-        comptime_val_klass,
-        self_asm_opnd,
-        SelfOpnd,
-        comptime_val,
-        GET_IVAR_MAX_DEPTH,
-        side_exit,
-    );
 
     gen_get_ivar(
         jit,
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 25149ab730..65f398f075 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -120,7 +120,7 @@ extern "C" {
         obj: VALUE,
         v: VALUE,
     ) -> bool;
-    pub fn rb_vm_set_ivar_idx(obj: VALUE, idx: u32, val: VALUE) -> VALUE;
+    pub fn rb_vm_set_ivar_id(obj: VALUE, idx: u32, val: VALUE) -> VALUE;
     pub fn rb_vm_setinstancevariable(iseq: IseqPtr, obj: VALUE, id: ID, val: VALUE, ic: IVC);
     pub fn rb_aliased_callable_method_entry(
         me: *const rb_callable_method_entry_t,
@@ -354,18 +354,26 @@ impl VALUE {
 
     /// Read the flags bits from the RBasic object, then return a Ruby type enum (e.g. RUBY_T_ARRAY)
     pub fn builtin_type(self) -> ruby_value_type {
+        (self.builtin_flags() & (RUBY_T_MASK as usize)) as ruby_value_type
+    }
+
+    pub fn builtin_flags(self) -> usize {
         assert!(!self.special_const_p());
 
         let VALUE(cval) = self;
         let rbasic_ptr = cval as *const RBasic;
         let flags_bits: usize = unsafe { (*rbasic_ptr).flags }.as_usize();
-        (flags_bits & (RUBY_T_MASK as usize)) as ruby_value_type
+        return flags_bits;
     }
 
     pub fn class_of(self) -> VALUE {
         unsafe { CLASS_OF(self) }
     }
 
+    pub fn shape_of(self) -> u32 {
+        unsafe { rb_shape_get_shape_id(self) }
+    }
+
     pub fn as_isize(self) -> isize {
         let VALUE(is) = self;
         is as isize
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index f58bf1ca05..04380f5bd5 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -269,6 +269,29 @@ extern "C" {
 extern "C" {
     pub fn rb_reg_new_ary(ary: VALUE, options: ::std::os::raw::c_int) -> VALUE;
 }
+pub type attr_index_t = u32;
+pub type shape_id_t = u32;
+#[repr(C)]
+pub struct rb_shape {
+    pub parent: *mut rb_shape,
+    pub edges: *mut rb_id_table,
+    pub edge_name: ID,
+    pub iv_count: attr_index_t,
+    pub type_: u8,
+}
+pub type rb_shape_t = rb_shape;
+extern "C" {
+    pub fn rb_shape_get_shape_by_id(shape_id: shape_id_t) -> *mut rb_shape_t;
+}
+extern "C" {
+    pub fn rb_shape_get_shape_id(obj: VALUE) -> shape_id_t;
+}
+extern "C" {
+    pub fn rb_shape_get_iv_index(shape: *mut rb_shape_t, id: ID, value: *mut attr_index_t) -> bool;
+}
+extern "C" {
+    pub fn rb_shape_flags_mask() -> VALUE;
+}
 pub const idDot2: ruby_method_ids = 128;
 pub const idDot3: ruby_method_ids = 129;
 pub const idUPlus: ruby_method_ids = 132;
@@ -572,6 +595,11 @@ pub const OPTIMIZED_METHOD_TYPE_STRUCT_AREF: method_optimized_type = 3;
 pub const OPTIMIZED_METHOD_TYPE_STRUCT_ASET: method_optimized_type = 4;
 pub const OPTIMIZED_METHOD_TYPE__MAX: method_optimized_type = 5;
 pub type method_optimized_type = u32;
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct rb_id_table {
+    _unused: [u8; 0],
+}
 extern "C" {
     pub fn rb_method_entry_at(obj: VALUE, id: ID) -> *const rb_method_entry_t;
 }
@@ -600,9 +628,10 @@ pub struct iseq_inline_constant_cache {
     pub segments: *const ID,
 }
 #[repr(C)]
-#[derive(Debug, Copy, Clone)]
 pub struct iseq_inline_iv_cache_entry {
-    pub entry: *mut rb_iv_index_tbl_entry,
+    pub source_shape_id: shape_id_t,
+    pub dest_shape_id: shape_id_t,
+    pub attr_index: attr_index_t,
 }
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
@@ -698,12 +727,6 @@ extern "C" {
     ) -> *const rb_callable_method_entry_t;
 }
 #[repr(C)]
-pub struct rb_iv_index_tbl_entry {
-    pub index: u32,
-    pub class_serial: rb_serial_t,
-    pub class_value: VALUE,
-}
-#[repr(C)]
 pub struct rb_cvar_class_tbl_entry {
     pub index: u32,
     pub global_cvar_state: rb_serial_t,