YJIT: Implement new struct accessors (#5161)

* YJIT: Implement optimized_method_struct_aref

* YJIT: Implement struct_aref without method call

Struct member reads can be compiled directly into a memory read (with
either one or two levels of indirection).

* YJIT: Implement optimized struct aset

* YJIT: Update tests for struct access

* YJIT: Add counters for remaining optimized methods

* Check for INT32_MAX overflow

It only takes a struct with 0x7fffffff/8+1 members. Also add some
cheap compile time checks.

* Add tests for non-embedded struct aref/aset

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

5 files changed, 192 insertions, 5 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 8286be74e7..28fe9446ec 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -2455,3 +2455,80 @@ assert_equal 'new', %q{
 
   test
 } if false # disabled for now since OOM crashes in the test harness
+
+# struct aref embedded
+assert_equal '2', %q{
+  def foo(s)
+    s.foo
+  end
+
+  S = Struct.new(:foo)
+  foo(S.new(1))
+  foo(S.new(2))
+}
+
+# struct aref non-embedded
+assert_equal '4', %q{
+  def foo(s)
+    s.d
+  end
+
+  S = Struct.new(:a, :b, :c, :d, :e)
+  foo(S.new(1,2,3,4,5))
+  foo(S.new(1,2,3,4,5))
+}
+
+# struct aset embedded
+assert_equal '123', %q{
+  def foo(s)
+    s.foo = 123
+  end
+
+  s = Struct.new(:foo).new
+  foo(s)
+  s = Struct.new(:foo).new
+  foo(s)
+  s.foo
+}
+
+# struct aset non-embedded
+assert_equal '[1, 2, 3, 4, 5]', %q{
+  def foo(s)
+    s.a = 1
+    s.b = 2
+    s.c = 3
+    s.d = 4
+    s.e = 5
+  end
+
+  S = Struct.new(:a, :b, :c, :d, :e)
+  s = S.new
+  foo(s)
+  s = S.new
+  foo(s)
+  [s.a, s.b, s.c, s.d, s.e]
+}
+
+# struct aref too many args
+assert_equal 'ok', %q{
+  def foo(s)
+    s.foo(:bad)
+  end
+
+  s = Struct.new(:foo).new
+  foo(s) rescue :ok
+  foo(s) rescue :ok
+}
+
+# struct aset too many args
+assert_equal 'ok', %q{
+  def foo(s)
+    s.set_foo(123, :bad)
+  end
+
+  s = Struct.new(:foo) do
+    alias :set_foo :foo=
+  end
+  foo(s) rescue :ok
+  foo(s) rescue :ok
+}
diff --git a/common.mk b/common.mk
index 94d803bc96..9243285689 100644
--- a/common.mk
+++ b/common.mk
@@ -17483,6 +17483,7 @@ yjit.$(OBJEXT): $(top_srcdir)/internal/sanitizers.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/serial.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/static_assert.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/string.h
+yjit.$(OBJEXT): $(top_srcdir)/internal/struct.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/variable.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/vm.h
 yjit.$(OBJEXT): $(top_srcdir)/internal/warnings.h
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index d951f79cfc..0106a09166 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -412,9 +412,20 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
-  def test_invokebuiltin
-    skip "Struct's getter/setter doesn't use invokebuiltin and YJIT doesn't support new logic"
+  def test_struct_aref
+    assert_compiles(<<~RUBY)
+      def foo(obj)
+        obj.foo
+        obj.bar
+      end
+
+      Foo = Struct.new(:foo, :bar)
+      foo(Foo.new(123))
+      foo(Foo.new(123))
+    RUBY
+  end
 
+  def test_struct_aset
     assert_compiles(<<~RUBY)
       def foo(obj)
         obj.foo = 123
diff --git a/yjit.c b/yjit.c
index cef7492e34..33517ca36d 100644
--- a/yjit.c
+++ b/yjit.c
@@ -69,6 +69,9 @@ YJIT_DECLARE_COUNTERS(
     send_zsuper_method,
     send_undef_method,
     send_optimized_method,
+    send_optimized_method_send,
+    send_optimized_method_call,
+    send_optimized_method_block_call,
     send_missing_method,
     send_bmethod,
     send_refined_method,
diff --git a/yjit_codegen.c b/yjit_codegen.c
index add1e2012a..1076b48e8b 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -7,6 +7,7 @@
 #include "internal/object.h"
 #include "internal/sanitizers.h"
 #include "internal/string.h"
+#include "internal/struct.h"
 #include "internal/variable.h"
 #include "internal/re.h"
 #include "probes.h"
@@ -3901,6 +3902,83 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
     return YJIT_END_BLOCK;
 }
 
+static codegen_status_t
+gen_struct_aref(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const rb_callable_method_entry_t *cme, VALUE comptime_recv, VALUE comptime_recv_klass) {
+    if (vm_ci_argc(ci) != 0) {
+        return YJIT_CANT_COMPILE;
+    }
+
+    const unsigned int off = cme->def->body.optimized.index;
+
+    // Confidence checks
+    RUBY_ASSERT_ALWAYS(RB_TYPE_P(comptime_recv, T_STRUCT));
+    RUBY_ASSERT_ALWAYS((long)off < RSTRUCT_LEN(comptime_recv));
+
+    // We are going to use an encoding that takes a 4-byte immediate which
+    // limits the offset to INT32_MAX.
+    {
+        uint64_t native_off = (uint64_t)off * (uint64_t)SIZEOF_VALUE;
+        if (native_off > (uint64_t)INT32_MAX) {
+            return YJIT_CANT_COMPILE;
+        }
+    }
+
+    // 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.
+    bool embedded = FL_TEST_RAW(comptime_recv, RSTRUCT_EMBED_LEN_MASK);
+
+    ADD_COMMENT(cb, "struct aref");
+
+    x86opnd_t recv = ctx_stack_pop(ctx, 1);
+
+    mov(cb, REG0, recv);
+
+    if (embedded) {
+        mov(cb, REG0, member_opnd_idx(REG0, struct RStruct, as.ary, off));
+    }
+    else {
+        mov(cb, REG0, member_opnd(REG0, struct RStruct, as.heap.ptr));
+        mov(cb, REG0, mem_opnd(64, REG0, SIZEOF_VALUE * off));
+    }
+
+    x86opnd_t ret = ctx_stack_push(ctx, TYPE_UNKNOWN);
+    mov(cb, ret, REG0);
+
+    jit_jump_to_next_insn(jit, ctx);
+    return YJIT_END_BLOCK;
+}
+
+static codegen_status_t
+gen_struct_aset(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const rb_callable_method_entry_t *cme, VALUE comptime_recv, VALUE comptime_recv_klass) {
+    if (vm_ci_argc(ci) != 1) {
+        return YJIT_CANT_COMPILE;
+    }
+
+    const unsigned int off = cme->def->body.optimized.index;
+
+    // Confidence checks
+    RUBY_ASSERT_ALWAYS(RB_TYPE_P(comptime_recv, T_STRUCT));
+    RUBY_ASSERT_ALWAYS((long)off < RSTRUCT_LEN(comptime_recv));
+
+    ADD_COMMENT(cb, "struct aset");
+
+    x86opnd_t val = ctx_stack_pop(ctx, 1);
+    x86opnd_t recv = ctx_stack_pop(ctx, 1);
+
+    mov(cb, C_ARG_REGS[0], recv);
+    mov(cb, C_ARG_REGS[1], imm_opnd(off));
+    mov(cb, C_ARG_REGS[2], val);
+    call_ptr(cb, REG0, (void *)RSTRUCT_SET);
+
+    x86opnd_t ret = ctx_stack_push(ctx, TYPE_UNKNOWN);
+    mov(cb, ret, RAX);
+
+    jit_jump_to_next_insn(jit, ctx);
+    return YJIT_END_BLOCK;
+}
+
 const rb_callable_method_entry_t *
 rb_aliased_callable_method_entry(const rb_callable_method_entry_t *me);
 
@@ -4064,8 +4142,24 @@ gen_send_general(jitstate_t *jit, ctx_t *ctx, struct rb_call_data *cd, rb_iseq_t
             return YJIT_CANT_COMPILE;
           // Send family of methods, e.g. call/apply
           case VM_METHOD_TYPE_OPTIMIZED:
-            GEN_COUNTER_INC(cb, send_optimized_method);
-            return YJIT_CANT_COMPILE;
+            switch (cme->def->body.optimized.type) {
+              case OPTIMIZED_METHOD_TYPE_SEND:
+                GEN_COUNTER_INC(cb, send_optimized_method_send);
+                return YJIT_CANT_COMPILE;
+              case OPTIMIZED_METHOD_TYPE_CALL:
+                GEN_COUNTER_INC(cb, send_optimized_method_call);
+                return YJIT_CANT_COMPILE;
+              case OPTIMIZED_METHOD_TYPE_BLOCK_CALL:
+                GEN_COUNTER_INC(cb, send_optimized_method_block_call);
+                return YJIT_CANT_COMPILE;
+              case OPTIMIZED_METHOD_TYPE_STRUCT_AREF:
+                return gen_struct_aref(jit, ctx, ci, cme, comptime_recv, comptime_recv_klass);
+              case OPTIMIZED_METHOD_TYPE_STRUCT_ASET:
+                return gen_struct_aset(jit, ctx, ci, cme, comptime_recv, comptime_recv_klass);
+              default:
+                rb_bug("unknown optimized method type (%d)", cme->def->body.optimized.type);
+                UNREACHABLE_RETURN(YJIT_CANT_COMPILE);
+            }
           case VM_METHOD_TYPE_MISSING:
             GEN_COUNTER_INC(cb, send_missing_method);
             return YJIT_CANT_COMPILE;
@@ -4347,7 +4441,8 @@ gen_objtostring(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
         jit_guard_known_klass(jit, ctx, CLASS_OF(comptime_recv), OPND_STACK(0), comptime_recv, SEND_MAX_DEPTH, side_exit);
         // No work needed. The string value is already on the top of the stack.
         return YJIT_KEEP_COMPILING;
-    } else {
+    }
+    else {
         struct rb_call_data *cd = (struct rb_call_data *)jit_get_arg(jit, 0);
         return gen_send_general(jit, ctx, cd, NULL);
     }