From ba482d02580d3cfaaf3001e8cbb7e81a9085aed7 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Sun, 22 Jan 2023 20:50:26 -0500
Subject: [PATCH] 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 5f8be84d3c31fa..bc13eb916b52e6 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());
+                }
+            }
         }
     }