Weak field processing

We allow the C part to discover weak fields during tracing so that we
will clear or forward them during the weak reference processing phase.

Author: wks

mmtk/Cargo.toml

@@ -12,7 +12,7 @@ edition = "2021"
 # Metadata for the Ruby repository
 [package.metadata.ci-repos.ruby]
 repo = "mmtk/ruby" # This is used by actions/checkout, so the format is "owner/repo", not URL.
-rev = "64c07dcd7ca763e25607848df3be909135e02589"
+rev = "fe51e6277b4df810a90c4474f951ed0a263acea0"
 
 [lib]
 name = "mmtk_ruby"

mmtk/cbindgen.toml

@@ -8,8 +8,9 @@ tab_width = 4
 
 usize_is_size_t = true
 
-after_includes = """
+includes = ["ruby/internal/value.h"]
 
+after_includes = """
 typedef struct MMTk_Builder MMTk_Builder;
 typedef struct MMTk_Mutator MMTk_Mutator;
 
@@ -26,6 +27,7 @@ typedef uint32_t MMTk_AllocationSemantics;
 
 [export]
 include = ["HiddenHeader"]
+exclude = ["VALUE"]
 
 [export.rename]
 "MMTKBuilder" = "MMTk_Builder"

mmtk/src/abi.rs

@@ -15,6 +15,27 @@ pub const HIDDEN_SIZE_MASK: usize = 0x0000FFFFFFFFFFFF;
 pub const MMTK_WEAK_CONCURRENT_SET_KIND_FSTRING: u8 = 0;
 pub const MMTK_WEAK_CONCURRENT_SET_KIND_GLOBAL_SYMBOLS: u8 = 1;
 
+pub(crate) const RUBY_IMMEDIATE_MASK: usize = 0x07;
+
+#[allow(non_upper_case_globals)] // Match Ruby definition
+pub(crate) const Qundef: VALUE = VALUE(0x24);
+
+#[repr(transparent)]
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub struct VALUE(pub usize);
+
+impl VALUE {
+    pub fn is_special_const(&self) -> bool {
+        self.0 == 0 || self.0 & RUBY_IMMEDIATE_MASK != 0
+    }
+}
+
+impl From<ObjectReference> for VALUE {
+    fn from(value: ObjectReference) -> Self {
+        Self(value.to_raw_address().as_usize())
+    }
+}
+
 // An opaque type for the C counterpart.
 #[allow(non_camel_case_types)]
 pub struct st_table;

mmtk/src/api.rs

@@ -9,6 +9,7 @@ use crate::abi;
 use crate::abi::HiddenHeader;
 use crate::abi::RawVecOfObjRef;
 use crate::abi::RubyBindingOptions;
+use crate::abi::VALUE;
 use crate::binding;
 use crate::binding::RubyBinding;
 use crate::mmtk;
@@ -421,3 +422,8 @@ pub extern "C" fn mmtk_current_gc_is_nursery() -> bool {
         .generational()
         .is_some_and(|gen| gen.is_current_gc_nursery())
 }
+
+#[no_mangle]
+pub extern "C" fn mmtk_discover_weak_field(field: *mut VALUE) {
+    crate::binding().weak_proc.discover_weak_field(field)
+}

mmtk/src/weak_proc.rs

@@ -2,12 +2,12 @@ use std::sync::Mutex;
 
 use mmtk::{
     scheduler::{GCWork, GCWorker, WorkBucketStage},
-    util::ObjectReference,
+    util::{Address, ObjectReference},
     vm::ObjectTracerContext,
 };
 
 use crate::{
-    abi::{self, GCThreadTLS},
+    abi::{self, GCThreadTLS, Qundef, VALUE},
     extra_assert, is_mmtk_object_safe, upcalls, Ruby,
 };
 
@@ -27,11 +27,15 @@ pub enum WeakConcurrentSetKind {
     GlobalSymbols = abi::MMTK_WEAK_CONCURRENT_SET_KIND_GLOBAL_SYMBOLS,
 }
 
