Fix assertions in `invalidate_block_version()`, add small repro (#14)

* Fix block invalidation assertions

* Add Alan's small repro for double invalidation bug

2 files changed, 27 insertions, 4 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 5c7e74635a..d02506b2b8 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -172,6 +172,28 @@ assert_equal '1', %q{
     retval
 }
 
+# Code invalidation and opt_getinlinecache
+assert_normal_exit %q{
+  class Foo; end
+
+  # Uses the class constant Foo
+  def use_constant(arg)
+    [Foo.new, arg]
+  end
+
+  def propagate_type
+    i = Array.new
+    i.itself # make it remember that i is on-heap
+    use_constant(i)
+  end
+
+  propagate_type
+  propagate_type
+  use_constant(Foo.new)
+  class Jo; end # bump global constant state
+  use_constant(3)
+}
+
 # Method redefinition (code invalidation) and GC
 assert_equal '7', %q{
     def bar()
diff --git a/yjit_core.c b/yjit_core.c
index 140d0dc64a..eae43c38f7 100644
--- a/yjit_core.c
+++ b/yjit_core.c
@@ -888,7 +888,11 @@ invalidate_block_version(block_t* block)
     {
         branch_t* branch = rb_darray_get(block->incoming, incoming_idx);
         uint32_t target_idx = (branch->dst_addrs[0] == code_ptr)? 0:1;
-        RUBY_ASSERT(!branch->blocks[target_idx] || branch->blocks[target_idx] == block);
+        RUBY_ASSERT(branch->dst_addrs[target_idx] == code_ptr);
+        RUBY_ASSERT(branch->blocks[target_idx] == block);
+
+        // Mark this target as being a stub
+        branch->blocks[target_idx] = NULL;
 
         // Create a stub for this branch target
         branch->dst_addrs[target_idx] = get_branch_target(
@@ -898,9 +902,6 @@ invalidate_block_version(block_t* block)
             target_idx
         );
 
-        // Mark this target as being a stub
-        branch->blocks[target_idx] = NULL;
-
         // Check if the invalidated block immediately follows
         bool target_next = block->start_pos == branch->end_pos;