Implement Array#[]

2 files changed, 51 insertions, 3 deletions

diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index ca5397217d..5ea24266fc 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -575,6 +575,22 @@ module RubyVM::MJIT
       insn(opcode: 0xc3)
     end
 
+    def sar(dst, src)
+      case [dst, src]
+      in [Symbol => dst_reg, Integer => src_imm] if r64?(dst_reg) && imm8?(src_imm)
+        # REX.W + C1 /7 ib
+        # MI: Operand 1: ModRM:r/m (r, w), Operand 2: imm8
+        insn(
+          prefix: REX_W,
+          opcode: 0xc1,
+          mod_rm: ModRM[mod: Mod11, reg: 7, rm: dst_reg],
+          imm: imm8(src_imm),
+        )
+      else
+        raise NotImplementedError, "sar: not-implemented operands: #{dst.inspect}, #{src.inspect}"
+      end
+    end
+
     def sub(dst, src)
       case [dst, src]
       # SUB r/m64, imm8
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index dd00ff3e39..15a6efcd11 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -631,8 +631,40 @@ module RubyVM::MJIT
       side_exit = side_exit(jit, ctx)
 
       if comptime_recv.class == Array && fixnum?(comptime_obj)
-        asm.incr_counter(:optaref_array)
-        CantCompile
+        unless Invariants.assume_bop_not_redefined(jit, C.ARRAY_REDEFINED_OP_FLAG, C.BOP_AREF)
+          return CantCompile
+        end
+
+        idx_opnd = ctx.stack_opnd(0)
+        recv_opnd = ctx.stack_opnd(1)
+
+        not_array_exit = counted_exit(side_exit, :optaref_recv_not_array)
+        if jit_guard_known_class(jit, ctx, asm, comptime_recv.class, recv_opnd, comptime_recv, not_array_exit) == CantCompile
+          return CantCompile
+        end
+
+        # Bail if idx is not a FIXNUM
+        asm.mov(:rax, idx_opnd)
+        asm.test(:rax, C.RUBY_FIXNUM_FLAG)
+        asm.jz(counted_exit(side_exit, :optaref_arg_not_fixnum))
+
+        # Call VALUE rb_ary_entry_internal(VALUE ary, long offset).
+        # It never raises or allocates, so we don't need to write to cfp->pc.
+        asm.sar(:rax, 1) # Convert fixnum to int
+        asm.mov(C_ARGS[0], recv_opnd)
+        asm.mov(C_ARGS[1], :rax)
+        asm.call(C.rb_ary_entry_internal)
+
+        # Pop the argument and the receiver
+        ctx.stack_pop(2)
+
+        # Push the return value onto the stack
+        stack_ret = ctx.stack_push
+        asm.mov(stack_ret, C_RET)
+
+        # Let guard chains share the same successor
+        jump_to_next_insn(jit, ctx, asm)
+        EndBlock
       elsif comptime_recv.class == Hash
         unless Invariants.assume_bop_not_redefined(jit, C.HASH_REDEFINED_OP_FLAG, C.BOP_AREF)
           return CantCompile
@@ -641,7 +673,7 @@ module RubyVM::MJIT
         recv_opnd = ctx.stack_opnd(1)
 
         # Guard that the receiver is a Hash
-        not_hash_exit = counted_exit(side_exit, :optaref_not_hash)
+        not_hash_exit = counted_exit(side_exit, :optaref_recv_not_hash)
         if jit_guard_known_class(jit, ctx, asm, comptime_recv.class, recv_opnd, comptime_recv, not_hash_exit) == CantCompile
           return CantCompile
         end