YJIT: Fix incomplete invalidation from opt_setinlinecache

As part of YJIT's strategy for promoting Ruby constant expressions into
constants in the output native code, the interpreter calls
rb_yjit_constant_ic_update() from opt_setinlinecache.

The block invalidation loop indirectly calls rb_darray_remove_unordered(),
which does a shuffle remove. Because of this, looping with an
incrementing counter like done previously can miss some elements in the
array. Repeatedly invalidate the first element instead.

The bug this commit resolves does not seem to cause crashes or divergent
behaviors.

Co-authored-by: Jemma Issroff <jemmaissroff@gmail.com>

4 files changed, 110 insertions, 12 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 18d1f85b93..140365ee0c 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1,3 +1,88 @@
+assert_equal '0', %q{
+  # This is a regression test for incomplete invalidation from
+  # opt_setinlinecache. This test might be brittle, so
+  # feel free to remove it in the future if it's too annoying.
+  # This test assumes --yjit-call-threshold=2.
+  module M
+    Foo = 1
+    def foo
+      Foo
+    end
+
+    def pin_self_type_then_foo
+      _ = @foo
+      foo
+    end
+
+    def only_ints
+      1 + self
+      foo
+    end
+  end
+
+  class Integer
+    include M
+  end
+
+  class Sub
+    include M
+  end
+
+  foo_method = M.instance_method(:foo)
+
+  dbg = ->(message) do
+    return # comment this out to get printouts
+
+    $stderr.puts RubyVM::YJIT.disasm(foo_method)
+    $stderr.puts message
+  end
+
+  2.times { 42.only_ints }
+
+  dbg["There should be two versions of getinlineache"]
+
+  module M
+    remove_const(:Foo)
+  end
+
+  dbg["There should be no getinlinecaches"]
+
+  2.times do
+    42.only_ints
+  rescue NameError => err
+    _ = "caught name error #{err}"
+  end
+
+  dbg["There should be one version of getinlineache"]
+
+  2.times do
+    Sub.new.pin_self_type_then_foo
+  rescue NameError
+    _ = 'second specialization'
+  end
+
+  dbg["There should be two versions of getinlineache"]
+
+  module M
+    Foo = 1
+  end
+
+  dbg["There should still be two versions of getinlineache"]
+
+  42.only_ints
+
+  dbg["There should be no getinlinecaches"]
+
+  # Find name of the first VM instruction in M#foo.
+  insns = RubyVM::InstructionSequence.of(foo_method).to_a
+  if defined?(RubyVM::YJIT.blocks_for) && (insns.last.find { Array === _1 }&.first == :opt_getinlinecache)
+    RubyVM::YJIT.blocks_for(RubyVM::InstructionSequence.of(foo_method))
+      .filter { _1.iseq_start_index == 0 }.count
+  else
+    0 # skip the test
+  end
+}
+
 # Check that frozen objects are respected
 assert_equal 'great', %q{
   class Foo
diff --git a/yjit.c b/yjit.c
index 3354c24341..2f8ec039ea 100644
--- a/yjit.c
+++ b/yjit.c
@@ -187,7 +187,7 @@ void rb_yjit_iseq_mark(const struct rb_iseq_constant_body *body) {}
 void rb_yjit_iseq_update_references(const struct rb_iseq_constant_body *body) {}
 void rb_yjit_iseq_free(const struct rb_iseq_constant_body *body) {}
 void rb_yjit_before_ractor_spawn(void) {}
-void rb_yjit_constant_ic_update(const rb_iseq_t *iseq, IC ic) {}
+void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic) {}
 void rb_yjit_tracing_invalidate_all(void) {}
 
 #endif // if JIT_ENABLED && PLATFORM_SUPPORTED_P
