Implement Array#[]

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

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

mjit_c.h

@@ -149,10 +149,10 @@ MJIT_RUNTIME_COUNTERS(
     getivar_special_const,
     getivar_too_complex,
 
+    optaref_arg_not_fixnum,
     optaref_argc_not_one,
-    optaref_array,
-    optaref_hash,
-    optaref_not_hash,
+    optaref_recv_not_array,
+    optaref_recv_not_hash,
     optaref_send,
 
     compiled_block_count

mjit_c.rb

@@ -145,6 +145,10 @@ module RubyVM::MJIT # :nodoc: all
       }
     end
 
+    def rb_ary_entry_internal
+      Primitive.cexpr! 'SIZET2NUM((size_t)rb_ary_entry_internal)'
+    end
+
     def mjit_for_each_iseq(&block)
       Primitive.mjit_for_each_iseq(block)
     end
@@ -317,6 +321,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ INT2NUM(VM_ENV_DATA_INDEX_SPECVAL) }
   end
 
+  def C.ARRAY_REDEFINED_OP_FLAG
+    Primitive.cexpr! %q{ UINT2NUM(ARRAY_REDEFINED_OP_FLAG) }
+  end
+
   def C.BOP_AREF
     Primitive.cexpr! %q{ UINT2NUM(BOP_AREF) }
   end
@@ -1024,10 +1032,10 @@ module RubyVM::MJIT # :nodoc: all
       getivar_not_t_object: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_not_t_object)")],
       getivar_special_const: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_special_const)")],
       getivar_too_complex: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_too_complex)")],
+      optaref_arg_not_fixnum: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_arg_not_fixnum)")],
       optaref_argc_not_one: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_argc_not_one)")],
-      optaref_array: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_array)")],
-      optaref_hash: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_hash)")],
-      optaref_not_hash: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_not_hash)")],
+      optaref_recv_not_array: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_recv_not_array)")],
+      optaref_recv_not_hash: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_recv_not_hash)")],
       optaref_send: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_send)")],
       compiled_block_count: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), compiled_block_count)")],
     )

tool/mjit/bindgen.rb

@@ -355,6 +355,7 @@ generator = BindingGenerator.new(
       BOP_LT
       BOP_MINUS
       BOP_PLUS
+      ARRAY_REDEFINED_OP_FLAG
       HASH_REDEFINED_OP_FLAG
       INTEGER_REDEFINED_OP_FLAG
       METHOD_VISI_PRIVATE