Fix cloned rsx double drop (#2839)

* Implement Rc generational pointers

* Clean up error handling a bit

* start adding tests

* Simplify a few types

* Separate Rc and normal generational boxes

* Add tests for rc drop and read

* Move over the pointer instead of the data in ReadOnlySignal props

* Fix panic on ref drop and release mode builds

* final rc signal cleanup

* use rc generational boxes for event handlers as well

* add a regression test for read only signal cloning

* merge main

---------

Co-authored-by: Jonathan Kelley <jkelleyrtp@gmail.com>

packages/generational-box/src/entry.rs

@@ -2,36 +2,66 @@ use crate::{
     BorrowError, BorrowMutError, GenerationalLocation, GenerationalRefBorrowGuard,
     GenerationalRefBorrowMutGuard,
 };
+use std::{
+    num::NonZeroU64,
+    sync::atomic::{AtomicU64, Ordering},
+};
 
-pub(crate) struct StorageEntry<T> {
-    generation: u64,
-
-    pub data: Option<T>,
+pub(crate) struct RcStorageEntry<T> {
+    ref_count: AtomicU64,
+    pub data: T,
 }
 
-impl<T> Default for StorageEntry<T> {
-    fn default() -> Self {
+impl<T> RcStorageEntry<T> {
+    pub const fn new(data: T) -> Self {
         Self {
-            generation: 0,
-            data: None,
+            ref_count: AtomicU64::new(0),
+            data,
         }
     }
+
+    pub fn add_ref(&self) {
+        self.ref_count.fetch_add(1, Ordering::SeqCst);
+    }
+
+    pub fn drop_ref(&self) -> bool {
+        let new_ref_count = self.ref_count.fetch_sub(1, Ordering::SeqCst);
+        new_ref_count == 0
+    }
+}
+
+pub(crate) struct StorageEntry<T> {
+    generation: NonZeroU64,
+    pub(crate) data: T,
 }
 
 impl<T> StorageEntry<T> {
+    pub const fn new(data: T) -> Self {
+        Self {
+            generation: NonZeroU64::MIN,
+            data,
+        }
+    }
+
     pub fn valid(&self, location: &GenerationalLocation) -> bool {
         self.generation == location.generation
     }
 
     pub fn increment_generation(&mut self) {
-        self.generation += 1;
+        self.generation = self.generation.checked_add(1).unwrap();
     }
 
-    pub fn generation(&self) -> u64 {
+    pub fn generation(&self) -> NonZeroU64 {
         self.generation
     }
 }
 
+impl<T: Default + 'static> Default for StorageEntry<T> {
+    fn default() -> Self {
+        Self::new(T::default())
+    }
+}
+
 #[derive(Default)]
 pub(crate) struct MemoryLocationBorrowInfo(
     #[cfg(any(debug_assertions, feature = "debug_borrows"))]

packages/generational-box/src/error.rs

@@ -5,7 +5,10 @@ use std::fmt::Display;
 use crate::GenerationalLocation;
 
 /// A result that can be returned from a borrow operation.
-pub type BorrowResult<T> = std::result::Result<T, BorrowError>;
+pub type BorrowResult<T = ()> = std::result::Result<T, BorrowError>;
+
+/// A result that can be returned from a borrow mut operation.
+pub type BorrowMutResult<T = ()> = std::result::Result<T, BorrowMutError>;
 
 #[derive(Debug, Clone, PartialEq)]
 /// An error that can occur when trying to borrow a value.
@@ -50,6 +53,15 @@ impl Display for BorrowMutError {
 
 impl Error for BorrowMutError {}
 
+impl From<BorrowError> for BorrowMutError {
+    fn from(error: BorrowError) -> Self {
+        match error {
+            BorrowError::Dropped(error) => BorrowMutError::Dropped(error),
+            BorrowError::AlreadyBorrowedMut(error) => BorrowMutError::AlreadyBorrowedMut(error),
+        }
+    }
+}
+
 /// An error that can occur when trying to use a value that has been dropped.
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub struct ValueDroppedError {

packages/generational-box/src/lib.rs

@@ -5,6 +5,7 @@ use parking_lot::Mutex;
 use std::{
     fmt::Debug,
     marker::PhantomData,
+    num::NonZeroU64,
     ops::{Deref, DerefMut},
     sync::Arc,
 };
@@ -24,7 +25,7 @@ mod unsync;
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct GenerationalBoxId {
     data_ptr: *const (),
-    generation: u64,
+    generation: NonZeroU64,
 }
 
 // Safety: GenerationalBoxId is Send and Sync because there is no way to access the pointer.
@@ -54,9 +55,19 @@ impl<T, S: Storage<T>> GenerationalBox<T, S> {
     /// Create a new generational box by leaking a value into the storage. This is useful for creating
     /// a box that needs to be manually dropped with no owners.
     #[track_caller]
-    pub fn leak(value: T) -> Self {
-        let location = S::claim(std::panic::Location::caller());
-        location.set(value);
+    pub fn leak(value: T, location: &'static std::panic::Location<'static>) -> Self {
+        let location = S::new(value, location);
+        Self {
+            raw: location,
+            _marker: PhantomData,
+        }
+    }
+
+    /// Create a new reference counted generational box by leaking a value into the storage. This is useful for creating
+    /// a box that needs to be manually dropped with no owners.
+    #[track_caller]
+    pub fn leak_rc(value: T, location: &'static std::panic::Location<'static>) -> Self {
+        let location = S::new_rc(value, location);
         Self {
             raw: location,
             _marker: PhantomData,
@@ -98,8 +109,12 @@ impl<T, S: Storage<T>> GenerationalBox<T, S> {
     }
 
     /// Set the value. Panics if the value is no longer valid.
-    pub fn set(&self, value: T) {
-        S::set(self.raw, value);
+    #[track_caller]
+    pub fn set(&self, value: T)
+    where
+        T: 'static,
+    {
+        *self.write() = value;
     }
 
     /// Returns true if the pointer is equal to the other pointer.
@@ -108,24 +123,30 @@ impl<T, S: Storage<T>> GenerationalBox<T, S> {
     }
 
     /// Drop the value out of the generational box and invalidate the generational box.
-    /// This will return the value if the value was taken.
-    pub fn manually_drop(&self) -> Option<T>
+    pub fn manually_drop(&self)
     where
         T: 'static,
     {
-        self.raw.take()
+        self.raw.recycle();
     }
 
     /// Try to get the location the generational box was created at. In release mode this will always return None.
     pub fn created_at(&self) -> Option<&'static std::panic::Location<'static>> {
-        #[cfg(debug_assertions)]
-        {
-            Some(self.raw.location.created_at)
-        }
-        #[cfg(not(debug_assertions))]
-        {
-            None
-        }
+        self.raw.location.created_at()
+    }
+
+    /// Get a reference to the value
+    #[track_caller]
+    pub fn leak_reference(&self) -> BorrowResult<GenerationalBox<T, S>> {
+        Ok(Self {
+            raw: S::new_reference(self.raw)?,
+            _marker: std::marker::PhantomData,
+        })
+    }
+
+    /// Change this box to point to another generational box
+    pub fn point_to(&mut self, other: GenerationalBox<T, S>) -> BorrowResult {
+        S::change_reference(self.raw, other.raw)
     }
 }
 
@@ -140,17 +161,35 @@ impl<T, S> Clone for GenerationalBox<T, S> {
 /// A trait for a storage backing type. (RefCell, RwLock, etc.)
 pub trait Storage<Data = ()>: AnyStorage + 'static {
     /// Try to read the value. Returns None if the value is no longer valid.
-    fn try_read(
-        location: GenerationalPointer<Self>,
-    ) -> Result<Self::Ref<'static, Data>, BorrowError>;
+    fn try_read(pointer: GenerationalPointer<Self>) -> BorrowResult<Self::Ref<'static, Data>>;
 
     /// Try to write the value. Returns None if the value is no longer valid.
-    fn try_write(
-        location: GenerationalPointer<Self>,
-    ) -> Result<Self::Mut<'static, Data>, BorrowMutError>;
+    fn try_write(pointer: GenerationalPointer<Self>) -> BorrowMutResult<Self::Mut<'static, Data>>;
 
-    /// Set the value if the location is valid
-    fn set(location: GenerationalPointer<Self>, value: Data);
+    /// Create a new memory location. This will either create a new memory location or recycle an old one.
+    fn new(
+        value: Data,
+        caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalPointer<Self>;
+
+    /// Create a new reference counted memory location. This will either create a new memory location or recycle an old one.
+    fn new_rc(
+        value: Data,
+        caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalPointer<Self>;
+
+    /// Reference another location if the location is valid
+    ///
+    /// This method may return an error if the other box is no longer valid or it is already borrowed mutably.
+    fn new_reference(inner: GenerationalPointer<Self>) -> BorrowResult<GenerationalPointer<Self>>;
+
+    /// Change the reference a signal is pointing to
+    ///
+    /// This method may return an error if the other box is no longer valid or it is already borrowed mutably.
+    fn change_reference(
+        pointer: GenerationalPointer<Self>,
+        rc_pointer: GenerationalPointer<Self>,
+    ) -> BorrowResult;
 }
 
 /// A trait for any storage backing type.
@@ -206,10 +245,7 @@ pub trait AnyStorage: Default + 'static {
     fn data_ptr(&self) -> *const ();
 
     /// Recycle a memory location. This will drop the memory location and return it to the runtime.
-    fn recycle(location: GenerationalPointer<Self>) -> Option<Box<dyn std::any::Any>>;
-
-    /// Claim a new memory location. This will either create a new memory location or recycle an old one.
-    fn claim(caller: &'static std::panic::Location<'static>) -> GenerationalPointer<Self>;
+    fn recycle(location: GenerationalPointer<Self>);
 
     /// Create a new owner. The owner will be responsible for dropping all of the generational boxes that it creates.
     fn owner() -> Owner<Self> {
@@ -222,11 +258,24 @@ pub trait AnyStorage: Default + 'static {
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct GenerationalLocation {
     /// The generation this location is associated with. Using the location after this generation is invalidated will return errors.
-    generation: u64,
+    generation: NonZeroU64,
     #[cfg(any(debug_assertions, feature = "debug_ownership"))]
     created_at: &'static std::panic::Location<'static>,
 }
 
+impl GenerationalLocation {
+    pub(crate) fn created_at(&self) -> Option<&'static std::panic::Location<'static>> {
+        #[cfg(debug_assertions)]
+        {
+            Some(self.created_at)
+        }
+        #[cfg(not(debug_assertions))]
+        {
+            None
+        }
+    }
+}
+
 /// A pointer to a specific generational box and generation in that box.
 pub struct GenerationalPointer<S: 'static = UnsyncStorage> {
     /// The storage that is backing this location
@@ -261,20 +310,6 @@ impl<S: 'static> Clone for GenerationalPointer<S> {
 impl<S: 'static> Copy for GenerationalPointer<S> {}
 
 impl<S> GenerationalPointer<S> {
-    fn take<T: 'static>(self) -> Option<T>
-    where
-        S: Storage<T>,
-    {
-        S::recycle(self).map(|value| *(value.downcast().unwrap()))
-    }
-
-    fn set<T>(self, value: T)
-    where
-        S: Storage<T>,
-    {
-        S::set(self, value)
-    }
-
     #[track_caller]
     fn try_read<T>(self) -> Result<S::Ref<'static, T>, BorrowError>
     where
@@ -346,6 +381,32 @@ impl<S: AnyStorage> Owner<S> {
         self.insert_with_caller(value, std::panic::Location::caller())
     }
 
+    /// Create a new reference counted box. The box will be dropped when all references are dropped.
+    #[track_caller]
+    pub fn insert_rc<T: 'static>(&self, value: T) -> GenerationalBox<T, S>
+    where
+        S: Storage<T>,
+    {
+        self.insert_rc_with_caller(value, std::panic::Location::caller())
+    }
+
+    /// Insert a value into the store with a specific location blamed for creating the value. The value will be dropped when the owner is dropped.
+    pub fn insert_rc_with_caller<T: 'static>(
+        &self,
+        value: T,
+        caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalBox<T, S>
+    where
+        S: Storage<T>,
+    {
+        let location = S::new_rc(value, caller);
+        self.0.lock().owned.push(location);
+        GenerationalBox {
+            raw: location,
+            _marker: std::marker::PhantomData,
+        }
+    }
+
     /// Insert a value into the store with a specific location blamed for creating the value. The value will be dropped when the owner is dropped.
     pub fn insert_with_caller<T: 'static>(
         &self,
@@ -355,8 +416,7 @@ impl<S: AnyStorage> Owner<S> {
     where
         S: Storage<T>,
     {
-        let location = S::claim(caller);
-        location.set(value);
+        let location = S::new(value, caller);
         self.0.lock().owned.push(location);
         GenerationalBox {
             raw: location,
@@ -364,14 +424,19 @@ impl<S: AnyStorage> Owner<S> {
         }
     }
 
-    /// Creates an invalid handle. This is useful for creating a handle that will be filled in later. If you use this before the value is filled in, you will get may get a panic or an out of date value.
+    /// Create a new reference to an existing box. The reference will be dropped when the owner is dropped.
+    ///
+    /// This method may return an error if the other box is no longer valid or it is already borrowed mutably.
     #[track_caller]
-    pub fn invalid<T: 'static>(&self) -> GenerationalBox<T, S> {
-        let location = S::claim(std::panic::Location::caller());
-        self.0.lock().owned.push(location);
-        GenerationalBox {
-            raw: location,
-            _marker: PhantomData,
-        }
+    pub fn insert_reference<T: 'static>(
+        &self,
+        other: GenerationalBox<T, S>,
+    ) -> BorrowResult<GenerationalBox<T, S>>
+    where
+        S: Storage<T>,
+    {
+        let location = other.leak_reference()?;
+        self.0.lock().owned.push(location.raw);
+        Ok(location)
     }
 }

packages/generational-box/src/sync.rs

@@ -1,20 +1,174 @@
 use parking_lot::{
     MappedRwLockReadGuard, MappedRwLockWriteGuard, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard,
 };
-use std::sync::{Arc, OnceLock};
+use std::{
+    any::Any,
+    fmt::Debug,
+    num::NonZeroU64,
+    sync::{Arc, OnceLock},
+};
 
 use crate::{
-    entry::{MemoryLocationBorrowInfo, StorageEntry},
+    entry::{MemoryLocationBorrowInfo, RcStorageEntry, StorageEntry},
     error::{self, ValueDroppedError},
     references::{GenerationalRef, GenerationalRefMut},
-    AnyStorage, GenerationalLocation, GenerationalPointer, Storage,
+    AnyStorage, BorrowError, BorrowMutError, BorrowMutResult, BorrowResult, GenerationalLocation,
+    GenerationalPointer, Storage,
 };
 
+type RwLockStorageEntryRef = RwLockReadGuard<'static, StorageEntry<RwLockStorageEntryData>>;
+type RwLockStorageEntryMut = RwLockWriteGuard<'static, StorageEntry<RwLockStorageEntryData>>;
+
+pub(crate) enum RwLockStorageEntryData {
+    Reference(GenerationalPointer<SyncStorage>),
+    Rc(RcStorageEntry<Box<dyn Any + Send + Sync>>),
+    Data(Box<dyn Any + Send + Sync>),
+    Empty,
+}
+
+impl Default for RwLockStorageEntryData {
+    fn default() -> Self {
+        Self::Empty
+    }
+}
+
+impl Debug for RwLockStorageEntryData {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Reference(location) => write!(f, "Reference({:?})", location),
+            Self::Rc(_) => write!(f, "Rc"),
+            Self::Data(_) => write!(f, "Data"),
+            Self::Empty => write!(f, "Empty"),
+        }
+    }
+}
+
+impl RwLockStorageEntryData {
+    pub const fn new_full(data: Box<dyn Any + Send + Sync>) -> Self {
+        Self::Data(data)
+    }
+}
+
 /// A thread safe storage. This is slower than the unsync storage, but allows you to share the value between threads.
 #[derive(Default)]
 pub struct SyncStorage {
     borrow_info: MemoryLocationBorrowInfo,
-    data: RwLock<StorageEntry<Box<dyn std::any::Any + Send + Sync>>>,
+    data: RwLock<StorageEntry<RwLockStorageEntryData>>,
+}
+
+impl SyncStorage {
+    pub(crate) fn read(
+        pointer: GenerationalPointer<Self>,
+    ) -> BorrowResult<MappedRwLockReadGuard<'static, Box<dyn Any + Send + Sync + 'static>>> {
+        Self::get_split_ref(pointer).map(|(_, guard)| {
+            RwLockReadGuard::map(guard, |data| match &data.data {
+                RwLockStorageEntryData::Data(data) => data,
+                RwLockStorageEntryData::Rc(data) => &data.data,
+                _ => unreachable!(),
+            })
+        })
+    }
+
+    pub(crate) fn get_split_ref(
+        mut pointer: GenerationalPointer<Self>,
+    ) -> BorrowResult<(GenerationalPointer<Self>, RwLockStorageEntryRef)> {
+        loop {
+            let borrow = pointer.storage.data.read();
+            if !borrow.valid(&pointer.location) {
+                return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                    pointer.location,
+                )));
+            }
+            match &borrow.data {
+                // If this is a reference, keep traversing the pointers
+                RwLockStorageEntryData::Reference(data) => {
+                    pointer = *data;
+                }
+                // Otherwise return the value
+                RwLockStorageEntryData::Data(_) | RwLockStorageEntryData::Rc(_) => {
+                    return Ok((pointer, borrow));
+                }
+                RwLockStorageEntryData::Empty => {
+                    return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                        pointer.location,
+                    )));
+                }
+            }
+        }
+    }
+
+    pub(crate) fn write(
+        pointer: GenerationalPointer<Self>,
+    ) -> BorrowMutResult<MappedRwLockWriteGuard<'static, Box<dyn Any + Send + Sync + 'static>>>
+    {
+        Self::get_split_mut(pointer).map(|(_, guard)| {
+            RwLockWriteGuard::map(guard, |data| match &mut data.data {
+                RwLockStorageEntryData::Data(data) => data,
+                RwLockStorageEntryData::Rc(data) => &mut data.data,
+                _ => unreachable!(),
+            })
+        })
+    }
+
+    pub(crate) fn get_split_mut(
+        mut pointer: GenerationalPointer<Self>,
+    ) -> BorrowMutResult<(GenerationalPointer<Self>, RwLockStorageEntryMut)> {
+        loop {
+            let borrow = pointer.storage.data.write();
+            if !borrow.valid(&pointer.location) {
+                return Err(BorrowMutError::Dropped(
+                    ValueDroppedError::new_for_location(pointer.location),
+                ));
+            }
+            match &borrow.data {
+                // If this is a reference, keep traversing the pointers
+                RwLockStorageEntryData::Reference(data) => {
+                    pointer = *data;
+                }
+                // Otherwise return the value
+                RwLockStorageEntryData::Data(_) | RwLockStorageEntryData::Rc(_) => {
+                    return Ok((pointer, borrow));
+                }
+                RwLockStorageEntryData::Empty => {
+                    return Err(BorrowMutError::Dropped(
+                        ValueDroppedError::new_for_location(pointer.location),
+                    ));
+                }
+            }
+        }
+    }
+
+    fn create_new(
+        value: RwLockStorageEntryData,
+        #[allow(unused)] caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalPointer<Self> {
+        match sync_runtime().lock().pop() {
+            Some(storage) => {
+                let mut write = storage.data.write();
+                let location = GenerationalLocation {
+                    generation: write.generation(),
+                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+                    created_at: caller,
+                };
+                write.data = value;
+                GenerationalPointer { storage, location }
+            }
+            None => {
+                let storage: &'static Self = &*Box::leak(Box::new(Self {
+                    borrow_info: Default::default(),
+                    data: RwLock::new(StorageEntry::new(value)),
+                }));
+
+                let location = GenerationalLocation {
+                    generation: NonZeroU64::MIN,
+                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+                    created_at: caller,
+                };
+
+                GenerationalPointer { storage, location }
+            }
+        }
+    }
 }
 
 static SYNC_RUNTIME: OnceLock<Arc<Mutex<Vec<&'static SyncStorage>>>> = OnceLock::new();
@@ -71,42 +225,50 @@ impl AnyStorage for SyncStorage {
         self.data.data_ptr() as *const ()
     }
 
-    #[track_caller]
-    #[allow(unused)]
-    fn claim(caller: &'static std::panic::Location<'static>) -> GenerationalPointer<Self> {
-        match sync_runtime().lock().pop() {
-            Some(mut storage) => {
-                let location = GenerationalLocation {
-                    generation: storage.data.read().generation(),
-                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                    created_at: caller,
-                };
-                GenerationalPointer { storage, location }
-            }
-            None => {
-                let storage: &'static Self = &*Box::leak(Box::default());
-
-                let location = GenerationalLocation {
-                    generation: 0,
-                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                    created_at: caller,
-                };
-
-                GenerationalPointer { storage, location }
-            }
-        }
-    }
-
-    fn recycle(pointer: GenerationalPointer<Self>) -> Option<Box<dyn std::any::Any>> {
+    fn recycle(pointer: GenerationalPointer<Self>) {
         let mut borrow_mut = pointer.storage.data.write();
+
         // First check if the generation is still valid
         if !borrow_mut.valid(&pointer.location) {
-            return None;
+            return;
         }
+
         borrow_mut.increment_generation();
-        let old_data = borrow_mut.data.take();
+
+        // Then decrement the reference count or drop the value if it's the last reference
+        match &mut borrow_mut.data {
+            // If this is the original reference, drop the value
+            RwLockStorageEntryData::Data(_) => borrow_mut.data = RwLockStorageEntryData::Empty,
+            // If this is a rc, just ignore the drop
+            RwLockStorageEntryData::Rc(_) => {}
+            // If this is a reference, decrement the reference count
+            RwLockStorageEntryData::Reference(reference) => {
+                drop_ref(*reference);
+            }
+            RwLockStorageEntryData::Empty => {}
+        }
+
         sync_runtime().lock().push(pointer.storage);
-        old_data.map(|data| data as Box<dyn std::any::Any>)
+    }
+}
+
+fn drop_ref(pointer: GenerationalPointer<SyncStorage>) {
+    let mut borrow_mut = pointer.storage.data.write();
+
+    // First check if the generation is still valid
+    if !borrow_mut.valid(&pointer.location) {
+        return;
+    }
+
+    if let RwLockStorageEntryData::Rc(entry) = &mut borrow_mut.data {
+        // Decrement the reference count
+        if entry.drop_ref() {
+            // If the reference count is now zero, drop the value
+            borrow_mut.data = RwLockStorageEntryData::Empty;
+            sync_runtime().lock().push(pointer.storage);
+        }
+    } else {
+        unreachable!("References should always point to a data entry directly");
     }
 }
 
@@ -115,15 +277,11 @@ impl<T: Sync + Send + 'static> Storage<T> for SyncStorage {
     fn try_read(
         pointer: GenerationalPointer<Self>,
     ) -> Result<Self::Ref<'static, T>, error::BorrowError> {
-        let read = pointer.storage.data.read();
+        let read = Self::read(pointer)?;
 
-        let read = RwLockReadGuard::try_map(read, |any| {
-            // Verify the generation is still correct
-            if !any.valid(&pointer.location) {
-                return None;
-            }
+        let read = MappedRwLockReadGuard::try_map(read, |any| {
             // Then try to downcast
-            any.data.as_ref()?.downcast_ref()
+            any.downcast_ref()
         });
         match read {
             Ok(guard) => Ok(GenerationalRef::new(
@@ -140,15 +298,11 @@ impl<T: Sync + Send + 'static> Storage<T> for SyncStorage {
     fn try_write(
         pointer: GenerationalPointer<Self>,
     ) -> Result<Self::Mut<'static, T>, error::BorrowMutError> {
-        let write = pointer.storage.data.write();
+        let write = Self::write(pointer)?;
 
-        let write = RwLockWriteGuard::try_map(write, |any| {
-            // Verify the generation is still correct
-            if !any.valid(&pointer.location) {
-                return None;
-            }
+        let write = MappedRwLockWriteGuard::try_map(write, |any| {
             // Then try to downcast
-            any.data.as_mut()?.downcast_mut()
+            any.downcast_mut()
         });
         match write {
             Ok(guard) => Ok(GenerationalRefMut::new(
@@ -161,12 +315,78 @@ impl<T: Sync + Send + 'static> Storage<T> for SyncStorage {
         }
     }
 
-    fn set(pointer: GenerationalPointer<Self>, value: T) {
-        let mut write = pointer.storage.data.write();
-        // First check if the generation is still valid
-        if !write.valid(&pointer.location) {
-            return;
+    fn new(value: T, caller: &'static std::panic::Location<'static>) -> GenerationalPointer<Self> {
+        Self::create_new(RwLockStorageEntryData::new_full(Box::new(value)), caller)
+    }
+
+    fn new_rc(
+        value: T,
+        caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalPointer<Self> {
+        // Create the data that the rc points to
+        let data = Self::create_new(
+            RwLockStorageEntryData::Rc(RcStorageEntry::new(Box::new(value))),
+            caller,
+        );
+        Self::create_new(RwLockStorageEntryData::Reference(data), caller)
+    }
+
+    fn new_reference(
+        location: GenerationalPointer<Self>,
+    ) -> BorrowResult<GenerationalPointer<Self>> {
+        // Chase the reference to get the final location
+        let (location, value) = Self::get_split_ref(location)?;
+        if let RwLockStorageEntryData::Rc(data) = &value.data {
+            data.add_ref();
+        } else {
+            unreachable!()
         }
-        write.data = Some(Box::new(value));
+        Ok(Self::create_new(
+            RwLockStorageEntryData::Reference(location),
+            location
+                .location
+                .created_at()
+                .unwrap_or(std::panic::Location::caller()),
+        ))
+    }
+
+    fn change_reference(
+        location: GenerationalPointer<Self>,
+        other: GenerationalPointer<Self>,
+    ) -> BorrowResult {
+        if location == other {
+            return Ok(());
+        }
+
+        let (other_final, other_write) = Self::get_split_ref(other)?;
+
+        let mut write = location.storage.data.write();
+        // First check if the generation is still valid
+        if !write.valid(&location.location) {
+            return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                location.location,
+            )));
+        }
+
+        if let (RwLockStorageEntryData::Reference(reference), RwLockStorageEntryData::Rc(data)) =
+            (&mut write.data, &other_write.data)
+        {
+            if reference == &other_final {
+                return Ok(());
+            }
+            drop_ref(*reference);
+            *reference = other_final;
+            data.add_ref();
+        } else {
+            tracing::trace!(
+                "References should always point to a data entry directly found {:?} instead",
+                other_write.data
+            );
+            return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                other_final.location,
+            )));
+        }
+
+        Ok(())
     }
 }

packages/generational-box/src/unsync.rs

@@ -1,35 +1,175 @@
 use crate::{
-    entry::{MemoryLocationBorrowInfo, StorageEntry},
+    entry::{MemoryLocationBorrowInfo, RcStorageEntry, StorageEntry},
     error,
     references::{GenerationalRef, GenerationalRefMut},
-    AnyStorage, BorrowError, BorrowMutError, GenerationalLocation, GenerationalPointer, Storage,
+    AnyStorage, BorrowError, BorrowMutError, BorrowMutResult, BorrowResult, GenerationalLocation,
+    GenerationalPointer, Storage, ValueDroppedError,
 };
-use std::cell::{Ref, RefCell, RefMut};
+use std::{
+    any::Any,
+    cell::{Ref, RefCell, RefMut},
+    fmt::Debug,
+    num::NonZeroU64,
+};
+
+type RefCellStorageEntryRef = Ref<'static, StorageEntry<RefCellStorageEntryData>>;
+
+type RefCellStorageEntryMut = RefMut<'static, StorageEntry<RefCellStorageEntryData>>;
 
 thread_local! {
     static UNSYNC_RUNTIME: RefCell<Vec<&'static UnsyncStorage>> = const { RefCell::new(Vec::new()) };
 }
 
+pub(crate) enum RefCellStorageEntryData {
+    Reference(GenerationalPointer<UnsyncStorage>),
+    Rc(RcStorageEntry<Box<dyn Any>>),
+    Data(Box<dyn Any>),
+    Empty,
+}
+
+impl Debug for RefCellStorageEntryData {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Reference(pointer) => write!(f, "Reference({:?})", pointer.location),
+            Self::Rc(_) => write!(f, "Rc"),
+            Self::Data(_) => write!(f, "Data"),
+            Self::Empty => write!(f, "Empty"),
+        }
+    }
+}
+
+impl Default for RefCellStorageEntryData {
+    fn default() -> Self {
+        Self::Empty
+    }
+}
+
 /// A unsync storage. This is the default storage type.
 #[derive(Default)]
 pub struct UnsyncStorage {
     borrow_info: MemoryLocationBorrowInfo,
-    data: RefCell<StorageEntry<Box<dyn std::any::Any>>>,
+    data: RefCell<StorageEntry<RefCellStorageEntryData>>,
 }
 
 impl UnsyncStorage {
-    fn try_borrow_mut(
-        &self,
-    ) -> Result<RefMut<'_, StorageEntry<Box<dyn std::any::Any>>>, BorrowMutError> {
-        self.data
-            .try_borrow_mut()
-            .map_err(|_| self.borrow_info.borrow_mut_error())
+    pub(crate) fn read(
+        pointer: GenerationalPointer<Self>,
+    ) -> BorrowResult<Ref<'static, Box<dyn Any>>> {
+        Self::get_split_ref(pointer).map(|(_, guard)| {
+            Ref::map(guard, |data| match &data.data {
+                RefCellStorageEntryData::Data(data) => data,
+                RefCellStorageEntryData::Rc(data) => &data.data,
+                _ => unreachable!(),
+            })
+        })
     }
 
-    fn try_borrow(&self) -> Result<Ref<'_, StorageEntry<Box<dyn std::any::Any>>>, BorrowError> {
-        self.data
-            .try_borrow()
-            .map_err(|_| self.borrow_info.borrow_error())
+    pub(crate) fn get_split_ref(
+        mut pointer: GenerationalPointer<Self>,
+    ) -> BorrowResult<(GenerationalPointer<Self>, RefCellStorageEntryRef)> {
+        loop {
+            let borrow = pointer
+                .storage
+                .data
+                .try_borrow()
+                .map_err(|_| pointer.storage.borrow_info.borrow_error())?;
+            if !borrow.valid(&pointer.location) {
+                return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                    pointer.location,
+                )));
+            }
+            match &borrow.data {
+                // If this is a reference, keep traversing the pointers
+                RefCellStorageEntryData::Reference(data) => {
+                    pointer = *data;
+                }
+                // Otherwise return the value
+                RefCellStorageEntryData::Rc(_) | RefCellStorageEntryData::Data(_) => {
+                    return Ok((pointer, borrow));
+                }
+                RefCellStorageEntryData::Empty => {
+                    return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                        pointer.location,
+                    )));
+                }
+            }
+        }
+    }
+
+    pub(crate) fn write(
+        pointer: GenerationalPointer<Self>,
+    ) -> BorrowMutResult<RefMut<'static, Box<dyn Any>>> {
+        Self::get_split_mut(pointer).map(|(_, guard)| {
+            RefMut::map(guard, |data| match &mut data.data {
+                RefCellStorageEntryData::Data(data) => data,
+                RefCellStorageEntryData::Rc(data) => &mut data.data,
+                _ => unreachable!(),
+            })
+        })
+    }
+
+    pub(crate) fn get_split_mut(
+        mut pointer: GenerationalPointer<Self>,
+    ) -> BorrowMutResult<(GenerationalPointer<Self>, RefCellStorageEntryMut)> {
+        loop {
+            let borrow = pointer
+                .storage
+                .data
+                .try_borrow_mut()
+                .map_err(|_| pointer.storage.borrow_info.borrow_mut_error())?;
+            if !borrow.valid(&pointer.location) {
+                return Err(BorrowMutError::Dropped(
+                    ValueDroppedError::new_for_location(pointer.location),
+                ));
+            }
+            match &borrow.data {
+                // If this is a reference, keep traversing the pointers
+                RefCellStorageEntryData::Reference(data) => {
+                    pointer = *data;
+                }
+                // Otherwise return the value
+                RefCellStorageEntryData::Data(_) | RefCellStorageEntryData::Rc(_) => {
+                    return Ok((pointer, borrow));
+                }
+                RefCellStorageEntryData::Empty => {
+                    return Err(BorrowMutError::Dropped(
+                        ValueDroppedError::new_for_location(pointer.location),
+                    ));
+                }
+            }
+        }
+    }
+
+    fn create_new(
+        value: RefCellStorageEntryData,
+        #[allow(unused)] caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalPointer<Self> {
+        UNSYNC_RUNTIME.with(|runtime| match runtime.borrow_mut().pop() {
+            Some(storage) => {
+                let mut write = storage.data.borrow_mut();
+                let location = GenerationalLocation {
+                    generation: write.generation(),
+                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+                    created_at: caller,
+                };
+                write.data = value;
+                GenerationalPointer { storage, location }
+            }
+            None => {
+                let storage: &'static Self = &*Box::leak(Box::new(Self {
+                    borrow_info: Default::default(),
+                    data: RefCell::new(StorageEntry::new(value)),
+                }));
+
+                let location = GenerationalLocation {
+                    generation: NonZeroU64::MIN,
+                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+                    created_at: caller,
+                };
+
+                GenerationalPointer { storage, location }
+            }
+        })
     }
 }
 
