Introduce ArcFlexMut for Space references

Cargo.toml

@@ -40,6 +40,7 @@ mimalloc-sys = { version = "0.1.6", optional = true }
 mmtk-macros = { version = "0.20.0", path = "macros/" }
 num_cpus = "1.8"
 num-traits = "0.2"
+peace-lock = "0.1.3"
 pfm = { version = "0.1.0-beta.3", optional = true }
 probe = "0.5"
 portable-atomic = "1.4.3"
@@ -72,6 +73,9 @@ perf_counter = ["pfm"]
 # Do not modify the following line - ci-common.sh matches it
 # -- Non mutually exclusive features --
 
+# Turn on checks for the type ArcFlexMut (we use peace-lock for the type so just turn on check for peace-lock).
+check_flex_mut = ["peace-lock/check"]
+
 # spaces with different semantics
 
 # A VM-allocated/managed space. A binding could use this for their boot image, metadata space, etc.
@@ -129,7 +133,7 @@ nogc_no_zeroing = ["nogc_lock_free"]
 single_worker = []
 
 # To run expensive comprehensive runtime checks, such as checking duplicate edges
-extreme_assertions = []
+extreme_assertions = ["check_flex_mut"]
 
 # Enable multiple spaces for NoGC, each allocator maps to an individual ImmortalSpace.
 nogc_multi_space = []

docs/userguide/src/tutorial/code/mygc_semispace/gc_work.rs

@@ -61,15 +61,15 @@ impl<VM: VMBinding> ProcessEdgesWork for MyGCProcessEdges<VM> {
         }
         let worker = self.worker();
         let queue = &mut self.base.nodes;