+pub type FieldType = *mut VALUE;
+
 pub struct WeakProcessor {
     /// Objects that needs `obj_free` called when dying.
     /// If it is a bottleneck, replace it with a lock-free data structure,
     /// or add candidates in batch.
     obj_free_candidates: Mutex<Vec<ObjectReference>>,
+    /// Weak fields discovered during the current GC.
+    weak_fields: Mutex<Vec<FieldType>>,
 }
 
 impl Default for WeakProcessor {
@@ -44,6 +48,7 @@ impl WeakProcessor {
     pub fn new() -> Self {
         Self {
             obj_free_candidates: Mutex::new(Vec::new()),
+            weak_fields: Default::default(),
         }
     }
 
@@ -70,6 +75,28 @@ impl WeakProcessor {
         std::mem::take(obj_free_candidates.as_mut())
     }
 
+    pub fn clear_weak_fields(&self) {
+        let mut weak_fields = self
+            .weak_fields
+            .try_lock()
+            .expect("Should not have contention.");
+        weak_fields.clear();
+    }
+
+    pub fn discover_weak_field(&self, field: FieldType) {
+        let mut weak_fields = self.weak_fields.lock().unwrap();
+        weak_fields.push(field);
+        trace!("Pushed weak field {field:?}");
+    }
+
+    pub fn get_all_weak_fields(&self) -> Vec<FieldType> {
+        let mut weak_fields = self
+            .weak_fields
+            .try_lock()
+            .expect("Should not have contention.");
+        std::mem::take(&mut weak_fields)
+    }
+
     pub fn process_weak_stuff(
         &self,
         worker: &mut GCWorker<Ruby>,
@@ -88,6 +115,7 @@ impl WeakProcessor {
             Box::new(UpdateCCRefinementTable) as _,
             // END: Weak tables
             Box::new(UpdateWbUnprotectedObjectsList) as _,
+            Box::new(UpdateWeakFields) as _,
         ]);
 
         if SPECIALIZE_FSTRING_TABLE_PROCESSING {
@@ -317,3 +345,49 @@ impl Forwardable for ObjectReference {
         self.get_forwarded_object().unwrap_or(*self)
     }
 }
+
+struct UpdateWeakFields;
+
+impl GCWork<Ruby> for UpdateWeakFields {
+    fn do_work(&mut self, _worker: &mut GCWorker<Ruby>, _mmtk: &'static mmtk::MMTK<Ruby>) {
+        let weak_fields = crate::binding().weak_proc.get_all_weak_fields();
+
+        let num_fields = weak_fields.len();
+        let mut live = 0usize;
+        let mut forwarded = 0usize;
+
+        debug!("Updating {num_fields} weak fields...");
+
+        for field in weak_fields {
+            let old_value = unsafe { *field };
+            trace!("  Visiting weak field {field:?} -> {old_value:?}");
+
+            if old_value.is_special_const() {
+                continue;
+            }
+
+            let addr = unsafe { Address::from_usize(old_value.0) };
+            if !addr.is_mapped() {
+                panic!("Field {field:?} value {addr} points to unmapped area");
+            }
+            let Some(old_objref) = mmtk::memory_manager::is_mmtk_object(addr) else {
+                panic!("Field {field:?} value {addr} is an invalid object reference");
+            };
+
+            if old_objref.is_reachable() {
+                live += 1;
+                if let Some(new_objref) = old_objref.get_forwarded_object() {
+                    forwarded += 1;
+                    let new_value = VALUE::from(new_objref);
+                    trace!("    Updated weak field {field:?} to {new_value:?}");
+                    unsafe { *field = new_value };
+                }
+            } else {
+                unsafe { *field = Qundef };
+            }
+        }
+
+        debug!("Updated {num_fields} weak fields.  {live} live, {forwarded} forwarded.");
+        probe!(mmtk_ruby, update_weak_fields, num_fields, live, forwarded);
+    }
+}