@@ -81,44 +221,51 @@ impl AnyStorage for UnsyncStorage {
         self.data.as_ptr() as *const ()
     }
 
-    #[allow(unused)]
-    fn claim(caller: &'static std::panic::Location<'static>) -> GenerationalPointer<Self> {
-        UNSYNC_RUNTIME.with(|runtime| {
-            if let Some(storage) = runtime.borrow_mut().pop() {
-                let location = GenerationalLocation {
-                    generation: storage.data.borrow().generation(),
-                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                    created_at: caller,
-                };
-                GenerationalPointer { storage, location }
-            } else {
-                let data: &'static Self = &*Box::leak(Box::default());
-                let location = GenerationalLocation {
-                    generation: 0,
-                    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                    created_at: caller,
-                };
-                GenerationalPointer {
-                    storage: data,
-                    location,
-                }
-            }
-        })
-    }
-
-    fn recycle(pointer: GenerationalPointer<Self>) -> Option<Box<dyn std::any::Any>> {
+    fn recycle(pointer: GenerationalPointer<Self>) {
         let mut borrow_mut = pointer.storage.data.borrow_mut();
 
         // First check if the generation is still valid
         if !borrow_mut.valid(&pointer.location) {
-            return None;
+            return;
         }
 
         borrow_mut.increment_generation();
-        let old_data = borrow_mut.data.take();
-        UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(pointer.storage));
+        // Then decrement the reference count or drop the value if it's the last reference
+        match &mut borrow_mut.data {
+            // If this is the original reference, drop the value
+            RefCellStorageEntryData::Data(_) => borrow_mut.data = RefCellStorageEntryData::Empty,
+            // If this is a rc, just ignore the drop
+            RefCellStorageEntryData::Rc(_) => {}
+            // If this is a reference, decrement the reference count
+            RefCellStorageEntryData::Reference(reference) => {
+                let reference = *reference;
+                drop(borrow_mut);
+                drop_ref(reference);
+            }
+            RefCellStorageEntryData::Empty => {}
+        }
 
-        old_data
+        UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(pointer.storage));
+    }
+}
+
+fn drop_ref(pointer: GenerationalPointer<UnsyncStorage>) {
+    let mut borrow_mut = pointer.storage.data.borrow_mut();
+
+    // First check if the generation is still valid
+    if !borrow_mut.valid(&pointer.location) {
+        return;
+    }
+
+    if let RefCellStorageEntryData::Rc(entry) = &mut borrow_mut.data {
+        // Decrement the reference count
+        if entry.drop_ref() {
+            // If the reference count is now zero, drop the value
+            borrow_mut.data = RefCellStorageEntryData::Empty;
+            UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(pointer.storage));
+        }
+    } else {
+        unreachable!("References should always point to a data entry directly",);
     }
 }
 
