Fix generation race condition with sync storage (#2638)

* Fix generation race condition with sync storage

* Document GenerationalPointer

* check generation before recycling generational box

* fix clippy

* make race_condition_regression test panic if it fails

packages/core/src/nodes.rs

@@ -203,7 +203,7 @@ impl Drop for VNode {
             for attrs in self.vnode.dynamic_attrs.iter() {
                 for attr in attrs.iter() {
                     if let AttributeValue::Listener(listener) = &attr.value {
-                        listener.callback.recycle();
+                        listener.callback.manually_drop();
                     }
                 }
             }

packages/generational-box/Cargo.toml

@@ -17,8 +17,6 @@ rand = "0.8.5"
 criterion = "0.3"
 
 [features]
-default = ["check_generation"]
-check_generation = []
 debug_borrows = []
 debug_ownership = []
 

packages/generational-box/src/entry.rs

@@ -0,0 +1,135 @@
+use crate::{
+    BorrowError, BorrowMutError, GenerationalLocation, GenerationalRefBorrowGuard,
+    GenerationalRefBorrowMutGuard,
+};
+
+pub(crate) struct StorageEntry<T> {
+    generation: u64,
+
+    pub data: Option<T>,
+}
+
+impl<T> Default for StorageEntry<T> {
+    fn default() -> Self {
+        Self {
+            generation: 0,
+            data: None,
+        }
+    }
+}
+
+impl<T> StorageEntry<T> {
+    pub fn valid(&self, location: &GenerationalLocation) -> bool {
+        self.generation == location.generation
+    }
+
+    pub fn increment_generation(&mut self) {
+        self.generation += 1;
+    }
+
+    pub fn generation(&self) -> u64 {
+        self.generation
+    }
+}
+
+#[derive(Default)]
+pub(crate) struct MemoryLocationBorrowInfo(
+    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+    parking_lot::RwLock<MemoryLocationBorrowInfoInner>,
+);
+
+impl MemoryLocationBorrowInfo {
+    pub(crate) fn borrow_mut_error(&self) -> BorrowMutError {
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        {
+            let borrow = self.0.read();
+            if let Some(borrowed_mut_at) = borrow.borrowed_mut_at.as_ref() {
+                BorrowMutError::AlreadyBorrowedMut(crate::error::AlreadyBorrowedMutError {
+                    borrowed_mut_at,
+                })
+            } else {
+                BorrowMutError::AlreadyBorrowed(crate::error::AlreadyBorrowedError {
+                    borrowed_at: borrow.borrowed_at.clone(),
+                })
+            }
+        }
+        #[cfg(not(any(debug_assertions, feature = "debug_borrows")))]
+        {
+            BorrowMutError::AlreadyBorrowed(crate::error::AlreadyBorrowedError {})
+        }
+    }
+
+    pub(crate) fn borrow_error(&self) -> BorrowError {
+        BorrowError::AlreadyBorrowedMut(crate::error::AlreadyBorrowedMutError {
+            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
+            borrowed_mut_at: self.0.read().borrowed_mut_at.unwrap(),
+        })
+    }
+
+    /// Start a new borrow
+    #[track_caller]
+    pub(crate) fn borrow_guard(&'static self) -> GenerationalRefBorrowGuard {
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        let borrowed_at = std::panic::Location::caller();
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        {
+            let mut borrow = self.0.write();
+            borrow.borrowed_at.push(borrowed_at);
+        }
+
+        GenerationalRefBorrowGuard {
+            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+            borrowed_at,
+            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+            borrowed_from: self,
+        }
+    }
+
+    /// Start a new mutable borrow
+    #[track_caller]
+    pub(crate) fn borrow_mut_guard(&'static self) -> GenerationalRefBorrowMutGuard {
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        let borrowed_mut_at = std::panic::Location::caller();
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        {
+            let mut borrow = self.0.write();
+            borrow.borrowed_mut_at = Some(borrowed_mut_at);
+        }
+
+        GenerationalRefBorrowMutGuard {
+            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+            borrowed_mut_at,
+            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+            borrowed_from: self,
+        }
+    }
+
+    #[allow(unused)]
+    pub(crate) fn drop_borrow(&self, borrowed_at: &'static std::panic::Location<'static>) {
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        {
+            let mut borrow = self.0.write();
+            borrow
+                .borrowed_at
+                .retain(|location| *location != borrowed_at);
+        }
+    }
+
+    #[allow(unused)]
+    pub(crate) fn drop_borrow_mut(&self, borrowed_mut_at: &'static std::panic::Location<'static>) {
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        {
+            let mut borrow = self.0.write();
+            if borrow.borrowed_mut_at == Some(borrowed_mut_at) {
+                borrow.borrowed_mut_at = None;
+            }
+        }
+    }
+}
+
+#[cfg(any(debug_assertions, feature = "debug_borrows"))]
+#[derive(Default)]
+struct MemoryLocationBorrowInfoInner {
+    borrowed_at: Vec<&'static std::panic::Location<'static>>,
+    borrowed_mut_at: Option<&'static std::panic::Location<'static>>,
+}

packages/generational-box/src/error.rs

@@ -2,6 +2,8 @@ use std::error::Error;
 use std::fmt::Debug;
 use std::fmt::Display;
 
+use crate::GenerationalLocation;
+
 #[derive(Debug, Clone, PartialEq)]
 /// An error that can occur when trying to borrow a value.
 pub enum BorrowError {
@@ -61,6 +63,15 @@ impl ValueDroppedError {
             created_at,
         }
     }
+
+    /// Create a new `ValueDroppedError` for a [`GenerationalLocation`].
+    #[allow(unused)]
+    pub(crate) fn new_for_location(location: GenerationalLocation) -> Self {
+        Self {
+            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+            created_at: location.created_at,
+        }
+    }
 }
 
 impl Display for ValueDroppedError {

packages/generational-box/src/lib.rs

@@ -14,6 +14,7 @@ pub use references::*;
 pub use sync::SyncStorage;
 pub use unsync::UnsyncStorage;
 
+mod entry;
 mod error;
 mod references;
 mod sync;
@@ -23,8 +24,7 @@ mod unsync;
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct GenerationalBoxId {
     data_ptr: *const (),
-    #[cfg(any(debug_assertions, feature = "check_generation"))]
-    generation: u32,
+    generation: u64,
 }
 
 // Safety: GenerationalBoxId is Send and Sync because there is no way to access the pointer.
@@ -33,35 +33,20 @@ unsafe impl Sync for GenerationalBoxId {}
 
 impl Debug for GenerationalBoxId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        #[cfg(any(debug_assertions, feature = "check_generation"))]
         f.write_fmt(format_args!("{:?}@{:?}", self.data_ptr, self.generation))?;
-        #[cfg(not(any(debug_assertions, feature = "check_generation")))]
-        f.write_fmt(format_args!("{:?}", self.data_ptr))?;
         Ok(())
     }
 }
 
 /// The core Copy state type. The generational box will be dropped when the [Owner] is dropped.
 pub struct GenerationalBox<T, S: 'static = UnsyncStorage> {
-    raw: MemoryLocation<S>,
-    #[cfg(any(debug_assertions, feature = "check_generation"))]
-    generation: u32,
-    #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-    created_at: &'static std::panic::Location<'static>,
+    raw: GenerationalPointer<S>,
     _marker: PhantomData<T>,
 }
 
 impl<T, S: AnyStorage> Debug for GenerationalBox<T, S> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        #[cfg(any(debug_assertions, feature = "check_generation"))]
