Implement struct aref

4 files changed, 77 insertions, 11 deletions

diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index ceafe65351..21f5a610cb 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -3372,7 +3372,7 @@ module RubyVM::MJIT
       proc_block = proc_t.block
 
       if proc_block.type != C.block_type_iseq
-        asm.incr_counter(:send_optimized_bmethod_not_iseq)
+        asm.incr_counter(:send_bmethod_not_iseq)
         return CantCompile
       end
 
@@ -3389,7 +3389,7 @@ module RubyVM::MJIT
       # Passing a block to a block needs logic different from passing
       # a block to a method and sometimes requires allocation. Bail for now.
       if block_handler != C.VM_BLOCK_HANDLER_NONE
-        asm.incr_counter(:send_optimized_bmethod_blockarg)
+        asm.incr_counter(:send_bmethod_blockarg)
         return CantCompile
       end
 
@@ -3427,8 +3427,7 @@ module RubyVM::MJIT
         asm.incr_counter(:send_optimized_block_call)
         return CantCompile
       when C.OPTIMIZED_METHOD_TYPE_STRUCT_AREF
-        asm.incr_counter(:send_optimized_struct_aref)
-        return CantCompile
+        jit_call_opt_struct_aref(jit, ctx, asm, cme, flags, argc, block_handler, known_recv_class, send_shift:)
       when C.OPTIMIZED_METHOD_TYPE_STRUCT_ASET
         asm.incr_counter(:send_optimized_struct_aset)
         return CantCompile
@@ -3524,6 +3523,50 @@ module RubyVM::MJIT
       return KeepCompiling
     end
 
+    # vm_call_opt_struct_aref
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_call_opt_struct_aref(jit, ctx, asm, cme, flags, argc, block_handler, known_recv_class, send_shift:)
+      if argc != 0
+        asm.incr_counter(:send_optimized_struct_aref_error)
+        return CantCompile
+      end
+
+      off = cme.def.body.optimized.index
+
+      recv_idx = argc # blockarg is not supported
+      recv_idx += send_shift
+      comptime_recv = jit.peek_at_stack(recv_idx)
+
+      # This is a .send call and we need to adjust the stack
+      if flags & C.VM_CALL_OPT_SEND != 0
+        jit_call_opt_send_shift_stack(ctx, asm, argc, send_shift:)
+      end
+
+      # All structs from the same Struct class should have the same
+      # length. So if our comptime_recv is embedded all runtime
+      # structs of the same class should be as well, and the same is
+      # true of the converse.
+      embedded = C.FL_TEST_RAW(comptime_recv, C.RSTRUCT_EMBED_LEN_MASK)
+
+      asm.comment('struct aref')
+      asm.mov(:rax, ctx.stack_pop(1)) # recv
+
+      if embedded
+        asm.mov(:rax, [:rax, C.RStruct.offsetof(:as, :ary) + (C.VALUE.size * off)])
+      else
+        asm.mov(:rax, [:rax, C.RStruct.offsetof(:as, :heap, :ptr)])
+        asm.mov(:rax, [:rax, C.VALUE.size * off])
+      end
+
+      ret = ctx.stack_push
+      asm.mov(ret, :rax)
+
+      jump_to_next_insn(jit, ctx, asm)
+      EndBlock
+    end
+
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def jit_call_opt_send_shift_stack(ctx, asm, argc, send_shift:)
diff --git a/mjit_c.h b/mjit_c.h
index b73da56d0d..4d66251acd 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -4,6 +4,7 @@
 
 #include "ruby/internal/config.h"
 #include "internal/string.h"
+#include "internal/struct.h"
 #include "internal/variable.h"
 #include "vm_core.h"
 #include "vm_callinfo.h"
