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/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(())
     }
 }