-        f.write_fmt(format_args!(
-            "{:?}@{:?}",
-            self.raw.0.data.data_ptr(),
-            self.generation
-        ))?;
-        #[cfg(not(any(debug_assertions, feature = "check_generation")))]
-        f.write_fmt(format_args!("{:?}", self.raw.0.data.data_ptr()))?;
-        Ok(())
+        self.raw.fmt(f)
     }
 }
 
@@ -70,73 +55,28 @@ impl<T, S: Storage<T>> GenerationalBox<T, S> {
     /// a box that needs to be manually dropped with no owners.
     #[track_caller]
     pub fn leak(value: T) -> Self {
-        let mut location = S::claim();
-        location.replace_with_caller(
-            value,
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            std::panic::Location::caller(),
-        )
-    }
-
-    #[inline(always)]
-    pub(crate) fn validate(&self) -> bool {
-        #[cfg(any(debug_assertions, feature = "check_generation"))]
-        {
-            self.raw
-                .0
-                .generation
-                .load(std::sync::atomic::Ordering::Relaxed)
-                == self.generation
-        }
-        #[cfg(not(any(debug_assertions, feature = "check_generation")))]
-        {
-            true
+        let location = S::claim(std::panic::Location::caller());
+        location.set(value);
+        Self {
+            raw: location,
+            _marker: PhantomData,
         }
     }
 
     /// Get the raw pointer to the value.
     pub fn raw_ptr(&self) -> *const () {
-        self.raw.0.data.data_ptr()
+        self.raw.storage.data_ptr()
     }
 
     /// Get the id of the generational box.
     pub fn id(&self) -> GenerationalBoxId {
-        GenerationalBoxId {
-            data_ptr: self.raw.0.data.data_ptr(),
-            #[cfg(any(debug_assertions, feature = "check_generation"))]
-            generation: self.generation,
-        }
+        self.raw.id()
     }
 
-    /// Try to read the value. Returns None if the value is no longer valid.
+    /// Try to read the value. Returns an error if the value is no longer valid.
     #[track_caller]
     pub fn try_read(&self) -> Result<S::Ref<'static, T>, BorrowError> {
-        if !self.validate() {
-            return Err(BorrowError::Dropped(ValueDroppedError {
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                created_at: self.created_at,
-            }));
-        }
-        let result = self.raw.0.data.try_read(
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            GenerationalRefBorrowInfo {
-                borrowed_at: std::panic::Location::caller(),
-                borrowed_from: &self.raw.0.borrow,
-                created_at: self.created_at,
-            },
-        );
-
-        if result.is_ok() {
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            self.raw
-                .0
-                .borrow
-                .borrowed_at
-                .write()
-                .push(std::panic::Location::caller());
-        }
-
-        result
+        self.raw.try_read()
     }
 
     /// Read the value. Panics if the value is no longer valid.
@@ -148,28 +88,7 @@ impl<T, S: Storage<T>> GenerationalBox<T, S> {
     /// Try to write the value. Returns None if the value is no longer valid.
     #[track_caller]
     pub fn try_write(&self) -> Result<S::Mut<'static, T>, BorrowMutError> {
-        if !self.validate() {
-            return Err(BorrowMutError::Dropped(ValueDroppedError {
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                created_at: self.created_at,
-            }));
-        }
-        let result = self.raw.0.data.try_write(
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            GenerationalRefMutBorrowInfo {
-                borrowed_from: &self.raw.0.borrow,
-                created_at: self.created_at,
-            },
-        );
-
-        if result.is_ok() {
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            {
-                *self.raw.0.borrow.borrowed_mut_at.write() = Some(std::panic::Location::caller());
-            }
-        }
-
-        result
+        self.raw.try_write()
     }
 
     /// Write the value. Panics if the value is no longer valid.
@@ -180,42 +99,21 @@ 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) {
-        self.validate().then(|| {
-            self.raw.0.data.set(value);
-        });
+        S::set(self.raw, value);
     }
 
     /// Returns true if the pointer is equal to the other pointer.
     pub fn ptr_eq(&self, other: &Self) -> bool {
-        #[cfg(any(debug_assertions, feature = "check_generation"))]
-        {
-            self.raw.0.data.data_ptr() == other.raw.0.data.data_ptr()
-                && self.generation == other.generation
-        }
-        #[cfg(not(any(debug_assertions, feature = "check_generation")))]
-        {
-            self.raw.0.data.data_ptr() == other.raw.0.data.data_ptr()
-        }
-    }
-
-    /// Recycle the generationalbox, dropping the value.
-    pub fn recycle(&self) {
-        if self.validate() {
-            <S as AnyStorage>::recycle(&self.raw);
-        }
+        self.raw == other.raw
     }
 
     /// 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> {
-        if self.validate() {
-            // TODO: Next breaking release we should change the take method to automatically recycle the box
-            let value = Storage::take(&self.raw.0.data)?;
-            <S as AnyStorage>::recycle(&self.raw);
-            Some(value)
-        } else {
-            None
-        }
+    pub fn manually_drop(&self) -> Option<T>
+    where
+        T: 'static,
+    {
+        self.raw.take()
     }
 }
 
@@ -231,25 +129,20 @@ impl<T, S> Clone for GenerationalBox<T, S> {
 pub trait Storage<Data = ()>: AnyStorage + 'static {
     /// Try to read the value. Returns None if the value is no longer valid.
     fn try_read(
-        &'static self,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))] at: GenerationalRefBorrowInfo,
+        location: GenerationalPointer<Self>,
     ) -> Result<Self::Ref<'static, Data>, BorrowError>;
 
     /// Try to write the value. Returns None if the value is no longer valid.
     fn try_write(
-        &'static self,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))] at: GenerationalRefMutBorrowInfo,
+        location: GenerationalPointer<Self>,
     ) -> Result<Self::Mut<'static, Data>, BorrowMutError>;
 
-    /// Set the value
-    fn set(&'static self, value: Data);
-
-    /// Take the value out of the storage. This will return the value if the value was taken.
-    fn take(&'static self) -> Option<Data>;
+    /// Set the value if the location is valid
+    fn set(location: GenerationalPointer<Self>, value: Data);
 }
 
 /// A trait for any storage backing type.
-pub trait AnyStorage: Default {
+pub trait AnyStorage: Default + 'static {
     /// The reference this storage type returns.
     type Ref<'a, T: ?Sized + 'static>: Deref<Target = T>;
     /// The mutable reference this storage type returns.
@@ -300,14 +193,11 @@ pub trait AnyStorage: Default {
     /// Get the data pointer. No guarantees are made about the data pointer. It should only be used for debugging.
     fn data_ptr(&self) -> *const ();
 
-    /// Drop the value from the storage. This will return true if the value was taken.
-    fn manually_drop(&self) -> bool;
-
     /// Recycle a memory location. This will drop the memory location and return it to the runtime.
-    fn recycle(location: &MemoryLocation<Self>);
+    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() -> MemoryLocation<Self>;
+    fn claim(caller: &'static std::panic::Location<'static>) -> 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> {
@@ -317,142 +207,104 @@ pub trait AnyStorage: Default {
     }
 }
 
-/// A dynamic memory location that can be used in a generational box.
-pub struct MemoryLocation<S: 'static = UnsyncStorage>(&'static MemoryLocationInner<S>);
-
-impl<S: 'static> Clone for MemoryLocation<S> {
-    fn clone(&self) -> Self {
-        *self
-    }
+#[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,
+    #[cfg(any(debug_assertions, feature = "debug_ownership"))]
+    created_at: &'static std::panic::Location<'static>,
 }
 
-impl<S: 'static> Copy for MemoryLocation<S> {}
-
-#[cfg(any(debug_assertions, feature = "debug_borrows"))]
-#[derive(Debug, Default)]
-struct MemoryLocationBorrowInfo {
-    pub(crate) borrowed_at: parking_lot::RwLock<Vec<&'static std::panic::Location<'static>>>,
-    pub(crate) borrowed_mut_at: parking_lot::RwLock<Option<&'static std::panic::Location<'static>>>,
+/// 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
+    storage: &'static S,
+    /// The location of the data
+    location: GenerationalLocation,
 }
 
-#[cfg(any(debug_assertions, feature = "debug_ownership"))]
-impl MemoryLocationBorrowInfo {
-    fn borrow_mut_error(&self) -> BorrowMutError {
-        if let Some(borrowed_mut_at) = self.borrowed_mut_at.read().as_ref() {
-            BorrowMutError::AlreadyBorrowedMut(crate::error::AlreadyBorrowedMutError {
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                borrowed_mut_at,
-            })
-        } else {
-            BorrowMutError::AlreadyBorrowed(crate::error::AlreadyBorrowedError {
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                borrowed_at: self.borrowed_at.read().clone(),
-            })
-        }
+impl<S: AnyStorage + 'static> PartialEq for GenerationalPointer<S> {
+    fn eq(&self, other: &Self) -> bool {
+        self.storage.data_ptr() == other.storage.data_ptr()
+            && self.location.generation == other.location.generation
     }
+}
 
-    fn borrow_error(&self) -> BorrowError {
-        BorrowError::AlreadyBorrowedMut(crate::error::AlreadyBorrowedMutError {
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            borrowed_mut_at: self.borrowed_mut_at.read().unwrap(),
-        })
+impl<S: AnyStorage + 'static> Debug for GenerationalPointer<S> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_fmt(format_args!(
+            "{:?}@{:?}",
+            self.storage.data_ptr(),
+            self.location.generation
+        ))
     }
 }
 
