feat: use bucket mutex instead of global mutex for dynamic set

src/atom.rs

@@ -200,8 +200,7 @@ impl<'a, Static: StaticAtomSet> From<Cow<'a, str>> for Atom<Static> {
                     phantom: PhantomData,
                 }
             } else {
-                let ptr: std::ptr::NonNull<Entry> =
-                    DYNAMIC_SET.lock().insert(string_to_add, hash.g);
+                let ptr: std::ptr::NonNull<Entry> = DYNAMIC_SET.insert(string_to_add, hash.g);
                 let data = ptr.as_ptr() as u64;
                 debug_assert!(0 == data & TAG_MASK);
                 Atom {
@@ -237,9 +236,7 @@ impl<Static> Drop for Atom<Static> {
 
         // Out of line to guide inlining.
         fn drop_slow<Static>(this: &mut Atom<Static>) {
-            DYNAMIC_SET
-                .lock()
-                .remove(this.unsafe_data.get() as *mut Entry);
+            DYNAMIC_SET.remove(this.unsafe_data.get() as *mut Entry);
         }
     }
 }

src/dynamic_set.rs

@@ -19,7 +19,7 @@ const NB_BUCKETS: usize = 1 << 12; // 4096
 const BUCKET_MASK: u32 = (1 << 12) - 1;
 
 pub(crate) struct Set {
-    buckets: Box<[Option<Box<Entry>>; NB_BUCKETS]>,
+    buckets: Box<[Mutex<Option<Box<Entry>>>]>,
 }
 
 pub(crate) struct Entry {
@@ -38,22 +38,24 @@ fn entry_alignment_is_sufficient() {
     assert!(mem::align_of::<Entry>() >= ENTRY_ALIGNMENT);
 }
 
-pub(crate) static DYNAMIC_SET: Lazy<Mutex<Set>> = Lazy::new(|| {
-    Mutex::new({
-        type T = Option<Box<Entry>>;
-        let _static_assert_size_eq = std::mem::transmute::<T, usize>;
-        let vec = std::mem::ManuallyDrop::new(vec![0_usize; NB_BUCKETS]);
-        Set {
-            buckets: unsafe { Box::from_raw(vec.as_ptr() as *mut [T; NB_BUCKETS]) },
-        }
-    })
+pub(crate) static DYNAMIC_SET: Lazy<Set> = Lazy::new(|| {
+    // 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]);
+    // ```
+    let buckets = (0..NB_BUCKETS).map(|_| Mutex::new(None)).collect();
+    Set { buckets }
 });
 
 impl Set {
-    pub(crate) fn insert(&mut self, string: Cow<str>, hash: u32) -> NonNull<Entry> {
+    pub(crate) fn insert(&self, string: Cow<str>, hash: u32) -> NonNull<Entry> {
         let bucket_index = (hash & BUCKET_MASK) as usize;
+        let mut linked_list = self.buckets[bucket_index].lock();
+
         {
-            let mut ptr: Option<&mut Box<Entry>> = self.buckets[bucket_index].as_mut();
+            let mut ptr: Option<&mut Box<Entry>> = linked_list.as_mut();
 
             while let Some(entry) = ptr.take() {
                 if entry.hash == hash && *entry.string == *string {
@@ -74,25 +76,25 @@ impl Set {
         debug_assert!(mem::align_of::<Entry>() >= ENTRY_ALIGNMENT);
         let string = string.into_owned();
         let mut entry = Box::new(Entry {
-            next_in_bucket: self.buckets[bucket_index].take(),
+            next_in_bucket: linked_list.take(),
             hash,
             ref_count: AtomicIsize::new(1),
             string: string.into_boxed_str(),
         });
         let ptr = NonNull::from(&mut *entry);
-        self.buckets[bucket_index] = Some(entry);
-
+        *linked_list = Some(entry);
         ptr
     }
 
-    pub(crate) fn remove(&mut self, ptr: *mut Entry) {
+    pub(crate) fn remove(&self, ptr: *mut Entry) {
         let bucket_index = {
             let value: &Entry = unsafe { &*ptr };
             debug_assert!(value.ref_count.load(SeqCst) == 0);
             (value.hash & BUCKET_MASK) as usize
         };
 
-        let mut current: &mut Option<Box<Entry>> = &mut self.buckets[bucket_index];
+        let mut linked_list = self.buckets[bucket_index].lock();
+        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;