feat: Add Slab.get_mut

This commit adds a Slab::get_mut method, which takes a `&mut self`, and
thus is able to get a mutable reference to an element without actually
using any synchronization structures. This is a similar concept to
`UnsafeCell::get_mut`, and may be useful for cases such as:

- A sharded `Slab` wrapped in a `RwLock`, such that hot path operations
  are done through read guards, and rarer and complex operations are
  done through write guards.
- Any other situation where code has both parallel and serial section,
  such as in fork-join systems.
- Gradually converting single-threaded systems to parallelizable systems.

The doc comment adds a test that helps ensure basic functionality works.

`test_println!` comments are not added since, as this doesn't actually
manipulate any internal structures except the user-provided types, they
don't seem that useful in this case.

src/lib.rs

@@ -623,6 +623,34 @@ impl<T, C: cfg::Config> Slab<T, C> {
         })
     }
 
+    /// Return a mutable reference to the value associated with the given key.
+    ///
+    /// This call borrows `self` mutably (at compile-time) which guarantees that we
+    /// possess the only reference.
+    ///
+    /// If the slab does not contain a value for the given key, `None` is returned
+    /// instead.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// let mut slab = sharded_slab::Slab::new();
+    /// let key = slab.insert(String::from("hello world")).unwrap();
+    /// slab.get_mut(key).unwrap().push('!');
+    ///
+    /// assert_eq!(*slab.get_mut(key).unwrap(), "hello world!");
+    /// assert!(slab.get_mut(12345).is_none());
+    /// ```
+    #[cfg(not(loom))]
+    pub fn get_mut(&mut self, key: usize) -> Option<&mut T> {
+        let tid = C::unpack_tid(key);
+
+        let shard = self.shards.get_mut(tid.as_usize())?;
+        shard
+            .slot_mut(key)
+            .map(|slot| slot.value_mut().as_mut().unwrap())
+    }
+
     /// Return an owned reference to the value at the given index.
     ///
     /// If the slab does not contain a value for the given key, `None` is

src/page/mod.rs

@@ -190,6 +190,13 @@ where
         })
     }
 
+    #[inline]
+    #[cfg(not(loom))]
+    pub(crate) fn slot_mut(&mut self, addr: Addr<C>) -> Option<&mut Slot<T, C>> {
+        let poff = addr.offset() - self.prev_sz;
+        self.slab.get_mut().as_mut()?.get_mut(poff)
+    }
+
     #[inline(always)]
     pub(crate) fn free_list(&self) -> &impl FreeList<C> {
         &self.remote

src/page/slot.rs

@@ -96,6 +96,12 @@ where
         self.item.with(|item| unsafe { &*item })
     }
 
+    #[inline(always)]
+    #[cfg(not(loom))]
+    pub(crate) fn value_mut(&mut self) -> &mut T {
+        self.item.get_mut()
+    }
+
     #[inline(always)]
     pub(super) fn set_next(&self, next: usize) {
         self.next.with_mut(|n| unsafe {

src/shard.rs

@@ -13,7 +13,11 @@ use crate::{
     Pack,
 };
 
-use std::{fmt, ptr, slice};
+use std::{
+    fmt,
+    ptr::{self, NonNull},
+    slice,
+};
 
 // ┌─────────────┐      ┌────────┐
 // │ page 1      │      │        │
@@ -84,6 +88,15 @@ where
         self.shared[page_index].with_slot(addr, f)
     }
 
+    #[inline(always)]
+    #[cfg(not(loom))]
+    pub(crate) fn slot_mut(&mut self, idx: usize) -> Option<&mut page::Slot<T, C>> {
+        let (addr, page_index) = page::indices::<C>(idx);
+        self.shared
+            .get_mut(page_index)
+            .and_then(|page| page.slot_mut(addr))
+    }
+
     pub(crate) fn new(tid: usize) -> Self {
         let mut total_sz = 0;
         let shared = (0..C::MAX_PAGES)
@@ -287,6 +300,12 @@ where
         self.shards.get(idx)?.load(Acquire)
     }
 
+    #[inline]
+    #[cfg(not(loom))]
+    pub(crate) fn get_mut(&mut self, idx: usize) -> Option<&mut Shard<T, C>> {
+        self.shards.get_mut(idx)?.get_mut()
+    }
+
     #[inline]
     pub(crate) fn current(&self) -> (Tid<C>, &Shard<T, C>) {
         let tid = Tid::<C>::current();
@@ -398,6 +417,13 @@ impl<T, C: cfg::Config> Ptr<T, C> {
         Some(track.get_ref())
     }
 
+    #[inline]
+    #[cfg(not(loom))]
+    fn get_mut(&mut self) -> Option<&mut Shard<T, C>> {
+        // MSRV: We could do `track.is_mut()`, but the current MSRV doesn't have NLL
+        NonNull::new(*self.0.get_mut()).map(|track| unsafe { &mut *track.as_ptr() }.get_mut())
+    }
+
     #[inline]
     fn set(&self, new: *mut alloc::Track<Shard<T, C>>) {
         self.0

src/sync.rs

@@ -102,6 +102,14 @@ mod inner {
         {
             f(self.0.get())
         }
+
+        #[inline(always)]
+        #[cfg(not(loom))]
+        pub fn get_mut(&mut self) -> &mut T {
+            // Safety: same as `std::UnsafeCell::get_mut`
+            // MSRV: `std::UnsafeCell::get_mut` stabilized in 1.50.0
+            unsafe { &mut *self.0.get() }
+        }
     }
 
     pub(crate) mod alloc {