-struct MemoryLocationInner<S = UnsyncStorage> {
-    data: S,
-
-    #[cfg(any(debug_assertions, feature = "check_generation"))]
-    generation: std::sync::atomic::AtomicU32,
-
-    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-    borrow: MemoryLocationBorrowInfo,
+impl<S: 'static> Clone for GenerationalPointer<S> {
+    fn clone(&self) -> Self {
+        *self
+    }
 }
 
-impl<S> MemoryLocation<S> {
-    #[allow(unused)]
-    fn drop(&self)
+impl<S: 'static> Copy for GenerationalPointer<S> {}
+
+impl<S> GenerationalPointer<S> {
+    fn take<T: 'static>(self) -> Option<T>
     where
-        S: AnyStorage,
+        S: Storage<T>,
     {
-        self.0.data.manually_drop();
-
-        #[cfg(any(debug_assertions, feature = "check_generation"))]
-        {
-            let new_generation = self.0.generation.load(std::sync::atomic::Ordering::Relaxed) + 1;
-            self.0
-                .generation
-                .store(new_generation, std::sync::atomic::Ordering::Relaxed);
-        }
+        S::recycle(self).map(|value| *(value.downcast().unwrap()))
     }
 
-    fn replace_with_caller<T>(
-        &mut self,
-        value: T,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        caller: &'static std::panic::Location<'static>,
-    ) -> GenerationalBox<T, S>
+    fn set<T>(self, value: T)
     where
         S: Storage<T>,
     {
-        self.0.data.set(value);
-        GenerationalBox {
-            raw: *self,
-            #[cfg(any(debug_assertions, feature = "check_generation"))]
-            generation: self.0.generation.load(std::sync::atomic::Ordering::Relaxed),
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            created_at: caller,
-            _marker: PhantomData,
-        }
+        S::set(self, value)
     }
-}
 
-// We track the generation along with the memory location so that when generational boxes are dropped early, we don't end up dropping the new occupant of the slot
-struct LocationKey<S: 'static> {
-    #[cfg(any(debug_assertions, feature = "check_generation"))]
-    generation: u32,
-    location: MemoryLocation<S>,
-}
+    #[track_caller]
+    fn try_read<T>(self) -> Result<S::Ref<'static, T>, BorrowError>
+    where
+        S: Storage<T>,
+    {
+        S::try_read(self)
+    }
 
