From 93b6997103b34750b2d84df07e09586fe1de0649 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Tue, 28 Mar 2023 17:21:40 -0400
Subject: [PATCH] YJIT: Fix overlapping &mut in Assembler::code_gc()

Making overlapping `&mut`s triggers Undefined Bahavior. This function
previously had them through `cb` and `ocb` aliasing with `self` or live
references in the caller.

To fix the overlap, take `ocb` as a parameter and don't use `get_inline_cb()`
in the body of the function.
---
 yjit/src/asm/mod.rs | 16 ++++++----------
 yjit/src/core.rs    |  6 +++---
 yjit/src/yjit.rs    |  3 ++-
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index bf18dd5672..346ada7719 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -583,7 +583,7 @@ impl CodeBlock {
     }
 
     /// Code GC. Free code pages that are not on stack and reuse them.
-    pub fn code_gc(&mut self) {
+    pub fn code_gc(&mut self, ocb: &mut OutlinedCb) {
         // The previous code GC failed to free any pages. Give up.
         if self.freed_pages.as_ref() == &Some(vec![]) {
             return;
@@ -631,15 +631,11 @@ impl CodeBlock {
         freed_pages.append(&mut virtual_pages);
 
         if let Some(&first_page) = freed_pages.first() {
-            let mut cb = CodegenGlobals::get_inline_cb();
-            cb.write_pos = cb.get_page_pos(first_page);
-            cb.dropped_bytes = false;
-            cb.clear_comments();
-
-            let mut ocb = CodegenGlobals::get_outlined_cb().unwrap();
-            ocb.write_pos = ocb.get_page_pos(first_page);
-            ocb.dropped_bytes = false;
-            ocb.clear_comments();
+            for cb in [&mut *self, ocb.unwrap()] {
+                cb.write_pos = cb.get_page_pos(first_page);
+                cb.dropped_bytes = false;
+                cb.clear_comments();
+            }
         }
 
         // Track which pages are free.
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 334b037d4c..f1b6217dd7 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2049,7 +2049,7 @@ pub fn gen_entry_point(iseq: IseqPtr, ec: EcPtr) -> Option<CodePtr> {
         // Compilation failed
         None => {
             // Trigger code GC. This entry point will be recompiled later.
-            cb.code_gc();
+            cb.code_gc(ocb);
             return None;
         }
 
@@ -2146,7 +2146,7 @@ fn entry_stub_hit_body(entry_ptr: *const c_void, ec: EcPtr) -> Option<*const u8>
             Some(blockref) => blockref,
             None => { // No space
                 // Trigger code GC. This entry point will be recompiled later.
-                cb.code_gc();
+                cb.code_gc(ocb);
                 return None;
             }
         }
@@ -2426,7 +2426,7 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
             // because incomplete code could be used when cb.dropped_bytes is flipped
             // by code GC. So this place, after all compilation, is the safest place
             // to hook code GC on branch_stub_hit.
-            cb.code_gc();
+            cb.code_gc(ocb);
 
             // Failed to service the stub by generating a new block so now we
             // need to exit to the interpreter at the stubbed location. We are
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 10bddcc600..d75b9db0ee 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -140,7 +140,8 @@ pub extern "C" fn rb_yjit_code_gc(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     }
 
     let cb = CodegenGlobals::get_inline_cb();
-    cb.code_gc();
+    let ocb = CodegenGlobals::get_outlined_cb();
+    cb.code_gc(ocb);
     Qnil
 }