YJIT: Fix BorrowMutError on GC.compact (#7176)

YJIT: Fix BorrowMutError
ruby · Jan 30, 2023 · 2e0f3b5 · 2e0f3b5
1 parent bc0dc9d
commit 2e0f3b5
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 10 deletions.
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
@@ -1058,6 +1058,22 @@ def foo(_, a, b, c)
     RUBY
   end
 
+  def test_gc_compact_cyclic_branch
+    assert_compiles(<<~'RUBY', result: 2)
+      def foo
+        i = 0
+        while i < 2
+          i += 1
+        end
+        i
+      end
+
+      foo
+      GC.compact
+      foo
+    RUBY
+  end
+
   private
 
   def code_gc_helpers

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
@@ -756,16 +756,6 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
                 *cme_dep = unsafe { rb_gc_location((*cme_dep).into()) }.as_cme();
             }
 
-            // Update outgoing branch entries
-            mem::drop(block); // end mut borrow: target.get_blockid() might borrow it
-            let block = version.borrow();
-            for branch in &block.outgoing {
-                let mut branch = branch.borrow_mut();
-                for target in branch.targets.iter_mut().flatten() {
-                    target.set_iseq(unsafe { rb_gc_location(target.get_blockid().iseq.into()) }.as_iseq());
-                }
-            }
-
             // Walk over references to objects in generated code.
             for offset in &block.gc_obj_offsets {
                 let offset_to_value = offset.as_usize();
@@ -787,6 +777,16 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
                     }
                 }
             }
+
+            // Update outgoing branch entries
+            let outgoing_branches = block.outgoing.clone(); // clone to use after borrow
+            mem::drop(block); // end mut borrow: target.set_iseq and target.get_blockid() might (mut) borrow it
+            for branch in &outgoing_branches {
+                let mut branch = branch.borrow_mut();
+                for target in branch.targets.iter_mut().flatten() {
+                    target.set_iseq(unsafe { rb_gc_location(target.get_blockid().iseq.into()) }.as_iseq());
+                }
+            }
         }
     }