-impl<S: AnyStorage> LocationKey<S> {
-    fn exists(&self) -> bool {
-        #[cfg(any(debug_assertions, feature = "check_generation"))]
-        return self.generation
-            == self
-                .location
-                .0
-                .generation
-                .load(std::sync::atomic::Ordering::Relaxed);
-        #[cfg(not(any(debug_assertions, feature = "check_generation")))]
-        return true;
+    #[track_caller]
+    fn try_write<T>(self) -> Result<S::Mut<'static, T>, BorrowMutError>
+    where
+        S: Storage<T>,
+    {
+        S::try_write(self)
     }
 
-    fn drop(self) {
-        // If this is the same box we own, we can just drop it
-        if self.exists() {
-            S::recycle(&self.location);
-        }
+    fn recycle(self)
+    where
+        S: AnyStorage,
+    {
+        S::recycle(self);
     }
-}
 
-impl<T, S: AnyStorage> From<GenerationalBox<T, S>> for LocationKey<S> {
-    fn from(value: GenerationalBox<T, S>) -> Self {
-        Self {
-            #[cfg(any(debug_assertions, feature = "check_generation"))]
-            generation: value.generation,
-            location: value.raw,
+    fn id(&self) -> GenerationalBoxId
+    where
+        S: AnyStorage,
+    {
+        GenerationalBoxId {
+            data_ptr: self.storage.data_ptr(),
+            generation: self.location.generation,
         }
     }
 }
 
 struct OwnerInner<S: AnyStorage + 'static> {
-    owned: Vec<LocationKey<S>>,
+    owned: Vec<GenerationalPointer<S>>,
 }
 
 impl<S: AnyStorage> Drop for OwnerInner<S> {
     fn drop(&mut self) {
-        for key in self.owned.drain(..) {
-            key.drop();
+        for location in self.owned.drain(..) {
+            location.recycle();
         }
     }
 }
@@ -479,48 +331,35 @@ impl<S: AnyStorage> Owner<S> {
     where
         S: Storage<T>,
     {
-        self.insert_with_caller(
-            value,
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            std::panic::Location::caller(),
-        )
+        self.insert_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_with_caller<T: 'static>(
         &self,
         value: T,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
         caller: &'static std::panic::Location<'static>,
     ) -> GenerationalBox<T, S>
     where
         S: Storage<T>,
     {
-        let mut location = S::claim();
-        let key = location.replace_with_caller(
-            value,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            caller,
-        );
-        self.0.lock().owned.push(key.into());
-        key
+        let location = S::claim(caller);
+        location.set(value);
+        self.0.lock().owned.push(location);
+        GenerationalBox {
+            raw: location,
+            _marker: PhantomData,
+        }
     }
 
     /// 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.
+    #[track_caller]
     pub fn invalid<T: 'static>(&self) -> GenerationalBox<T, S> {
-        let location = S::claim();
-        let generational_box = GenerationalBox {
+        let location = S::claim(std::panic::Location::caller());
+        self.0.lock().owned.push(location);
+        GenerationalBox {
             raw: location,
-            #[cfg(any(debug_assertions, feature = "check_generation"))]
-            generation: location
-                .0
-                .generation
-                .load(std::sync::atomic::Ordering::Relaxed),
-            #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-            created_at: std::panic::Location::caller(),
             _marker: PhantomData,
-        };
-        self.0.lock().owned.push(generational_box.into());
-        generational_box
+        }
     }
 }

packages/generational-box/src/references.rs

@@ -6,21 +6,29 @@ use std::{
 /// A reference to a value in a generational box.
 pub struct GenerationalRef<R> {
     pub(crate) inner: R,
-    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-    pub(crate) borrow: GenerationalRefBorrowInfo,
+    guard: GenerationalRefBorrowGuard,
 }
 
-impl<T: 'static, R: Deref<Target = T>> GenerationalRef<R> {
-    pub(crate) fn new(
-        inner: R,
-        #[cfg(any(debug_assertions, feature = "debug_borrows"))] borrow: GenerationalRefBorrowInfo,
-    ) -> Self {
-        Self {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow,
+impl<T: ?Sized + 'static, R: Deref<Target = T>> GenerationalRef<R> {
+    pub(crate) fn new(inner: R, guard: GenerationalRefBorrowGuard) -> Self {
+        Self { inner, guard }
+    }
+
+    /// Map the inner value to a new type
+    pub fn map<R2, F: FnOnce(R) -> R2>(self, f: F) -> GenerationalRef<R2> {
+        GenerationalRef {
+            inner: f(self.inner),
+            guard: self.guard,
         }
     }
+
+    /// Try to map the inner value to a new type
+    pub fn try_map<R2, F: FnOnce(R) -> Option<R2>>(self, f: F) -> Option<GenerationalRef<R2>> {
+        f(self.inner).map(|inner| GenerationalRef {
+            inner,
+            guard: self.guard,
+        })
+    }
 }
 
 impl<T: ?Sized + Debug, R: Deref<Target = T>> Debug for GenerationalRef<R> {
@@ -43,43 +51,46 @@ impl<T: ?Sized + 'static, R: Deref<Target = T>> Deref for GenerationalRef<R> {
     }
 }
 
-#[cfg(any(debug_assertions, feature = "debug_borrows"))]
-/// Information about a borrow.
-pub struct GenerationalRefBorrowInfo {
+pub(crate) struct GenerationalRefBorrowGuard {
+    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
     pub(crate) borrowed_at: &'static std::panic::Location<'static>,
-    pub(crate) borrowed_from: &'static crate::MemoryLocationBorrowInfo,
-    pub(crate) created_at: &'static std::panic::Location<'static>,
+    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+    pub(crate) borrowed_from: &'static crate::entry::MemoryLocationBorrowInfo,
 }
 
 #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-impl Drop for GenerationalRefBorrowInfo {
+impl Drop for GenerationalRefBorrowGuard {
     fn drop(&mut self) {
-        self.borrowed_from
-            .borrowed_at
-            .write()
-            .retain(|location| !std::ptr::eq(*location, self.borrowed_at as *const _));
+        self.borrowed_from.drop_borrow(self.borrowed_at);
     }
 }
 
 /// A mutable reference to a value in a generational box.
 pub struct GenerationalRefMut<W> {
     pub(crate) inner: W,
-    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
     pub(crate) borrow: GenerationalRefBorrowMutGuard,
 }
 
-impl<T, R: DerefMut<Target = T>> GenerationalRefMut<R> {
-    pub(crate) fn new(
-        inner: R,
-        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-        borrow: GenerationalRefMutBorrowInfo,
-    ) -> Self {
-        Self {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow: borrow.into(),
+impl<T: ?Sized + 'static, R: DerefMut<Target = T>> GenerationalRefMut<R> {
+    pub(crate) fn new(inner: R, borrow: GenerationalRefBorrowMutGuard) -> Self {
+        Self { inner, borrow }
+    }
+
+    /// Map the inner value to a new type
+    pub fn map<R2, F: FnOnce(R) -> R2>(self, f: F) -> GenerationalRefMut<R2> {
+        GenerationalRefMut {
+            inner: f(self.inner),
+            borrow: self.borrow,
         }
     }
+
+    /// Try to map the inner value to a new type
+    pub fn try_map<R2, F: FnOnce(R) -> Option<R2>>(self, f: F) -> Option<GenerationalRefMut<R2>> {
+        f(self.inner).map(|inner| GenerationalRefMut {
+            inner,
+            borrow: self.borrow,
+        })
+    }
 }
 
 impl<T: ?Sized, W: DerefMut<Target = T>> Deref for GenerationalRefMut<W> {
@@ -96,33 +107,17 @@ impl<T: ?Sized, W: DerefMut<Target = T>> DerefMut for GenerationalRefMut<W> {
     }
 }
 
-#[cfg(any(debug_assertions, feature = "debug_borrows"))]
-/// Information about a mutable borrow.
-pub struct GenerationalRefMutBorrowInfo {
-    /// The location where the borrow occurred.
-    pub(crate) borrowed_from: &'static crate::MemoryLocationBorrowInfo,
-    pub(crate) created_at: &'static std::panic::Location<'static>,
-}
-
-#[cfg(any(debug_assertions, feature = "debug_borrows"))]
 pub(crate) struct GenerationalRefBorrowMutGuard {
-    borrow_info: GenerationalRefMutBorrowInfo,
-}
-
-#[cfg(any(debug_assertions, feature = "debug_borrows"))]
-impl From<GenerationalRefMutBorrowInfo> for GenerationalRefBorrowMutGuard {
-    fn from(borrow_info: GenerationalRefMutBorrowInfo) -> Self {
-        Self { borrow_info }
-    }
+    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+    /// The location where the borrow occurred.
+    pub(crate) borrowed_from: &'static crate::entry::MemoryLocationBorrowInfo,
+    #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+    pub(crate) borrowed_mut_at: &'static std::panic::Location<'static>,
 }
 
 #[cfg(any(debug_assertions, feature = "debug_borrows"))]
 impl Drop for GenerationalRefBorrowMutGuard {
     fn drop(&mut self) {
-        self.borrow_info
-            .borrowed_from
-            .borrowed_mut_at
-            .write()
-            .take();
+        self.borrowed_from.drop_borrow_mut(self.borrowed_mut_at);
     }
 }

packages/generational-box/src/sync.rs

@@ -4,18 +4,22 @@ use parking_lot::{
 use std::sync::{Arc, OnceLock};
 
 use crate::{
+    entry::{MemoryLocationBorrowInfo, StorageEntry},
     error::{self, ValueDroppedError},
     references::{GenerationalRef, GenerationalRefMut},
-    AnyStorage, MemoryLocation, MemoryLocationInner, Storage,
+    AnyStorage, GenerationalLocation, GenerationalPointer, Storage,
 };
 
 /// 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(RwLock<Option<Box<dyn std::any::Any + Send + Sync>>>);
+pub struct SyncStorage {
+    borrow_info: MemoryLocationBorrowInfo,
+    data: RwLock<StorageEntry<Box<dyn std::any::Any + Send + Sync>>>,
+}
 
-static SYNC_RUNTIME: OnceLock<Arc<Mutex<Vec<MemoryLocation<SyncStorage>>>>> = OnceLock::new();
+static SYNC_RUNTIME: OnceLock<Arc<Mutex<Vec<&'static SyncStorage>>>> = OnceLock::new();
 
-fn sync_runtime() -> &'static Arc<Mutex<Vec<MemoryLocation<SyncStorage>>>> {
+fn sync_runtime() -> &'static Arc<Mutex<Vec<&'static SyncStorage>>> {
     SYNC_RUNTIME.get_or_init(|| Arc::new(Mutex::new(Vec::new())))
 }
 
@@ -35,129 +39,134 @@ impl AnyStorage for SyncStorage {
         mut_
     }
 
+    fn map<T: ?Sized + 'static, U: ?Sized + 'static>(
+        ref_: Self::Ref<'_, T>,
+        f: impl FnOnce(&T) -> &U,
+    ) -> Self::Ref<'_, U> {
+        ref_.map(|inner| MappedRwLockReadGuard::map(inner, f))
+    }
+
+    fn map_mut<T: ?Sized + 'static, U: ?Sized + 'static>(
+        mut_ref: Self::Mut<'_, T>,
+        f: impl FnOnce(&mut T) -> &mut U,
+    ) -> Self::Mut<'_, U> {
+        mut_ref.map(|inner| MappedRwLockWriteGuard::map(inner, f))
+    }
+
     fn try_map<I: ?Sized + 'static, U: ?Sized + 'static>(
         ref_: Self::Ref<'_, I>,
         f: impl FnOnce(&I) -> Option<&U>,
     ) -> Option<Self::Ref<'_, U>> {
-        let GenerationalRef {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow,
-            ..
-        } = ref_;
-        MappedRwLockReadGuard::try_map(inner, f)
-            .ok()
-            .map(|inner| GenerationalRef {
-                inner,
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                borrow: crate::GenerationalRefBorrowInfo {
-                    borrowed_at: borrow.borrowed_at,
-                    borrowed_from: borrow.borrowed_from,
-                    created_at: borrow.created_at,
-                },
-            })
+        ref_.try_map(|inner| MappedRwLockReadGuard::try_map(inner, f).ok())
     }
 
     fn try_map_mut<I: ?Sized + 'static, U: ?Sized + 'static>(
         mut_ref: Self::Mut<'_, I>,
         f: impl FnOnce(&mut I) -> Option<&mut U>,
     ) -> Option<Self::Mut<'_, U>> {
-        let GenerationalRefMut {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow,
-            ..
-        } = mut_ref;
-        MappedRwLockWriteGuard::try_map(inner, f)
-            .ok()
-            .map(|inner| GenerationalRefMut {
-                inner,
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                borrow,
-            })
+        mut_ref.try_map(|inner| MappedRwLockWriteGuard::try_map(inner, f).ok())
     }
 
     fn data_ptr(&self) -> *const () {
-        self.0.data_ptr() as *const ()
+        self.data.data_ptr() as *const ()
     }
 
-    fn manually_drop(&self) -> bool {
-        self.0.write().take().is_some()
-    }
-
-    fn claim() -> MemoryLocation<Self> {
-        sync_runtime().lock().pop().unwrap_or_else(|| {
-            let data: &'static MemoryLocationInner<Self> =
-                &*Box::leak(Box::new(MemoryLocationInner {
-                    data: Self::default(),
-                    #[cfg(any(debug_assertions, feature = "check_generation"))]
-                    generation: 0.into(),
+    #[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"))]
-                    borrow: Default::default(),
-                }));
-            MemoryLocation(data)
-        })
+                    created_at: caller,
+                };
+
+                GenerationalPointer { storage, location }
+            }
+        }
     }
 
-    fn recycle(location: &MemoryLocation<Self>) {
-        let location = *location;
-        location.drop();
-        sync_runtime().lock().push(location);
+    fn recycle(pointer: GenerationalPointer<Self>) -> Option<Box<dyn std::any::Any>> {
+        let mut borrow_mut = pointer.storage.data.write();
+        // First check if the generation is still valid
+        if !borrow_mut.valid(&pointer.location) {
+            return None;
+        }
+        borrow_mut.increment_generation();
+        let old_data = borrow_mut.data.take();
+        sync_runtime().lock().push(pointer.storage);
+        old_data.map(|data| data as Box<dyn std::any::Any>)
     }
 }
 
 impl<T: Sync + Send + 'static> Storage<T> for SyncStorage {
+    #[track_caller]
     fn try_read(
-        &'static self,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        at: crate::GenerationalRefBorrowInfo,
+        pointer: GenerationalPointer<Self>,
     ) -> Result<Self::Ref<'static, T>, error::BorrowError> {
-        let read = self.0.read();
-
-        RwLockReadGuard::try_map(read, |any| any.as_ref()?.downcast_ref())
-            .map_err(|_| {
-                error::BorrowError::Dropped(ValueDroppedError {
-                    #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                    created_at: at.created_at,
-                })
-            })
-            .map(|guard| {
-                GenerationalRef::new(
-                    guard,
-                    #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                    at,
-                )
-            })
+        let read = pointer.storage.data.read();
+
+        let read = RwLockReadGuard::try_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()
+        });
+        match read {
+            Ok(guard) => Ok(GenerationalRef::new(
+                guard,
+                pointer.storage.borrow_info.borrow_guard(),
+            )),
+            Err(_) => Err(error::BorrowError::Dropped(
+                ValueDroppedError::new_for_location(pointer.location),
+            )),
+        }
     }
 
