From 2e0f3b5546854c98174d4210c0e06c78c8294f98 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Mon, 30 Jan 2023 11:16:33 -0800
Subject: [PATCH] YJIT: Fix BorrowMutError on GC.compact (#7176)

YJIT: Fix BorrowMutError
---
 test/ruby/test_yjit.rb | 16 ++++++++++++++++
 yjit/src/core.rs       | 20 ++++++++++----------
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 6b6ea6619e3787..d390a19cbc60e3 100644
--- 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
index cd6144919d75ca..d10654f1253b71 100644
--- 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());
+                }
+            }
         }
     }