Fix mutable aliasing UB: Replace Box<Entry> with NonNull<Entry>

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<Option<Box<Entry>>>]>,
+    buckets: Box<[Mutex<Option<NonNull<Entry>>>]>,
 }
 
 pub(crate) struct Entry {
     pub(crate) string: Box<str>,
     pub(crate) hash: u32,
     pub(crate) ref_count: AtomicIsize,
-    next_in_bucket: Option<Box<Entry>>,
+    next_in_bucket: Option<NonNull<Entry>>,
 }
 
+// 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<Option<Box<Entry>>>; NB_BUCKETS],
-    // const MUTEX: Mutex<Option<Box<Entry>>> = Mutex::new(None);
-    // let buckets = Box::new([MUTEX; NB_BUCKETS]);
-    // ```
     static DYNAMIC_SET: OnceLock<Set> = 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<Entry>> = linked_list.as_mut();
+            let mut ptr: Option<NonNull<Entry>> = *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>() >= 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<Box<Entry>> = &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<NonNull<Entry>> = &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();
+            }
         }
     }
 }