6 files changed, 127 insertions, 3 deletions
diff --git a/lib/ruby_vm/mjit/context.rb b/lib/ruby_vm/mjit/context.rb
index c5c5e987bd..86aa03514a 100644
--- a/lib/ruby_vm/mjit/context.rb
+++ b/lib/ruby_vm/mjit/context.rb
@@ -18,5 +18,13 @@ module RubyVM::MJIT
       self.sp_offset -= size
       [SP, C.VALUE.size * self.sp_offset]
     end
+
+    def stack_opnd(depth_from_top)
+      [SP, C.VALUE.size * (self.sp_offset - 1 - depth_from_top)]
+    end
+
+    def sp_opnd(offset_bytes)
+      [SP, (C.VALUE.size * self.sp_offset) + offset_bytes]
+    end
   end
 end
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index 75b033a801..4bc1dde9be 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -18,7 +18,7 @@ module RubyVM::MJIT
       asm.incr_counter(:mjit_insns_count)
       asm.comment("Insn: #{insn.name}")
 
-      # 15/101
+      # 16/101
       case insn.name
       when :nop then nop(jit, ctx, asm)
       # getlocal
@@ -102,7 +102,7 @@ module RubyVM::MJIT
       # opt_ltlt
       # opt_and
       # opt_or
-      # opt_aref
+      when :opt_aref then opt_aref(jit, ctx, asm)
       # opt_aset
       # opt_aset_with
       # opt_aref_with
@@ -488,7 +488,70 @@ module RubyVM::MJIT
     # opt_ltlt
     # opt_and
     # opt_or
-    # opt_aref
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def opt_aref(jit, ctx, asm)
+      cd = C.rb_call_data.new(jit.operand(0))
+      argc = C.vm_ci_argc(cd.ci)
+
+      if argc != 1
+        asm.incr_counter(:optaref_argc_not_one)
+        return CantCompile
+      end
+
+      unless jit.at_current_insn?
+        defer_compilation(jit, ctx, asm)
+        return EndBlock
+      end
+
+      comptime_recv = jit.peek_at_stack(1)
+      comptime_obj  = jit.peek_at_stack(0)
+
+      side_exit = side_exit(jit, ctx)
+
+      if comptime_recv.class == Array && fixnum?(comptime_obj)
+        asm.incr_counter(:optaref_array)
+        CantCompile
+      elsif comptime_recv.class == Hash
+        unless @invariants.assume_bop_not_redefined(jit, C.HASH_REDEFINED_OP_FLAG, C.BOP_AREF)
+          return CantCompile
+        end
+
+        recv_opnd = ctx.stack_opnd(1)
+
+        # Guard that the receiver is a Hash
+        not_hash_exit = counted_exit(side_exit, :optaref_not_hash)
+        if jit_guard_known_class(jit, ctx, asm, comptime_recv.class, recv_opnd, comptime_recv, not_hash_exit) == CantCompile
+          return CantCompile
+        end
+
+        # Prepare to call rb_hash_aref(). It might call #hash on the key.
+        jit_prepare_routine_call(jit, ctx, asm)
+
+        asm.comment('call rb_hash_aref')
+        key_opnd = ctx.stack_opnd(0)
+        recv_opnd = ctx.stack_opnd(1)
+        asm.mov(:rdi, recv_opnd)
+        asm.mov(:rsi, key_opnd)
+        asm.call(C.rb_hash_aref)
+
+        # Pop the key and the receiver
+        ctx.stack_pop(2)
+
+        stack_ret = ctx.stack_push
+        asm.mov(stack_ret, :rax)
+
+        # Let guard chains share the same successor
+        jump_to_next_insn(jit, ctx, asm)
+        EndBlock
+      else
+        asm.incr_counter(:optaref_send)
+        CantCompile
+      end
+    end
+
     # opt_aset
     # opt_aset_with
     # opt_aref_with