-        if self.plan.tospace().in_space(object) {
-            self.plan.tospace().trace_object(
+        if self.plan.tospace().read().in_space(object) {
+            self.plan.tospace().read().trace_object(
                 queue,
                 object,
                 Some(CopySemantics::DefaultCopy),
                 worker,
             )
-        } else if self.plan.fromspace().in_space(object) {
-            self.plan.fromspace().trace_object(
+        } else if self.plan.fromspace().read().in_space(object) {
+            self.plan.fromspace().read().trace_object(
                 queue,
                 object,
                 Some(CopySemantics::DefaultCopy),

docs/userguide/src/tutorial/code/mygc_semispace/global.rs

@@ -15,6 +15,7 @@ use crate::util::copy::*;
 use crate::util::heap::VMRequest;
 use crate::util::metadata::side_metadata::SideMetadataContext;
 use crate::util::opaque_pointer::*;
+use crate::util::rust_util::flex_mut::ArcFlexMut;
 use crate::vm::VMBinding;
 use enum_map::EnumMap;
 use std::sync::atomic::{AtomicBool, Ordering}; // Add
@@ -32,10 +33,10 @@ pub struct MyGC<VM: VMBinding> {
     pub hi: AtomicBool,
     #[space]
     #[copy_semantics(CopySemantics::DefaultCopy)]
-    pub copyspace0: CopySpace<VM>,
+    pub copyspace0: ArcFlexMut<CopySpace<VM>>,
     #[space]
     #[copy_semantics(CopySemantics::DefaultCopy)]
-    pub copyspace1: CopySpace<VM>,
+    pub copyspace1: ArcFlexMut<CopySpace<VM>>,
     #[parent]
     pub common: CommonPlan<VM>,
 }
@@ -66,7 +67,7 @@ impl<VM: VMBinding> Plan for MyGC<VM> {
             },
             space_mapping: vec![
                 // The tospace argument doesn't matter, we will rebind before a GC anyway.
-                (CopySelector::CopySpace(0), &self.copyspace0)
+                (CopySelector::CopySpace(0), self.copyspace0.clone().into_dyn_space())
             ],
             constraints: &MYGC_CONSTRAINTS,
         }
@@ -92,48 +93,48 @@ impl<VM: VMBinding> Plan for MyGC<VM> {
 
     // Modify
     // ANCHOR: prepare
-    fn prepare(&mut self, tls: VMWorkerThread) {
+    fn prepare(&self, tls: VMWorkerThread) {
         self.common.prepare(tls, true);
 
         self.hi
             .store(!self.hi.load(Ordering::SeqCst), Ordering::SeqCst);
         // Flips 'hi' to flip space definitions
         let hi = self.hi.load(Ordering::SeqCst);
-        self.copyspace0.prepare(hi);
-        self.copyspace1.prepare(!hi);
+        self.copyspace0.read().prepare(hi);
+        self.copyspace1.read().prepare(!hi);
 
-        self.fromspace_mut()
+        self.fromspace().write()
             .set_copy_for_sft_trace(Some(CopySemantics::DefaultCopy));
-        self.tospace_mut().set_copy_for_sft_trace(None);
+        self.tospace().write().set_copy_for_sft_trace(None);
     }
     // ANCHOR_END: prepare
 
     // Add
     // ANCHOR: prepare_worker
     fn prepare_worker(&self, worker: &mut GCWorker<VM>) {
-        unsafe { worker.get_copy_context_mut().copy[0].assume_init_mut() }.rebind(self.tospace());
+        unsafe { worker.get_copy_context_mut().copy[0].assume_init_mut() }.rebind(self.tospace().clone());
     }
     // ANCHOR_END: prepare_worker
 
     // Modify
     // ANCHOR: release
-    fn release(&mut self, tls: VMWorkerThread) {
+    fn release(&self, tls: VMWorkerThread) {
         self.common.release(tls, true);
-        self.fromspace().release();
+        self.fromspace().read().release();
     }
     // ANCHOR_END: release
 
     // Modify
     // ANCHOR: plan_get_collection_reserve
     fn get_collection_reserved_pages(&self) -> usize {
-        self.tospace().reserved_pages()
+        self.tospace().read().reserved_pages()
     }
     // ANCHOR_END: plan_get_collection_reserve
 
     // Modify
     // ANCHOR: plan_get_used_pages
     fn get_used_pages(&self) -> usize {
-        self.tospace().reserved_pages() + self.common.get_used_pages()
+        self.tospace().read().reserved_pages() + self.common.get_used_pages()
     }
     // ANCHOR_END: plan_get_used_pages
 
@@ -170,9 +171,9 @@ impl<VM: VMBinding> MyGC<VM> {
         let res = MyGC {
             hi: AtomicBool::new(false),
             // ANCHOR: copyspace_new
-            copyspace0: CopySpace::new(plan_args.get_space_args("copyspace0", true, VMRequest::discontiguous()), false),
+            copyspace0: ArcFlexMut::new(CopySpace::new(plan_args.get_space_args("copyspace0", true, VMRequest::discontiguous()), false)),
             // ANCHOR_END: copyspace_new
-            copyspace1: CopySpace::new(plan_args.get_space_args("copyspace1", true, VMRequest::discontiguous()), true),
+            copyspace1: ArcFlexMut::new(CopySpace::new(plan_args.get_space_args("copyspace1", true, VMRequest::discontiguous()), true)),
             common: CommonPlan::new(plan_args),
         };
 
@@ -183,36 +184,20 @@ impl<VM: VMBinding> MyGC<VM> {
     // ANCHOR_END: plan_new
 
     // ANCHOR: plan_space_access
-    pub fn tospace(&self) -> &CopySpace<VM> {
+    pub fn tospace(&self) -> &ArcFlexMut<CopySpace<VM>> {
         if self.hi.load(Ordering::SeqCst) {
             &self.copyspace1
         } else {
             &self.copyspace0
         }
     }
 
-    pub fn fromspace(&self) -> &CopySpace<VM> {
+    pub fn fromspace(&self) -> &ArcFlexMut<CopySpace<VM>> {
         if self.hi.load(Ordering::SeqCst) {
             &self.copyspace0
         } else {
             &self.copyspace1
         }
     }
-
-    pub fn tospace_mut(&mut self) -> &mut CopySpace<VM> {
-        if self.hi.load(Ordering::SeqCst) {
-            &mut self.copyspace1
-        } else {
-            &mut self.copyspace0
-        }
-    }
-
-    pub fn fromspace_mut(&mut self) -> &mut CopySpace<VM> {
-        if self.hi.load(Ordering::SeqCst) {
-            &mut self.copyspace0
-        } else {
-            &mut self.copyspace1
-        }
-    }
     // ANCHOR_END: plan_space_access
 }

docs/userguide/src/tutorial/code/mygc_semispace/mutator.rs

@@ -44,7 +44,7 @@ pub fn mygc_mutator_release<VM: VMBinding>(
             .plan
             .downcast_ref::<MyGC<VM>>()
             .unwrap()
-            .tospace(),
+            .tospace().clone().into_dyn_space(),
     );
 }
 // ANCHOR_END: release
@@ -78,7 +78,7 @@ pub fn create_mygc_mutator<VM: VMBinding>(
         // ANCHOR: space_mapping
         space_mapping: Box::new({
             let mut vec = create_space_mapping(RESERVED_ALLOCATORS, true, mygc);
-            vec.push((AllocatorSelector::BumpPointer(0), mygc.tospace()));
+            vec.push((AllocatorSelector::BumpPointer(0), mygc.tospace().clone().into_dyn_space()));
             vec
         }),
         // ANCHOR_END: space_mapping

docs/userguide/src/tutorial/mygc/ss/alloc.md

@@ -56,8 +56,11 @@ Change `pub struct MyGC<VM: VMBinding>` to add new instance variables.
    1. Delete the existing fields in the constructor.
    2. Add `pub hi: AtomicBool,`. This is a thread-safe bool, indicating which 
    copyspace is the tospace.
-   3. Add `pub copyspace0: CopySpace<VM>,` 
-   and `pub copyspace1: CopySpace<VM>,`. These are the two copyspaces.
+   3. Add `pub copyspace0: ArcFlexMut<CopySpace<VM>>,` 
+   and `pub copyspace1: ArcFlexMut<CopySpace<VM>>,`. These are the two copyspaces.
+   We use the type `ArcFlexMut` from `mmtk::util::rust_util::flex_mut`, which
+   allows us to share the reference among different types and allows us to
+   flexibly acquire mutable references.
    4. Add `pub common: CommonPlan<VM>,`.
     This holds an instance of the common plan.
 

macros/src/has_spaces_impl.rs

@@ -45,11 +45,13 @@ pub(crate) fn generate_impl_items<'a>(
         let f_ident = f.ident.as_ref().unwrap();
 
         let visitor = quote! {
-            __func(&self.#f_ident);
+            let space = self.#f_ident.read();
+            __func(&*space);
         };
 
         let visitor_mut = quote! {
-            __func(&mut self.#f_ident);
+            let mut space = self.#f_ident.write();
+            __func(&mut *space);
         };
 
         space_visitors.push(visitor);
@@ -75,7 +77,7 @@ pub(crate) fn generate_impl_items<'a>(
             #parent_visitor
         }
 
-        fn for_each_space_mut(&mut self, __func: &mut dyn FnMut(&mut dyn Space<VM>)) {
+        fn for_each_space_mut(&self, __func: &mut dyn FnMut(&mut dyn Space<VM>)) {
             #(#space_visitors_mut)*
             #parent_visitor_mut
         }

macros/src/plan_trace_object_impl.rs

@@ -45,7 +45,6 @@ pub(crate) fn generate_trace_object<'a>(
     // Generate a check with early return for each space
     let space_field_handler = space_fields.iter().map(|f| {
         let f_ident = f.ident.as_ref().unwrap();
-        let f_ty = &f.ty;
 
         // Figure out copy
         let maybe_copy_semantics_attr = util::get_field_attribute(f, "copy_semantics");
@@ -71,8 +70,11 @@ pub(crate) fn generate_trace_object<'a>(
         };
 
         quote! {
-            if self.#f_ident.in_space(__mmtk_objref) {
-                return <#f_ty as PolicyTraceObject #ty_generics>::trace_object::<Q, KIND>(&self.#f_ident, __mmtk_queue, __mmtk_objref, #copy, __mmtk_worker);
+            {
+                let space = self.#f_ident.read();
+                if space.in_space(__mmtk_objref) {
+                    return PolicyTraceObject::trace_object::<Q, KIND>(&*space, __mmtk_queue, __mmtk_objref, #copy, __mmtk_worker);
+                }
             }
         }
     });
@@ -108,13 +110,15 @@ pub(crate) fn generate_post_scan_object<'a>(
 ) -> TokenStream2 {
     let scan_field_handler = post_scan_object_fields.iter().map(|f| {
         let f_ident = f.ident.as_ref().unwrap();
-        let f_ty = &f.ty;
 
         quote! {
-            if self.#f_ident.in_space(__mmtk_objref) {
-                use crate::policy::gc_work::PolicyTraceObject;
-                <#f_ty as PolicyTraceObject #ty_generics>::post_scan_object(&self.#f_ident, __mmtk_objref);
-                return;
+            {
+                let space = self.#f_ident.read();
+                if space.in_space(__mmtk_objref) {
+                    use crate::policy::gc_work::PolicyTraceObject;
+                    PolicyTraceObject::post_scan_object(&*space, __mmtk_objref);
+                    return;
+                }
             }
         }
     });
@@ -148,10 +152,23 @@ pub(crate) fn generate_may_move_objects<'a>(
 ) -> TokenStream2 {
     // If any space or the parent may move objects, the plan may move objects
     let space_handlers = space_fields.iter().map(|f| {
-        let f_ty = &f.ty;
-
-        quote! {
-            || <#f_ty as PolicyTraceObject #ty_generics>::may_move_objects::<KIND>()
+        use syn::{Type, PathArguments};
+        // We assume the space field is `ArcFlexMut<T>`
+        if let Type::Path(type_path) = &f.ty {
+            if type_path.path.segments[0].ident == "ArcFlexMut" {
+                if let PathArguments::AngleBracketed(angle_bracketed_args) = &type_path.path.segments[0].arguments {
+                    let inner_type = &angle_bracketed_args.args.first().unwrap();
+                    quote! {
+                        || <#inner_type as PolicyTraceObject #ty_generics>::may_move_objects::<KIND>()
+                    }
+                } else {
+                    unreachable!("Failed to get the inner type of ArcFlexMut: {:?}", f.ty)
+                }
+            } else {
+                panic!("Expected a space to be ArcFlexMut<T>, found {:?}", f.ty)
+            }
+        } else {
+            panic!("Failed to get the type of a space: {:?}", f.ty)
         }
     });
 

src/memory_manager.rs

@@ -92,9 +92,10 @@ pub fn mmtk_init<VM: VMBinding>(builder: &MMTKBuilder) -> Box<MMTK<VM>> {
 /// Currently we do not allow removing regions from VM space.
 #[cfg(feature = "vm_space")]
 pub fn set_vm_space<VM: VMBinding>(mmtk: &'static mut MMTK<VM>, start: Address, size: usize) {
-    unsafe { mmtk.get_plan_mut() }
-        .base_mut()
+    mmtk.get_plan()
+        .base()
         .vm_space
+        .write()
         .set_vm_region(start, size);
 }
 

src/mmtk.rs

@@ -23,7 +23,6 @@ use crate::util::sanity::sanity_checker::SanityChecker;
 use crate::util::statistics::stats::Stats;
 use crate::vm::ReferenceGlue;
 use crate::vm::VMBinding;
-use std::cell::UnsafeCell;
 use std::default::Default;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -107,7 +106,7 @@ impl Default for MMTKBuilder {
 pub struct MMTK<VM: VMBinding> {
     pub(crate) options: Arc<Options>,
     pub(crate) state: Arc<GlobalState>,
-    pub(crate) plan: UnsafeCell<Box<dyn Plan<VM = VM>>>,
+    pub(crate) plan: Box<dyn Plan<VM = VM>>,
     pub(crate) reference_processors: ReferenceProcessors,
     pub(crate) finalizable_processor:
         Mutex<FinalizableProcessor<<VM::VMReferenceGlue as ReferenceGlue<VM>>::FinalizableType>>,
@@ -198,7 +197,7 @@ impl<VM: VMBinding> MMTK<VM> {
         MMTK {
             options,
             state,
-            plan: UnsafeCell::new(plan),
+            plan,
             reference_processors: ReferenceProcessors::new(),
             finalizable_processor: Mutex::new(FinalizableProcessor::<
                 <VM::VMReferenceGlue as ReferenceGlue<VM>>::FinalizableType,
@@ -339,17 +338,7 @@ impl<VM: VMBinding> MMTK<VM> {
     }
 
     pub fn get_plan(&self) -> &dyn Plan<VM = VM> {
-        unsafe { &**(self.plan.get()) }
-    }
-
-    /// Get the plan as mutable reference.
-    ///
-    /// # Safety
-    ///
-    /// This is unsafe because the caller must ensure that the plan is not used by other threads.
-    #[allow(clippy::mut_from_ref)]
-    pub unsafe fn get_plan_mut(&self) -> &mut dyn Plan<VM = VM> {
-        &mut **(self.plan.get())
+        &*self.plan
     }
 
     pub fn get_options(&self) -> &Options {