+    #[track_caller]
     fn try_write(
-        &'static self,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        at: crate::GenerationalRefMutBorrowInfo,
+        pointer: GenerationalPointer<Self>,
     ) -> Result<Self::Mut<'static, T>, error::BorrowMutError> {
-        let write = self.0.write();
-
-        match RwLockWriteGuard::try_map(write, |any| any.as_mut()?.downcast_mut()) {
+        let write = pointer.storage.data.write();
+
+        let write = RwLockWriteGuard::try_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()
+        });
+        match write {
             Ok(guard) => Ok(GenerationalRefMut::new(
                 guard,
-                #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                at,
+                pointer.storage.borrow_info.borrow_mut_guard(),
+            )),
+            Err(_) => Err(error::BorrowMutError::Dropped(
+                ValueDroppedError::new_for_location(pointer.location),
             )),
-            Err(_) => Err(error::BorrowMutError::Dropped(ValueDroppedError {
-                #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                created_at: at.created_at,
-            })),
         }
     }
 
-    fn set(&self, value: T) {
-        *self.0.write() = Some(Box::new(value));
-    }
-
-    fn take(&'static self) -> Option<T> {
-        self.0
-            .write()
-            .take()
-            .and_then(|any| any.downcast().ok().map(|boxed| *boxed))
+    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;
+        }
+        write.data = Some(Box::new(value));
     }
 }

packages/generational-box/src/unsync.rs

@@ -1,89 +1,38 @@
 use crate::{
+    entry::{MemoryLocationBorrowInfo, StorageEntry},
     error,
     references::{GenerationalRef, GenerationalRefMut},
-    AnyStorage, MemoryLocation, MemoryLocationInner, Storage,
+    AnyStorage, BorrowError, BorrowMutError, GenerationalLocation, GenerationalPointer, Storage,
 };
 use std::cell::{Ref, RefCell, RefMut};
 
+thread_local! {
+    static UNSYNC_RUNTIME: RefCell<Vec<&'static UnsyncStorage>> = const { RefCell::new(Vec::new()) };
+}
+
 /// A unsync storage. This is the default storage type.
 #[derive(Default)]