@@ -668,6 +731,35 @@ module RubyVM::MJIT
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
+    def jit_prepare_routine_call(jit, ctx, asm)
+      jit_save_pc(jit, asm)
+      jit_save_sp(jit, ctx, asm)
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_save_pc(jit, asm)
+      next_pc = jit.pc + jit.insn.len * C.VALUE.size # Use the next one for backtrace and side exits
+      asm.comment('save PC to CFP')
+      asm.mov(:rax, next_pc)
+      asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax)
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_save_sp(jit, ctx, asm)
+      if ctx.sp_offset != 0
+        asm.comment('save SP to CFP')
+        asm.lea(SP, ctx.sp_opnd(0))
+        asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP)
+        ctx.sp_offset = 0
+      end
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
     def jump_to_next_insn(jit, ctx, asm)
       reset_depth = ctx.dup
       reset_depth.chain_depth = 0
diff --git a/lib/ruby_vm/mjit/stats.rb b/lib/ruby_vm/mjit/stats.rb
index 9a73d2bbdf..343505efb9 100644
--- a/lib/ruby_vm/mjit/stats.rb
+++ b/lib/ruby_vm/mjit/stats.rb
@@ -36,6 +36,7 @@ module RubyVM::MJIT
 
       print_counters(stats, prefix: 'send_', prompt: 'method call exit reasons')
       print_counters(stats, prefix: 'getivar_', prompt: 'getinstancevariable exit reasons')
+      print_counters(stats, prefix: 'optaref_', prompt: 'opt_aref exit reasons')
 
       $stderr.puts "compiled_block_count:  #{format('%10d', stats[:compiled_block_count])}"
       $stderr.puts "side_exit_count:       #{format('%10d', stats[:side_exit_count])}"
diff --git a/mjit_c.h b/mjit_c.h
index 99fd9df8e7..4ce846009a 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -132,6 +132,11 @@ MJIT_RUNTIME_COUNTERS(
     getivar_special_const,
     getivar_too_complex,
 
+    optaref_argc_not_one,
+    optaref_array,
+    optaref_not_hash,
+    optaref_send,
+
     compiled_block_count
 )
 #undef MJIT_RUNTIME_COUNTERS
diff --git a/mjit_c.rb b/mjit_c.rb
index 0b90071cac..d6704e762c 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -127,6 +127,10 @@ module RubyVM::MJIT # :nodoc: all
       Primitive.cexpr! 'RBOOL(FL_TEST_RAW((VALUE)NUM2SIZET(_value), (VALUE)NUM2SIZET(flags)))'
     end
 
+    def rb_hash_aref
+      Primitive.cexpr! 'SIZET2NUM((size_t)rb_hash_aref)'
+    end
+
     #========================================================================================
     #
     # Old stuff
@@ -287,6 +291,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ INT2NUM(VM_ENV_DATA_INDEX_SPECVAL) }
   end
 
+  def C.BOP_AREF
+    Primitive.cexpr! %q{ UINT2NUM(BOP_AREF) }
+  end
+
   def C.BOP_LT
     Primitive.cexpr! %q{ UINT2NUM(BOP_LT) }
   end
@@ -299,6 +307,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ UINT2NUM(BOP_PLUS) }
   end
 
+  def C.HASH_REDEFINED_OP_FLAG
+    Primitive.cexpr! %q{ UINT2NUM(HASH_REDEFINED_OP_FLAG) }
+  end
+
   def C.INTEGER_REDEFINED_OP_FLAG
     Primitive.cexpr! %q{ UINT2NUM(INTEGER_REDEFINED_OP_FLAG) }
   end
@@ -887,6 +899,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_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_not_hash: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), optaref_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)")],
     )
   end
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index 1e1fd96bda..5cee397cb3 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -348,9 +348,11 @@ generator = BindingGenerator.new(
       VM_ENV_DATA_INDEX_SPECVAL
     ],
     UINT: %w[
+      BOP_AREF
       BOP_LT
       BOP_MINUS
       BOP_PLUS
+      HASH_REDEFINED_OP_FLAG
       INTEGER_REDEFINED_OP_FLAG
       METHOD_VISI_PRIVATE
       METHOD_VISI_PROTECTED