3 files changed, 104 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])}"