@@ -161,15 +162,15 @@ MJIT_RUNTIME_COUNTERS(
     send_optimized_call_block,
     send_optimized_call_kwarg,
     send_optimized_call_splat,
-    send_optimized_blockarg,
+    send_optimized_struct_aref_error,
 
+    send_optimized_blockarg,
     send_optimized_block_call,
-    send_optimized_struct_aref,
     send_optimized_struct_aset,
     send_optimized_unknown_type,
 
-    send_optimized_bmethod_not_iseq,
-    send_optimized_bmethod_blockarg,
+    send_bmethod_not_iseq,
+    send_bmethod_blockarg,
 
     invokesuper_me_changed,
     invokesuper_same_me,
diff --git a/mjit_c.rb b/mjit_c.rb
index ece67ac6e3..337bb85714 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -897,6 +897,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ ULONG2NUM(RMODULE_IS_REFINEMENT) }
   end
 
+  def C.RSTRUCT_EMBED_LEN_MASK
+    Primitive.cexpr! %q{ ULONG2NUM(RSTRUCT_EMBED_LEN_MASK) }
+  end
+
   def C.RUBY_FIXNUM_FLAG
     Primitive.cexpr! %q{ ULONG2NUM(RUBY_FIXNUM_FLAG) }
   end
@@ -1042,6 +1046,22 @@ module RubyVM::MJIT # :nodoc: all
     )
   end
 
+  def C.RStruct
+    @RStruct ||= CType::Struct.new(
+      "RStruct", Primitive.cexpr!("SIZEOF(struct RStruct)"),
+      basic: [self.RBasic, Primitive.cexpr!("OFFSETOF((*((struct RStruct *)NULL)), basic)")],
+      as: [CType::Union.new(
+        "", Primitive.cexpr!("SIZEOF(((struct RStruct *)NULL)->as)"),
+        heap: CType::Struct.new(
+          "", Primitive.cexpr!("SIZEOF(((struct RStruct *)NULL)->as.heap)"),
+          len: [CType::Immediate.parse("long"), Primitive.cexpr!("OFFSETOF(((struct RStruct *)NULL)->as.heap, len)")],
+          ptr: [CType::Pointer.new { self.VALUE }, Primitive.cexpr!("OFFSETOF(((struct RStruct *)NULL)->as.heap, ptr)")],
+        ),
+        ary: CType::Pointer.new { self.VALUE },
+      ), Primitive.cexpr!("OFFSETOF((*((struct RStruct *)NULL)), as)")],
+    )
+  end
+
   def C.attr_index_t
     @attr_index_t ||= CType::Immediate.parse("uint32_t")
   end
@@ -1554,13 +1574,13 @@ module RubyVM::MJIT # :nodoc: all
       send_optimized_call_block: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_call_block)")],
       send_optimized_call_kwarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_call_kwarg)")],
       send_optimized_call_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_call_splat)")],
+      send_optimized_struct_aref_error: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_struct_aref_error)")],
       send_optimized_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_blockarg)")],
       send_optimized_block_call: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_block_call)")],
-      send_optimized_struct_aref: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_struct_aref)")],
       send_optimized_struct_aset: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_struct_aset)")],
       send_optimized_unknown_type: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_unknown_type)")],
-      send_optimized_bmethod_not_iseq: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_bmethod_not_iseq)")],
-      send_optimized_bmethod_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_bmethod_blockarg)")],
+      send_bmethod_not_iseq: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_bmethod_not_iseq)")],
+      send_bmethod_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_bmethod_blockarg)")],
       invokesuper_me_changed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), invokesuper_me_changed)")],
       invokesuper_same_me: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), invokesuper_same_me)")],
       getivar_megamorphic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_megamorphic)")],
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index 3f39ddfa48..c4f8bff329 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -448,6 +448,7 @@ generator = BindingGenerator.new(
       RUBY_T_STRING
       RMODULE_IS_REFINEMENT
       RUBY_FL_SINGLETON
+      RSTRUCT_EMBED_LEN_MASK
     ],
     PTR: %w[
       rb_cFalseClass
@@ -468,6 +469,7 @@ generator = BindingGenerator.new(
     RArray
     RBasic
     RObject
+    RStruct
     attr_index_t
     compile_branch
     compile_status