-pub struct UnsyncStorage(RefCell<Option<Box<dyn std::any::Any>>>);
-
-impl<T: 'static> Storage<T> for UnsyncStorage {
-    fn try_read(
-        &'static self,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        at: crate::GenerationalRefBorrowInfo,
-    ) -> Result<Self::Ref<'static, T>, error::BorrowError> {
-        let borrow = self.0.try_borrow();
-
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        let borrow = borrow.map_err(|_| at.borrowed_from.borrow_error())?;
-
-        #[cfg(not(any(debug_assertions, feature = "debug_ownership")))]
-        let borrow = borrow.map_err(|_| {
-            error::BorrowError::AlreadyBorrowedMut(error::AlreadyBorrowedMutError {})
-        })?;
-
-        match Ref::filter_map(borrow, |any| any.as_ref()?.downcast_ref()) {
-            Ok(guard) => Ok(GenerationalRef::new(
-                guard,
-                #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                at,
-            )),
-            Err(_) => Err(error::BorrowError::Dropped(error::ValueDroppedError {
-                #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                created_at: at.created_at,
-            })),
-        }
-    }
-
-    fn try_write(
-        &'static self,
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        at: crate::GenerationalRefMutBorrowInfo,
-    ) -> Result<Self::Mut<'static, T>, error::BorrowMutError> {
-        let borrow = self.0.try_borrow_mut();
-
-        #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-        let borrow = borrow.map_err(|_| at.borrowed_from.borrow_mut_error())?;
-
-        #[cfg(not(any(debug_assertions, feature = "debug_ownership")))]
-        let borrow = borrow
-            .map_err(|_| error::BorrowMutError::AlreadyBorrowed(error::AlreadyBorrowedError {}))?;
-
-        RefMut::filter_map(borrow, |any| any.as_mut()?.downcast_mut())
-            .map_err(|_| {
-                error::BorrowMutError::Dropped(error::ValueDroppedError {
-                    #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                    created_at: at.created_at,
-                })
-            })
-            .map(|guard| {
-                GenerationalRefMut::new(
-                    guard,
-                    #[cfg(any(debug_assertions, feature = "debug_ownership"))]
-                    at,
-                )
-            })
-    }
+pub struct UnsyncStorage {
+    borrow_info: MemoryLocationBorrowInfo,
+    data: RefCell<StorageEntry<Box<dyn std::any::Any>>>,
+}
 
-    fn set(&self, value: T) {
-        *self.0.borrow_mut() = Some(Box::new(value));
+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())
     }
 
-    fn take(&'static self) -> Option<T> {
-        self.0
-            .borrow_mut()
-            .take()
-            .map(|any| *any.downcast().unwrap())
+    fn try_borrow(&self) -> Result<Ref<'_, StorageEntry<Box<dyn std::any::Any>>>, BorrowError> {
+        self.data
+            .try_borrow()
+            .map_err(|_| self.borrow_info.borrow_error())
     }
 }
 
-thread_local! {
-    static UNSYNC_RUNTIME: RefCell<Vec<MemoryLocation<UnsyncStorage>>> = const { RefCell::new(Vec::new()) };
-}
-
 impl AnyStorage for UnsyncStorage {
     type Ref<'a, R: ?Sized + 'static> = GenerationalRef<Ref<'a, R>>;
     type Mut<'a, W: ?Sized + 'static> = GenerationalRefMut<RefMut<'a, W>>;
@@ -100,71 +49,136 @@ impl AnyStorage for UnsyncStorage {
         mut_
     }
 
+    fn map<T: ?Sized + 'static, U: ?Sized + 'static>(
+        ref_: Self::Ref<'_, T>,
+        f: impl FnOnce(&T) -> &U,
+    ) -> Self::Ref<'_, U> {
+        ref_.map(|inner| Ref::map(inner, f))
+    }
+
+    fn map_mut<T: ?Sized + 'static, U: ?Sized + 'static>(
+        mut_ref: Self::Mut<'_, T>,
+        f: impl FnOnce(&mut T) -> &mut U,
+    ) -> Self::Mut<'_, U> {
+        mut_ref.map(|inner| RefMut::map(inner, f))
+    }
+
     fn try_map<I: ?Sized + 'static, U: ?Sized + 'static>(
         _self: Self::Ref<'_, I>,
         f: impl FnOnce(&I) -> Option<&U>,
     ) -> Option<Self::Ref<'_, U>> {
-        let GenerationalRef {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow,
-            ..
-        } = _self;
-        Ref::filter_map(inner, f).ok().map(|inner| GenerationalRef {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow,
-        })
+        _self.try_map(|inner| Ref::filter_map(inner, f).ok())
     }
 
     fn try_map_mut<I: ?Sized + 'static, U: ?Sized + 'static>(
         mut_ref: Self::Mut<'_, I>,
         f: impl FnOnce(&mut I) -> Option<&mut U>,
     ) -> Option<Self::Mut<'_, U>> {
-        let GenerationalRefMut {
-            inner,
-            #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-            borrow,
-            ..
-        } = mut_ref;
-        RefMut::filter_map(inner, f)
-            .ok()
-            .map(|inner| GenerationalRefMut {
-                inner,
-                #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                borrow,
-            })
+        mut_ref.try_map(|inner| RefMut::filter_map(inner, f).ok())
     }
 
     fn data_ptr(&self) -> *const () {
-        self.0.as_ptr() as *const ()
-    }
-
-    fn manually_drop(&self) -> bool {
-        self.0.borrow_mut().take().is_some()
+        self.data.as_ptr() as *const ()
     }
 
-    fn claim() -> MemoryLocation<Self> {
+    #[allow(unused)]
+    fn claim(caller: &'static std::panic::Location<'static>) -> GenerationalPointer<Self> {
         UNSYNC_RUNTIME.with(|runtime| {
-            if let Some(location) = runtime.borrow_mut().pop() {
-                location
+            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 MemoryLocationInner =
-                    &*Box::leak(Box::new(MemoryLocationInner {
-                        data: Self::default(),
-                        #[cfg(any(debug_assertions, feature = "check_generation"))]
-                        generation: 0.into(),
-                        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
-                        borrow: Default::default(),
-                    }));
-                MemoryLocation(data)
+                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(location: &MemoryLocation<Self>) {
-        let location = *location;
-        location.drop();
-        UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(location));
+    fn recycle(pointer: GenerationalPointer<Self>) -> Option<Box<dyn std::any::Any>> {
+        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;
+        }
+
+        borrow_mut.increment_generation();
+        let old_data = borrow_mut.data.take();
+        UNSYNC_RUNTIME.with(|runtime| runtime.borrow_mut().push(pointer.storage));
+
+        old_data
+    }
+}
+
+impl<T: 'static> Storage<T> for UnsyncStorage {
+    #[track_caller]
+    fn try_read(
+        pointer: GenerationalPointer<Self>,
+    ) -> Result<Self::Ref<'static, T>, error::BorrowError> {
+        let read = pointer.storage.try_borrow()?;
+
+        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()
+        });
+        match ref_ {
+            Ok(guard) => Ok(GenerationalRef::new(
+                guard,
+                pointer.storage.borrow_info.borrow_guard(),
+            )),
+            Err(_) => Err(error::BorrowError::Dropped(
+                error::ValueDroppedError::new_for_location(pointer.location),
+            )),
+        }
+    }
+
+    #[track_caller]
+    fn try_write(
+        pointer: GenerationalPointer<Self>,
+    ) -> Result<Self::Mut<'static, T>, error::BorrowMutError> {
+        let write = pointer.storage.try_borrow_mut()?;
+
+        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()
+        });
+        match ref_mut {
+            Ok(guard) => Ok(GenerationalRefMut::new(
+                guard,
+                pointer.storage.borrow_info.borrow_mut_guard(),
+            )),
+            Err(_) => Err(error::BorrowMutError::Dropped(
+                error::ValueDroppedError::new_for_location(pointer.location),
+            )),
+        }
+    }
+
+    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;
+        }
+        borrow_mut.data = Some(Box::new(value));
     }
 }

packages/generational-box/tests/basic.rs

@@ -1,4 +1,4 @@
-use generational_box::{AnyStorage, GenerationalBox, UnsyncStorage};
+use generational_box::{GenerationalBox, Storage, SyncStorage, UnsyncStorage};
 
 /// # Example
 ///
