diff --git a/src/dynamic_set.rs b/src/dynamic_set.rs index 4442b4d..7767550 100644 --- a/src/dynamic_set.rs +++ b/src/dynamic_set.rs @@ -19,16 +19,25 @@ const NB_BUCKETS: usize = 1 << 12; // 4096 const BUCKET_MASK: u32 = (1 << 12) - 1; pub(crate) struct Set { - buckets: Box<[Mutex>>]>, + buckets: Box<[Mutex>>]>, } pub(crate) struct Entry { pub(crate) string: Box, pub(crate) hash: u32, pub(crate) ref_count: AtomicIsize, - next_in_bucket: Option>, + next_in_bucket: Option>, } +// SAFETY: Access to the global linked list is strictly guarded by a Mutex, +// and the reference counts are atomic. Even though `NonNull` is strictly +// `!Send` and `!Sync`, the surrounding architecture makes it safe to share. +unsafe impl Send for Entry {} +unsafe impl Sync for Entry {} + +unsafe impl Send for Set {} +unsafe impl Sync for Set {} + // Addresses are a multiples of this, // and therefore have have TAG_MASK bits unset, available for tagging. pub(crate) const ENTRY_ALIGNMENT: usize = 4; @@ -40,11 +49,6 @@ fn entry_alignment_is_sufficient() { pub(crate) fn dynamic_set() -> &'static Set { // NOTE: Using const initialization for buckets breaks the small-stack test. - // ``` - // // buckets: [Mutex>>; NB_BUCKETS], - // const MUTEX: Mutex>> = Mutex::new(None); - // let buckets = Box::new([MUTEX; NB_BUCKETS]); - // ``` static DYNAMIC_SET: OnceLock = OnceLock::new(); DYNAMIC_SET.get_or_init(|| { @@ -59,12 +63,19 @@ impl Set { let mut linked_list = self.buckets[bucket_index].lock(); { - let mut ptr: Option<&mut Box> = linked_list.as_mut(); + let mut ptr: Option> = *linked_list; - while let Some(entry) = ptr.take() { + while let Some(entry_ptr) = ptr { + // SAFETY: We hold the Mutex lock for this bucket, so no other thread can mutate + // the linked list. The `NonNull` pointer is guaranteed to point to a valid Entry. + let entry = unsafe { entry_ptr.as_ref() }; if entry.hash == hash && *entry.string == *string { - if entry.ref_count.fetch_add(1, SeqCst) > 0 { - return NonNull::from(&mut **entry); + let old_size = entry.ref_count.fetch_add(1, SeqCst); + if old_size > 0 { + if old_size == std::isize::MAX { + std::process::abort(); + } + return entry_ptr; } // Uh-oh. The pointer's reference count was zero, which means someone may try // to free it. (Naive attempts to defend against this, for example having the @@ -74,39 +85,56 @@ impl Set { entry.ref_count.fetch_sub(1, SeqCst); break; } - ptr = entry.next_in_bucket.as_mut(); + ptr = entry.next_in_bucket; } } debug_assert!(mem::align_of::() >= ENTRY_ALIGNMENT); let string = string.into_owned(); - let mut entry = Box::new(Entry { + let entry = Box::new(Entry { next_in_bucket: linked_list.take(), hash, ref_count: AtomicIsize::new(1), string: string.into_boxed_str(), }); - let ptr = NonNull::from(&mut *entry); - *linked_list = Some(entry); + let ptr = NonNull::new(Box::into_raw(entry)).unwrap(); + *linked_list = Some(ptr); ptr } pub(crate) fn remove(&self, ptr: *mut Entry) { + // SAFETY: The caller provides a pointer derived from a valid Atom. We hold the lock + // below, and `ptr` is guaranteed to be valid until we drop the `Box` later in this function. let value: &Entry = unsafe { &*ptr }; let bucket_index = (value.hash & BUCKET_MASK) as usize; let mut linked_list = self.buckets[bucket_index].lock(); debug_assert!(value.ref_count.load(SeqCst) == 0); - let mut current: &mut Option> = &mut linked_list; - - while let Some(entry_ptr) = current.as_mut() { - let entry_ptr: *mut Entry = &mut **entry_ptr; - if entry_ptr == ptr { - mem::drop(mem::replace(current, unsafe { - (*entry_ptr).next_in_bucket.take() - })); + let mut current: &mut Option> = &mut linked_list; + + while let Some(entry_ptr) = *current { + if entry_ptr.as_ptr() == ptr { + // SAFETY: The reference count has reached 0, and we hold the bucket lock. + // We have exclusive access to recreate the Box and deallocate the memory. + let mut unlinked_entry = unsafe { Box::from_raw(entry_ptr.as_ptr()) }; + *current = unlinked_entry.next_in_bucket.take(); break; } - current = unsafe { &mut (*entry_ptr).next_in_bucket }; + // SAFETY: We hold the bucket lock, so the pointer remains valid and unaliased here. + current = unsafe { &mut (*entry_ptr.as_ptr()).next_in_bucket }; + } + } +} + +impl Drop for Set { + fn drop(&mut self) { + for bucket in self.buckets.iter_mut() { + let mut current = bucket.get_mut().take(); + while let Some(ptr) = current { + // SAFETY: The Set is being dropped, meaning no other Atom references can exist. + // We own the pointers and must reconstruct the Box to safely free the memory. + let mut entry = unsafe { Box::from_raw(ptr.as_ptr()) }; + current = entry.next_in_bucket.take(); + } } } }