Refactor: Use Atomic<Address> where appropriate (#843)

Use `Atomic<Address>` instead of `AtomicUsize` where appropriate.

Closes #121

---------

Signed-off-by: ClSlaid <[email protected]>
Signed-off-by: 蔡略 <[email protected]>

Author: ClSlaid

src/policy/lockfreeimmortalspace.rs

@@ -1,19 +1,22 @@
+use atomic::Atomic;
+
+use std::marker::PhantomData;
+use std::sync::atomic::Ordering;
+
 use crate::mmtk::SFT_MAP;
+use crate::policy::sft::GCWorkerMutRef;
 use crate::policy::sft::SFT;
 use crate::policy::space::{CommonSpace, Space};
 use crate::util::address::Address;
-use crate::util::heap::PageResource;
-use crate::util::ObjectReference;
 
-use crate::policy::sft::GCWorkerMutRef;
 use crate::util::conversions;
 use crate::util::heap::layout::vm_layout_constants::{AVAILABLE_BYTES, AVAILABLE_START};
+use crate::util::heap::PageResource;
 use crate::util::metadata::side_metadata::SideMetadataContext;
 use crate::util::metadata::side_metadata::SideMetadataSanity;
 use crate::util::opaque_pointer::*;
+use crate::util::ObjectReference;
 use crate::vm::VMBinding;
-use std::marker::PhantomData;
-use std::sync::atomic::{AtomicUsize, Ordering};
 
 /// This type implements a lock free version of the immortal collection
 /// policy. This is close to the OpenJDK's epsilon GC.
@@ -25,10 +28,7 @@ pub struct LockFreeImmortalSpace<VM: VMBinding> {
     #[allow(unused)]
     name: &'static str,
     /// Heap range start
-    ///
-    /// We use `AtomicUsize` instead of `Address` here to atomically bumping this cursor.
-    /// TODO: Better address type here (Atomic<Address>?)
-    cursor: AtomicUsize,
+    cursor: Atomic<Address>,
     /// Heap range end
     limit: Address,
     /// start of this space
@@ -108,15 +108,20 @@ impl<VM: VMBinding> Space<VM> for LockFreeImmortalSpace<VM> {
     }
 
     fn reserved_pages(&self) -> usize {
-        let cursor = unsafe { Address::from_usize(self.cursor.load(Ordering::Relaxed)) };
+        let cursor = self.cursor.load(Ordering::Relaxed);
         let data_pages = conversions::bytes_to_pages_up(self.limit - cursor);
         let meta_pages = self.metadata.calculate_reserved_pages(data_pages);
         data_pages + meta_pages
     }
 
     fn acquire(&self, _tls: VMThread, pages: usize) -> Address {
         let bytes = conversions::pages_to_bytes(pages);
-        let start = unsafe { Address::from_usize(self.cursor.fetch_add(bytes, Ordering::Relaxed)) };
+        let start = self
+            .cursor
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |addr| {
+                Some(addr.add(bytes))
+            })
+            .expect("update cursor failed");
         if start + bytes > self.limit {
             panic!("OutOfMemory")
         }
@@ -182,7 +187,7 @@ impl<VM: VMBinding> LockFreeImmortalSpace<VM> {
         // https://github.com/mmtk/mmtk-core/issues/314
         let space = Self {
             name: args.name,
-            cursor: AtomicUsize::new(AVAILABLE_START.as_usize()),
+            cursor: Atomic::new(AVAILABLE_START),
             limit: AVAILABLE_START + total_bytes,
             start: AVAILABLE_START,
             extent: total_bytes,

src/policy/marksweepspace/malloc_ms/global.rs

@@ -1,3 +1,5 @@
+use atomic::Atomic;
+
 use super::metadata::*;
 use crate::plan::ObjectQueue;
 use crate::plan::VectorObjectQueue;
@@ -39,8 +41,8 @@ pub struct MallocSpace<VM: VMBinding> {
     phantom: PhantomData<VM>,
     active_bytes: AtomicUsize,
     active_pages: AtomicUsize,
-    pub chunk_addr_min: AtomicUsize, // XXX: have to use AtomicUsize to represent an Address
-    pub chunk_addr_max: AtomicUsize,
+    pub chunk_addr_min: Atomic<Address>,
+    pub chunk_addr_max: Atomic<Address>,
     metadata: SideMetadataContext,
     /// Work packet scheduler
     scheduler: Arc<GCWorkScheduler<VM>>,
@@ -255,8 +257,8 @@ impl<VM: VMBinding> MallocSpace<VM> {
             phantom: PhantomData,
             active_bytes: AtomicUsize::new(0),
             active_pages: AtomicUsize::new(0),
-            chunk_addr_min: AtomicUsize::new(usize::max_value()), // XXX: have to use AtomicUsize to represent an Address
-            chunk_addr_max: AtomicUsize::new(0),
+            chunk_addr_min: Atomic::new(Address::MAX),
+            chunk_addr_max: Atomic::new(Address::ZERO),
             metadata: SideMetadataContext {
                 global: args.global_side_metadata_specs.clone(),
                 local: metadata::extract_side_metadata(&[
@@ -428,12 +430,11 @@ impl<VM: VMBinding> MallocSpace<VM> {
         // Update chunk_addr_min, basing on the start of the allocation: addr.
         {
             let min_chunk_start = conversions::chunk_align_down(addr);
-            let min_chunk_usize = min_chunk_start.as_usize();
             let mut min = self.chunk_addr_min.load(Ordering::Relaxed);
-            while min_chunk_usize < min {
+            while min_chunk_start < min {
                 match self.chunk_addr_min.compare_exchange_weak(
                     min,
-                    min_chunk_usize,
+                    min_chunk_start,
                     Ordering::AcqRel,
                     Ordering::Relaxed,
                 ) {
@@ -446,12 +447,11 @@ impl<VM: VMBinding> MallocSpace<VM> {
         // Update chunk_addr_max, basing on the end of the allocation: addr + size.
         {
             let max_chunk_start = conversions::chunk_align_down(addr + size);
-            let max_chunk_usize = max_chunk_start.as_usize();
             let mut max = self.chunk_addr_max.load(Ordering::Relaxed);
-            while max_chunk_usize > max {
+            while max_chunk_start > max {
                 match self.chunk_addr_max.compare_exchange_weak(
                     max,
-                    max_chunk_usize,
+                    max_chunk_start,
                     Ordering::AcqRel,
                     Ordering::Relaxed,
                 ) {
@@ -467,9 +467,8 @@ impl<VM: VMBinding> MallocSpace<VM> {
     pub fn release(&mut self) {
         use crate::scheduler::WorkBucketStage;
         let mut work_packets: Vec<Box<dyn GCWork<VM>>> = vec![];
-        let mut chunk = unsafe { Address::from_usize(self.chunk_addr_min.load(Ordering::Relaxed)) }; // XXX: have to use AtomicUsize to represent an Address
-        let end = unsafe { Address::from_usize(self.chunk_addr_max.load(Ordering::Relaxed)) }
-            + BYTES_IN_CHUNK;
+        let mut chunk = self.chunk_addr_min.load(Ordering::Relaxed);
+        let end = self.chunk_addr_max.load(Ordering::Relaxed) + BYTES_IN_CHUNK;
 
         // Since only a single thread generates the sweep work packets as well as it is a Stop-the-World collector,
         // we can assume that the chunk mark metadata is not being accessed by anything else and hence we use

src/util/address.rs

@@ -1,4 +1,5 @@
 use atomic_traits::Atomic;
+
 use std::fmt;
 use std::mem;
 use std::ops::*;