@@ -14,88 +14,119 @@ fn compile_fail() {}
 
 #[test]
 fn reused() {
-    let first_ptr;
-    {
-        let owner = UnsyncStorage::owner();
-        first_ptr = owner.insert(1).raw_ptr();
-        drop(owner);
-    }
-    {
-        let owner = UnsyncStorage::owner();
-        let second_ptr = owner.insert(1234).raw_ptr();
-        assert_eq!(first_ptr, second_ptr);
-        drop(owner);
+    fn reused_test<S: Storage<i32>>() {
+        let first_ptr;
+        {
+            let owner = S::owner();
+            first_ptr = owner.insert(1).raw_ptr();
+            drop(owner);
+        }
+        {
+            let owner = S::owner();
+            let second_ptr = owner.insert(1234).raw_ptr();
+            assert_eq!(first_ptr, second_ptr);
+            drop(owner);
+        }
     }
+
+    reused_test::<UnsyncStorage>();
+    reused_test::<SyncStorage>();
 }
 
 #[test]
 fn leaking_is_ok() {
-    let data = String::from("hello world");
-    let key;
-    {
-        // create an owner
-        let owner = UnsyncStorage::owner();
-        // insert data into the store
-        key = owner.insert(data);
-        // don't drop the owner
-        std::mem::forget(owner);
+    fn leaking_is_ok_test<S: Storage<String>>() {
+        let data = String::from("hello world");
+        let key;
+        {
+            // create an owner
+            let owner = S::owner();
+            // insert data into the store
+            key = owner.insert(data);
+            // don't drop the owner
+            std::mem::forget(owner);
+        }
+        assert_eq!(
+            key.try_read().as_deref().unwrap(),
+            &"hello world".to_string()
+        );
     }
-    assert_eq!(
-        key.try_read().as_deref().unwrap(),
-        &"hello world".to_string()
-    );
+
+    leaking_is_ok_test::<UnsyncStorage>();
+    leaking_is_ok_test::<SyncStorage>();
 }
 
 #[test]
 fn drops() {
-    let data = String::from("hello world");
-    let key;
-    {
-        // create an owner
-        let owner = UnsyncStorage::owner();
-        // insert data into the store
-        key = owner.insert(data);
-        // drop the owner
+    fn drops_test<S: Storage<String>>() {
+        let data = String::from("hello world");
+        let key;
+        {
+            // create an owner
+            let owner = S::owner();
+            // insert data into the store
+            key = owner.insert(data);
+            // drop the owner
+        }
+        assert!(key.try_read().is_err());
     }
-    assert!(key.try_read().is_err());
+
+    drops_test::<UnsyncStorage>();
+    drops_test::<SyncStorage>();
 }
 
 #[test]
 fn works() {
-    let owner = UnsyncStorage::owner();
-    let key = owner.insert(1);
+    fn works_test<S: Storage<i32>>() {
+        let owner = S::owner();
+        let key = owner.insert(1);
 
-    assert_eq!(*key.read(), 1);
+        assert_eq!(*key.read(), 1);
+    }
+
+    works_test::<UnsyncStorage>();
+    works_test::<SyncStorage>();
 }
 
 #[test]
 fn insert_while_reading() {
-    let owner = UnsyncStorage::owner();
-    let key;
-    {
-        let data: String = "hello world".to_string();
-        key = owner.insert(data);
+    fn insert_while_reading_test<S: Storage<String> + Storage<&'static i32>>() {
+        let owner = S::owner();
+        let key;
+        {
+            let data: String = "hello world".to_string();
+            key = owner.insert(data);
+        }
+        let value = key.read();
+        owner.insert(&1);
+        assert_eq!(*value, "hello world");
     }
-    let value = key.read();
-    owner.insert(&1);
-    assert_eq!(*value, "hello world");
+
+    insert_while_reading_test::<UnsyncStorage>();
+    insert_while_reading_test::<SyncStorage>();
 }
 
 #[test]
 #[should_panic]
 fn panics() {
-    let owner = UnsyncStorage::owner();
-    let key = owner.insert(1);
-    drop(owner);
+    fn panics_test<S: Storage<i32>>() {
+        let owner = S::owner();
+
+        let key = owner.insert(1);
+        drop(owner);
 
-    assert_eq!(*key.read(), 1);
+        assert_eq!(*key.read(), 1);
+    }
+
+    panics_test::<UnsyncStorage>();
+    panics_test::<SyncStorage>();
 }
 
 #[test]
 fn fuzz() {
-    fn maybe_owner_scope(
-        valid_keys: &mut Vec<GenerationalBox<String>>,
-        invalid_keys: &mut Vec<GenerationalBox<String>>,
+    fn maybe_owner_scope<S: Storage<String>>(
+        valid_keys: &mut Vec<GenerationalBox<String, S>>,
+        invalid_keys: &mut Vec<GenerationalBox<String, S>>,
         path: &mut Vec<u8>,
     ) {
         let branch_cutoff = 5;
@@ -106,28 +137,36 @@ fn fuzz() {
         };
 
         for i in 0..children {
-            let owner = UnsyncStorage::owner();
+            let owner = S::owner();
             let key = owner.insert(format!("hello world {path:?}"));
+            println!("created new box {key:?}");
             valid_keys.push(key);
             path.push(i);
             // read all keys
             println!("{:?}", path);
             for key in valid_keys.iter() {
+                println!("reading {key:?}");
                 let value = key.read();
                 println!("{:?}", &*value);
                 assert!(value.starts_with("hello world"));
             }
-            #[cfg(any(debug_assertions, feature = "check_generation"))]
             for key in invalid_keys.iter() {
+                println!("reading invalid {key:?}");
                 assert!(key.try_read().is_err());
             }
             maybe_owner_scope(valid_keys, invalid_keys, path);
-            invalid_keys.push(valid_keys.pop().unwrap());
+            let invalid = valid_keys.pop().unwrap();
+            println!("popping {invalid:?}");
+            invalid_keys.push(invalid);
             path.pop();
         }
     }
 
     for _ in 0..10 {
-        maybe_owner_scope(&mut Vec::new(), &mut Vec::new(), &mut Vec::new());
+        maybe_owner_scope::<UnsyncStorage>(&mut Vec::new(), &mut Vec::new(), &mut Vec::new());
+    }
+
+    for _ in 0..10 {
+        maybe_owner_scope::<SyncStorage>(&mut Vec::new(), &mut Vec::new(), &mut Vec::new());
     }
 }

packages/generational-box/tests/errors.rs

@@ -1,38 +1,33 @@
-use std::cell::{Ref, RefMut};
-
 use generational_box::{
-    AlreadyBorrowedError, AlreadyBorrowedMutError, AnyStorage, BorrowError, BorrowMutError,
-    GenerationalBox, GenerationalRef, GenerationalRefMut, Owner, UnsyncStorage, ValueDroppedError,
+    AlreadyBorrowedError, AlreadyBorrowedMutError, BorrowError, BorrowMutError, GenerationalBox,
+    Owner, Storage, SyncStorage, UnsyncStorage, ValueDroppedError,
 };
 
 #[track_caller]
-fn read_at_location(
-    value: GenerationalBox<i32>,
-) -> (
-    GenerationalRef<Ref<'static, i32>>,
-    &'static std::panic::Location<'static>,
-) {
+fn read_at_location<S: Storage<i32>>(
+    value: GenerationalBox<i32, S>,
+) -> (S::Ref<'static, i32>, &'static std::panic::Location<'static>) {
     let location = std::panic::Location::caller();
     let read = value.read();
     (read, location)
 }
 
 #[track_caller]
-fn write_at_location(
-    value: GenerationalBox<i32>,
-) -> (
-    GenerationalRefMut<RefMut<'static, i32>>,
-    &'static std::panic::Location<'static>,
-) {
+fn write_at_location<S: Storage<i32>>(
+    value: GenerationalBox<i32, S>,
+) -> (S::Mut<'static, i32>, &'static std::panic::Location<'static>) {
     let location = std::panic::Location::caller();
     let write = value.write();
     (write, location)
 }
 
 #[track_caller]
-fn create_at_location(
-    owner: &Owner,
-) -> (GenerationalBox<i32>, &'static std::panic::Location<'static>) {
+fn create_at_location<S: Storage<i32>>(
+    owner: &Owner<S>,
+) -> (
+    GenerationalBox<i32, S>,
+    &'static std::panic::Location<'static>,
+) {
     let location = std::panic::Location::caller();
     let value = owner.insert(1);
     (value, location)
@@ -40,65 +35,85 @@ fn create_at_location(
 
 #[test]
 fn read_while_writing_error() {
-    let owner = UnsyncStorage::owner();
-    let value = owner.insert(1);
-
-    let (write, location) = write_at_location(value);
-
-    assert_eq!(
-        value.try_read().err(),
-        Some(BorrowError::AlreadyBorrowedMut(
-            AlreadyBorrowedMutError::new(location)
-        ))
-    );
-    drop(write);
+    fn read_while_writing_error_test<S: Storage<i32>>() {
+        let owner = S::owner();
+        let value = owner.insert(1);
+
+        let (write, location) = write_at_location(value);
+
+        assert_eq!(
+            value.try_read().err(),
+            Some(BorrowError::AlreadyBorrowedMut(
+                AlreadyBorrowedMutError::new(location)
+            ))
+        );
+        drop(write);
+    }
+
+    // For sync storage this will deadlock
+    read_while_writing_error_test::<UnsyncStorage>();
 }
 
 #[test]
 fn read_after_dropped_error() {
-    let owner = UnsyncStorage::owner();
-    let (value, location) = create_at_location(&owner);
-    drop(owner);
-    assert_eq!(
-        value.try_read().err(),
-        Some(BorrowError::Dropped(ValueDroppedError::new(location)))
-    );
+    fn read_after_dropped_error_test<S: Storage<i32>>() {
+        let owner = S::owner();
+        let (value, location) = create_at_location(&owner);
+        drop(owner);
+        assert_eq!(
+            value.try_read().err(),
+            Some(BorrowError::Dropped(ValueDroppedError::new(location)))
+        );
+    }
+
+    read_after_dropped_error_test::<UnsyncStorage>();
+    read_after_dropped_error_test::<SyncStorage>();
 }
 
 #[test]
 fn write_while_writing_error() {
-    let owner = UnsyncStorage::owner();
-    let value = owner.insert(1);
-
-    #[allow(unused)]
-    let (write, location) = write_at_location(value);
-
-    let write_result = value.try_write();
-    assert!(write_result.is_err());
-    #[cfg(debug_assertions)]
-    assert_eq!(
-        write_result.err(),
-        Some(BorrowMutError::AlreadyBorrowedMut(
-            AlreadyBorrowedMutError::new(location)
-        ))
-    );
-
-    drop(write);
+    fn write_while_writing_error_test<S: Storage<i32>>() {
+        let owner = S::owner();
+        let value = owner.insert(1);
+
+        #[allow(unused)]
+        let (write, location) = write_at_location(value);
+
+        let write_result = value.try_write();
+        assert!(write_result.is_err());
+        #[cfg(any(debug_assertions, feature = "debug_borrows"))]
+        assert_eq!(
+            write_result.err(),
+            Some(BorrowMutError::AlreadyBorrowedMut(
+                AlreadyBorrowedMutError::new(location)
+            ))
+        );
+
+        drop(write);
+    }
+
+    // For sync storage this will deadlock
+    write_while_writing_error_test::<UnsyncStorage>();
 }
 
 #[test]
 fn write_while_reading_error() {
-    let owner = UnsyncStorage::owner();
-    let value = owner.insert(1);
+    fn write_while_reading_error_test<S: Storage<i32>>() {
+        let owner = S::owner();
+        let value = owner.insert(1);
+
+        let (read, location) = read_at_location(value);
 
-    let (read, location) = read_at_location(value);
+        assert_eq!(
+            value.try_write().err(),
+            Some(BorrowMutError::AlreadyBorrowed(AlreadyBorrowedError::new(
+                vec![location]
+            )))
+        );
 
-    assert_eq!(
-        value.try_write().err(),
-        Some(BorrowMutError::AlreadyBorrowed(AlreadyBorrowedError::new(
-            vec![location]
-        )))
-    );
+        drop(read);
+    }
 
-    drop(read);
+    // For sync storage this will deadlock
+    write_while_reading_error_test::<UnsyncStorage>();
 }

packages/generational-box/tests/sync.rs

@@ -0,0 +1,42 @@
+// Regression test for https://github.com/DioxusLabs/dioxus/issues/2636
+
+use std::time::Duration;
+
+use generational_box::{AnyStorage, GenerationalBox, SyncStorage};
+
+#[test]
+fn race_condition_regression() {
+    for _ in 0..100 {
+        let handle = {
+            let owner = SyncStorage::owner();
+            let key = owner.insert(1u64);
+            let handle = std::thread::spawn(move || reader(key));
+
+            std::thread::sleep(Duration::from_millis(10));
+            handle
+            // owner is dropped now
+        };
+        // owner is *recycled*
+        let owner = SyncStorage::owner();
+        let _key = owner.insert(2u64);
+        let _ = handle.join();
+    }
+}
+
+fn reader(key: GenerationalBox<u64, SyncStorage>) {
+    for _ in 0..1000000 {
+        match key.try_read() {
+            Ok(value) => {
+                if *value == 2 {
+                    panic!("Read a new value with the old generation");
+                } else {
+                    // fine
+                }
+            }
+            Err(err) => {
+                eprintln!("bailing out - {err:?}");
+                break;
+            }
+        }
+    }
+}

packages/signals/src/copy_value.rs

@@ -75,16 +75,12 @@ impl<T: 'static, S: Storage<T>> CopyValue<T, S> {
 
     pub(crate) fn new_with_caller(
         value: T,
-        #[cfg(debug_assertions)] caller: &'static std::panic::Location<'static>,
+        caller: &'static std::panic::Location<'static>,
     ) -> Self {
         let owner = current_owner();
 
         Self {
-            value: owner.insert_with_caller(
-                value,
-                #[cfg(debug_assertions)]
-                caller,
-            ),
+            value: owner.insert_with_caller(value, caller),
             origin_scope: current_scope_id().expect("in a virtual dom"),
         }
     }

packages/signals/src/signal.rs

@@ -167,7 +167,6 @@ impl<T: 'static, S: Storage<SignalData<T>>> Signal<T, S> {
     ///     use_hook(move || Signal::new_with_caller(function(), caller))
     /// }
     /// ```
-    #[allow(unused)]
     pub fn new_with_caller(value: T, caller: &'static std::panic::Location<'static>) -> Self {
         Self {
             inner: CopyValue::new_with_caller(
@@ -175,7 +174,6 @@ impl<T: 'static, S: Storage<SignalData<T>>> Signal<T, S> {
                     subscribers: Default::default(),
                     value,
                 },
-                #[cfg(debug_assertions)]
                 caller,
             ),
         }