diff --git a/yjit.h b/yjit.h
index c254c02cd8..357acdd4b4 100644
--- a/yjit.h
+++ b/yjit.h
@@ -60,7 +60,7 @@ void rb_yjit_iseq_mark(const struct rb_iseq_constant_body *body);
 void rb_yjit_iseq_update_references(const struct rb_iseq_constant_body *body);
 void rb_yjit_iseq_free(const struct rb_iseq_constant_body *body);
 void rb_yjit_before_ractor_spawn(void);
-void rb_yjit_constant_ic_update(const rb_iseq_t *iseq, IC ic);
+void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic);
 void rb_yjit_tracing_invalidate_all(void);
 
 #endif // #ifndef YJIT_H
diff --git a/yjit_iface.c b/yjit_iface.c
index 5c9e024a5f..dee0c42500 100644
--- a/yjit_iface.c
+++ b/yjit_iface.c
@@ -604,7 +604,7 @@ rb_yjit_constant_state_changed(void)
 // Invalidate the block for the matching opt_getinlinecache so it could regenerate code
 // using the new value in the constant cache.
 void
-rb_yjit_constant_ic_update(const rb_iseq_t *iseq, IC ic)
+rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic)
 {
     if (!rb_yjit_enabled_p()) return;
 
@@ -617,27 +617,40 @@ rb_yjit_constant_ic_update(const rb_iseq_t *iseq, IC ic)
     RB_VM_LOCK_ENTER();
     rb_vm_barrier(); // Stop other ractors since we are going to patch machine code.
     {
-
         const struct rb_iseq_constant_body *const body = iseq->body;
         VALUE *code = body->iseq_encoded;
+        const unsigned get_insn_idx = ic->get_insn_idx;
 
         // This should come from a running iseq, so direct threading translation
         // should have been done
         RUBY_ASSERT(FL_TEST((VALUE)iseq, ISEQ_TRANSLATED));
-        RUBY_ASSERT(ic->get_insn_idx < body->iseq_size);
-        RUBY_ASSERT(rb_vm_insn_addr2insn((const void *)code[ic->get_insn_idx]) == BIN(opt_getinlinecache));
+        RUBY_ASSERT(get_insn_idx < body->iseq_size);
+        RUBY_ASSERT(rb_vm_insn_addr2insn((const void *)code[get_insn_idx]) == BIN(opt_getinlinecache));
 
         // Find the matching opt_getinlinecache and invalidate all the blocks there
         RUBY_ASSERT(insn_op_type(BIN(opt_getinlinecache), 1) == TS_IC);
-        if (ic == (IC)code[ic->get_insn_idx + 1 + 1]) {
-            rb_yjit_block_array_t getinlinecache_blocks = yjit_get_version_array(iseq, ic->get_insn_idx);
-            rb_darray_for(getinlinecache_blocks, i) {
-                block_t *block = rb_darray_get(getinlinecache_blocks, i);
-                invalidate_block_version(block);
+        if (ic == (IC)code[get_insn_idx + 1 + 1]) {
+            rb_yjit_block_array_t getinlinecache_blocks = yjit_get_version_array(iseq, get_insn_idx);
+
+            // Put a bound for loop below to be defensive
+            const int32_t initial_version_count = rb_darray_size(getinlinecache_blocks);
+            for (int32_t iteration=0; iteration<initial_version_count; ++iteration) {
+                getinlinecache_blocks = yjit_get_version_array(iseq, get_insn_idx);
+
+                if (rb_darray_size(getinlinecache_blocks) > 0) {
+                    block_t *block = rb_darray_get(getinlinecache_blocks, 0);
+                    invalidate_block_version(block);
 #if YJIT_STATS
-                yjit_runtime_counters.invalidate_constant_ic_fill++;
+                    yjit_runtime_counters.invalidate_constant_ic_fill++;
 #endif
+                }
+                else {
+                    break;
+                }
             }
+
+            // All versions at get_insn_idx should now be gone
+            RUBY_ASSERT(0 == rb_darray_size(yjit_get_version_array(iseq, get_insn_idx)));
         }
         else {
             RUBY_ASSERT(false && "ic->get_insn_diex not set properly");