fixes

Author: jkarneges

src/core/arena.rs

@@ -16,44 +16,26 @@
  */
 
 use slab::Slab;
-use std::cell::{Cell, RefCell, RefMut};
+use std::cell::{Cell, RefCell};
+use std::fmt;
 use std::marker::PhantomData;
 use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::process::abort;
 use std::ptr::NonNull;
 use std::sync::{Mutex, MutexGuard};
 
-pub struct EntryGuard<'a, T> {
-    entries: RefMut<'a, Slab<T>>,
-    entry: &'a mut T,
-    key: usize,
-}
-
-impl<T> EntryGuard<'_, T> {
-    fn remove(mut self) {
-        self.entries.remove(self.key);
-    }
-}
-
-impl<T> Deref for EntryGuard<'_, T> {
-    type Target = T;
+pub struct InsertError<T>(pub T);
 
-    fn deref(&self) -> &Self::Target {
-        self.entry
-    }
-}
-
-impl<T> DerefMut for EntryGuard<'_, T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.entry
+impl<T> fmt::Debug for InsertError<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("InsertError").finish_non_exhaustive()
     }
 }
 
 // This is essentially a sharable slab for use within a single thread.
-// operations are protected by a RefCell. When an element is retrieved for
-// reading or modification, it is wrapped in a EntryGuard which keeps the
-// entire slab borrowed until the caller is done working with the element
+// Operations are protected by a RefCell, however lookup operations return
+// pointers that can be used without RefCell protection.
 pub struct Memory<T> {
     entries: RefCell<Slab<T>>,
 }
@@ -70,58 +52,44 @@ impl<T> Memory<T> {
 
     #[cfg(test)]
     pub fn len(&self) -> usize {
-        let entries = self.entries.borrow();
-
-        entries.len()
+        self.entries.borrow().len()
     }
 
-    fn insert(&self, e: T) -> Result<usize, ()> {
+    fn insert(&self, e: T) -> Result<usize, InsertError<T>> {
         let mut entries = self.entries.borrow_mut();
 
         // Out of capacity. By preventing inserts beyond the capacity, we
-        // ensure the underlying memory won't get moved due to a realloc
+        // ensure the underlying memory won't get moved due to a realloc.
         if entries.len() == entries.capacity() {
-            return Err(());
+            return Err(InsertError(e));
         }
 
         Ok(entries.insert(e))
     }
 
-    fn get<'a>(&'a self, key: usize) -> Option<EntryGuard<'a, T>> {
-        let mut entries = self.entries.borrow_mut();
+    fn remove(&self, key: usize) {
+        self.entries.borrow_mut().remove(key);
+    }
 
-        let entry = entries.get_mut(key)?;
+    // Returns a pointer to an entry if it exists.
+    //
+    // SAFETY: The returned pointer is guaranteed to be valid until the entry
+    // is removed or the Memory is dropped.
+    fn addr_of(&self, key: usize) -> Option<*const T> {
+        let entries = self.entries.borrow();
+
+        let entry = entries.get(key)?;
 
         // Slab element addresses are guaranteed to be stable once created,
-        // and the only place we remove the element is in EntryGuard's
-        // remove method which consumes itself, therefore it is safe to
-        // assume the element will live at least as long as the EntryGuard
-        // and we can extend the lifetime of the reference beyond the
-        // RefMut
-        let entry = unsafe { mem::transmute::<&mut T, &'a mut T>(entry) };
+        // therefore we can return a pointer to the element and guarantee
+        // its validity until the element is removed.
 
-        Some(EntryGuard {
-            entries,
-            entry,
-            key,
-        })
+        Some(entry as *const T)
     }
 
     fn key_of(&self, present_element: &T) -> usize {
         self.entries.borrow().key_of(present_element)
     }
-
-    // For tests, as a way to confirm the memory isn't moving. Be careful
-    // with this. The very first element inserted will be at index 0, but
-    // if the slab has been used and cleared, then the next element
-    // inserted may not be at index 0 and calling this method afterward
-    // will panic
-    #[cfg(test)]
-    fn entry0_ptr(&self) -> *const T {
-        let entries = self.entries.borrow();
-
-        entries.get(0).unwrap() as *const T
-    }
 }
 
 pub struct SyncEntryGuard<'a, T> {
@@ -370,18 +338,21 @@ pub struct Rc<T> {
 
 impl<T> Rc<T> {
     #[allow(clippy::result_unit_err)]
-    pub fn new(v: T, memory: &std::rc::Rc<RcMemory<T>>) -> Result<Self, ()> {
-        let key = memory.insert(RcInner {
-            refs: Cell::new(1),
-            memory: std::rc::Rc::clone(memory),
-            value: v,
-        })?;
+    pub fn new(v: T, memory: &std::rc::Rc<RcMemory<T>>) -> Result<Self, InsertError<T>> {
+        let key = memory
+            .insert(RcInner {
+                refs: Cell::new(1),
+                memory: std::rc::Rc::clone(memory),
+                value: v,
+            })
+            .map_err(|e| InsertError(e.0.value))?;
 
         // Guaranteed to succeed since we inserted the element above
-        let mut inner = memory.get(key).unwrap();
+        let ptr: *const RcInner<T> = memory.addr_of(key).unwrap();
 
-        // Get a pointer to the inner data
-        let ptr = (&mut *inner).into();
+        // SAFETY: ptr is not null and we promise to only use it immutably
+        // despite casting it to *mut in order to construct NonNull
+        let ptr = unsafe { NonNull::new_unchecked(ptr as *mut RcInner<T>) };
 
         Ok(Self {
             ptr,
@@ -417,10 +388,21 @@ impl<T> Rc<T> {
 
     #[inline(never)]
     fn drop_slow(&mut self) {
-        let inner = self.inner();
-        let key = inner.memory.key_of(inner);
-        let entry = inner.memory.get(key).unwrap();
-        entry.remove();
+        // Entries contain a std::rc::Rc to the Memory they are contained in,
+        // and we need to be careful the Memory is not dropped while an entry
+        // is being removed. To ensure this, we clone the Rc to the Memory
+        // before removing the entry.
+
+        let memory = {
+            let inner = self.inner();
+            let memory = std::rc::Rc::clone(&inner.memory);
+            memory.remove(memory.key_of(inner));
+
+            memory
+        };
+
+        // The entry has been removed by this point
+        drop(memory);
     }
 }
 
@@ -626,15 +608,15 @@ mod tests {
 
         let e0a = Rc::new(123 as i32, &memory).unwrap();
         assert_eq!(memory.len(), 1);
-        let p = memory.entry0_ptr();
+        let p = memory.addr_of(0).unwrap();
 
         let e0b = Rc::clone(&e0a);
         assert_eq!(memory.len(), 1);
-        assert_eq!(memory.entry0_ptr(), p);
+        assert_eq!(memory.addr_of(0).unwrap(), p);
 
         let e1a = Rc::new(456 as i32, &memory).unwrap();
         assert_eq!(memory.len(), 2);
-        assert_eq!(memory.entry0_ptr(), p);
+        assert_eq!(memory.addr_of(0).unwrap(), p);
 
         // No room
         assert!(Rc::new(789 as i32, &memory).is_err());
@@ -645,7 +627,7 @@ mod tests {
 
         mem::drop(e0b);
         assert_eq!(memory.len(), 2);
-        assert_eq!(memory.entry0_ptr(), p);
+        assert_eq!(memory.addr_of(0).unwrap(), p);
 
         mem::drop(e0a);
         assert_eq!(memory.len(), 1);