Implement `opt_aref_with` instruction (#8118)

Implement gen_opt_aref_with

Vm opt_aref_with is available

Test opt_aref_with

Stats for opt_aref_with

Co-authored-by: jhawthorn <jhawthorn@github.com>

6 files changed, 61 insertions, 0 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 402b5c512a..80d0ef00d8 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1100,6 +1100,18 @@ assert_equal '[42, :default]', %q{
   ]
 }
 
+# Test default value block for Hash with opt_aref_with
+assert_equal "false", %q{
+  def index_with_string(h)
+    h["foo"]
+  end
+
+  h = Hash.new { |h, k| k.frozen? }
+
+  index_with_string(h)
+  index_with_string(h)
+}
+
 # A regression test for making sure cfp->sp is proper when
 # hitting stubs. See :stub-sp-flush:
 assert_equal 'ok', %q{
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 6e8070199d..61f86b7178 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1289,6 +1289,14 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_opt_aref_with
+    assert_compiles(<<~RUBY, insns: %i[opt_aref_with], result: "bar")
+      h = {"foo" => "bar"}
+
+      h["foo"]
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index c676399c59..d4474e7475 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -6259,6 +6259,12 @@ vm_opt_aref_with(VALUE recv, VALUE key)
     }
 }
 
+VALUE
+rb_vm_opt_aref_with(VALUE recv, VALUE key)
+{
+    return vm_opt_aref_with(recv, key);
+}
+
 static VALUE
 vm_opt_aset_with(VALUE recv, VALUE key, VALUE val)
 {
diff --git a/yjit.rb b/yjit.rb
index 394554b7e3..3308be60c6 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -259,6 +259,7 @@ module RubyVM::YJIT
       print_counters(stats, out: out, prefix: 'setivar_', prompt: 'setinstancevariable exit reasons:')
       print_counters(stats, out: out, prefix: 'definedivar_', prompt: 'definedivar exit reasons:')
       print_counters(stats, out: out, prefix: 'opt_aref_', prompt: 'opt_aref exit reasons: ')
+      print_counters(stats, out: out, prefix: 'opt_aref_with_', prompt: 'opt_aref_with exit reasons: ')
       print_counters(stats, out: out, prefix: 'expandarray_', prompt: 'expandarray exit reasons: ')
       print_counters(stats, out: out, prefix: 'opt_getinlinecache_', prompt: 'opt_getinlinecache exit reasons: ')
       print_counters(stats, out: out, prefix: 'invalidate_', prompt: 'invalidation reasons: ')
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index e6d65ef423..6f5f20769d 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -3064,6 +3064,37 @@ fn gen_opt_aset(
     }
 }
 
+fn gen_opt_aref_with(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    _ocb: &mut OutlinedCb,
+) -> Option<CodegenStatus>{
+    jit_prepare_routine_call(jit, asm);
+
+    let key_opnd = Opnd::Value(jit.get_arg(0));
+    let recv_opnd = asm.stack_pop(1);
+
+    extern "C" {
+        fn rb_vm_opt_aref_with(recv: VALUE, key: VALUE) -> VALUE;
+    }
+
+    let val_opnd = asm.ccall(
+        rb_vm_opt_aref_with as *const u8,
+        vec![
+            recv_opnd,
+            key_opnd
+        ],
+    );
+
+    asm.cmp(val_opnd, Qundef.into());
+    asm.je(Target::side_exit(Counter::opt_aref_with_qundef));
+
+    let top = asm.stack_push(Type::Unknown);
+    asm.mov(top, val_opnd);
+
+    return Some(KeepCompiling);
+}
+
 fn gen_opt_and(
     jit: &mut JITState,
     asm: &mut Assembler,
@@ -8040,6 +8071,7 @@ fn get_gen_fn(opcode: VALUE) -> Option<InsnGenFn> {
         YARVINSN_opt_neq => Some(gen_opt_neq),
         YARVINSN_opt_aref => Some(gen_opt_aref),
         YARVINSN_opt_aset => Some(gen_opt_aset),
+        YARVINSN_opt_aref_with => Some(gen_opt_aref_with),
         YARVINSN_opt_mult => Some(gen_opt_mult),
         YARVINSN_opt_div => Some(gen_opt_div),
         YARVINSN_opt_ltlt => Some(gen_opt_ltlt),
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index c113f428cd..f7bb5c4f33 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -344,6 +344,8 @@ make_counters! {
     opt_aset_not_fixnum,
     opt_aset_not_hash,
 
+    opt_aref_with_qundef,
+
     opt_case_dispatch_megamorphic,
 
     opt_getinlinecache_miss,