@@ -127,15 +274,11 @@ impl<T: 'static> Storage<T> for UnsyncStorage {
     fn try_read(
         pointer: GenerationalPointer<Self>,
     ) -> Result<Self::Ref<'static, T>, error::BorrowError> {
-        let read = pointer.storage.try_borrow()?;
+        let read = Self::read(pointer)?;
 
         let ref_ = Ref::filter_map(read, |any| {
-            // Verify the generation is still correct
-            if !any.valid(&pointer.location) {
-                return None;
-            }
             // Then try to downcast
-            any.data.as_ref()?.downcast_ref()
+            any.downcast_ref()
         });
         match ref_ {
             Ok(guard) => Ok(GenerationalRef::new(
@@ -152,15 +295,11 @@ impl<T: 'static> Storage<T> for UnsyncStorage {
     fn try_write(
         pointer: GenerationalPointer<Self>,
     ) -> Result<Self::Mut<'static, T>, error::BorrowMutError> {
-        let write = pointer.storage.try_borrow_mut()?;
+        let write = Self::write(pointer)?;
 
         let ref_mut = RefMut::filter_map(write, |any| {
-            // Verify the generation is still correct
-            if !any.valid(&pointer.location) {
-                return None;
-            }
             // Then try to downcast
-            any.data.as_mut()?.downcast_mut()
+            any.downcast_mut()
         });
         match ref_mut {
             Ok(guard) => Ok(GenerationalRefMut::new(
@@ -173,12 +312,78 @@ impl<T: 'static> Storage<T> for UnsyncStorage {
         }
     }
 
-    fn set(pointer: GenerationalPointer<Self>, value: T) {
-        let mut borrow_mut = pointer.storage.data.borrow_mut();
-        // First check if the generation is still valid
-        if !borrow_mut.valid(&pointer.location) {
-            return;
+    fn new(value: T, caller: &'static std::panic::Location<'static>) -> GenerationalPointer<Self> {
+        Self::create_new(RefCellStorageEntryData::Data(Box::new(value)), caller)
+    }
+
+    fn new_rc(
+        value: T,
+        caller: &'static std::panic::Location<'static>,
+    ) -> GenerationalPointer<Self> {
+        // Create the data that the rc points to
+        let data = Self::create_new(
+            RefCellStorageEntryData::Rc(RcStorageEntry::new(Box::new(value))),
+            caller,
+        );
+        Self::create_new(RefCellStorageEntryData::Reference(data), caller)
+    }
+
+    fn new_reference(
+        pointer: GenerationalPointer<Self>,
+    ) -> BorrowResult<GenerationalPointer<Self>> {
+        // Chase the reference to get the final location
+        let (pointer, value) = Self::get_split_ref(pointer)?;
+        if let RefCellStorageEntryData::Rc(data) = &value.data {
+            data.add_ref();
+        } else {
+            unreachable!()
         }
-        borrow_mut.data = Some(Box::new(value));
+        Ok(Self::create_new(
+            RefCellStorageEntryData::Reference(pointer),
+            pointer
+                .location
+                .created_at()
+                .unwrap_or(std::panic::Location::caller()),
+        ))
+    }
+
+    fn change_reference(
+        location: GenerationalPointer<Self>,
+        other: GenerationalPointer<Self>,
+    ) -> BorrowResult {
+        if location == other {
+            return Ok(());
+        }
+
+        let (other_final, other_write) = Self::get_split_ref(other)?;
+
+        let mut write = location.storage.data.borrow_mut();
+        // First check if the generation is still valid
+        if !write.valid(&location.location) {
+            return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                location.location,
+            )));
+        }
+
+        if let (RefCellStorageEntryData::Reference(reference), RefCellStorageEntryData::Rc(data)) =
+            (&mut write.data, &other_write.data)
+        {
+            if reference == &other_final {
+                return Ok(());
+            }
+            drop_ref(*reference);
+            *reference = other_final;
+            data.add_ref();
+        } else {
+            tracing::trace!(
+                "References should always point to a data entry directly found {:?} instead",
+                other_write.data
+            );
+            return Err(BorrowError::Dropped(ValueDroppedError::new_for_location(
+                other_final.location,
+            )));
+        }
+
+        Ok(())
     }
 }