From 064e3450536ad4c4fe49cb282d13713775d4ba8a Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Fri, 19 Feb 2021 11:20:55 -0500
Subject: [PATCH] Check for ::Array, not T_ARRAY in opt_aref

---
 bootstraptest/test_ujit.rb | 17 +++++++++++++++++
 ujit_codegen.c             | 13 +++++++------
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/bootstraptest/test_ujit.rb b/bootstraptest/test_ujit.rb
index 318d07c670..f28d1f37df 100644
--- a/bootstraptest/test_ujit.rb
+++ b/bootstraptest/test_ujit.rb
@@ -197,3 +197,20 @@ assert_equal "nil\n", %q{
 
   p other.read
 }
+
+# Test that opt_aref checks the class of the receiver
+assert_equal ":special\n", %q{
+  def foo(array)
+    array[30]
+  end
+
+  UJIT.install_entry(RubyVM::InstructionSequence.of(method(:foo)))
+
+  special = []
+  def special.[](idx)
+    :special
+  end
+
+  p foo(special)
+  nil
+}
diff --git a/ujit_codegen.c b/ujit_codegen.c
index 6e751a1257..0e311f19e2 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -717,7 +717,7 @@ gen_opt_aref(jitstate_t* jit, ctx_t* ctx)
     uint8_t* side_exit = ujit_side_exit(jit, ctx);
 
     // TODO: make a helper function for guarding on op-not-redefined
-    // Make sure that minus isn't redefined for integers
+    // Make sure that aref isn't redefined for arrays.
     mov(cb, RAX, const_ptr_opnd(ruby_current_vm_ptr));
     test(
         cb,
@@ -740,10 +740,11 @@ gen_opt_aref(jitstate_t* jit, ctx_t* ctx)
     cmp(cb, REG0, imm_opnd(Qnil));
     je_ptr(cb, side_exit);
 
-    // Bail if recv is not an array
-    mov(cb, REG1, mem_opnd(64, REG0, offsetof(struct RBasic, flags)));
-    and(cb, REG1, imm_opnd(T_MASK));
-    cmp(cb, REG1, imm_opnd(T_ARRAY));
+    // Bail if recv has a class other than ::Array.
+    // BOP_AREF check above is only good for ::Array.
+    mov(cb, REG1, mem_opnd(64, REG0, offsetof(struct RBasic, klass)));
+    mov(cb, REG0, const_ptr_opnd((void *)rb_cArray));
+    cmp(cb, REG0, REG1);
     jne_ptr(cb, side_exit);
 
     // Bail if idx is not a FIXNUM
@@ -758,7 +759,7 @@ gen_opt_aref(jitstate_t* jit, ctx_t* ctx)
     // Maintain 16-byte RSP alignment
     sub(cb, RSP, imm_opnd(8));
 
-    mov(cb, RDI, REG0);
+    mov(cb, RDI, recv_opnd);
     sar(cb, REG1, imm_opnd(1)); // Convert fixnum to int
     mov(cb, RSI, REG1);
     call_ptr(cb, REG0, (void *)rb_ary_entry_internal);