Use generational identifiers for tracked structs (#864)

* use generational identifiers for tracked structs

* increase ID generations to 32-bits

* remove `created_at` field from tracked structs

* clean up tracked struct IDs to handle overflow

* log tracing message for leaked tracked structs

Author: ibraheemdev

src/id.rs

@@ -1,32 +1,33 @@
 use std::fmt::Debug;
 use std::hash::Hash;
-use std::num::NonZeroU32;
+use std::num::NonZeroU64;
 
 use crate::zalsa::Zalsa;
 
 /// The `Id` of a salsa struct in the database [`Table`](`crate::table::Table`).
 ///
-/// The higher-order bits of an `Id` identify a [`Page`](`crate::table::Page`)
-/// and the low-order bits identify a slot within the page.
+/// The high-order bits of an `Id` store a 32-bit generation counter, while
+/// the low-order bits pack a [`PageIndex`](`crate::table::PageIndex`) and
+/// [`SlotIndex`](`crate::table::SlotIndex`) within the page.
 ///
-/// An Id is a newtype'd u32 ranging from `0..Id::MAX_U32`.
-/// The maximum range is smaller than a standard u32 to leave
+/// The low-order bits of `Id` are a `u32` ranging from `0..Id::MAX_U32`.
+/// The maximum range is smaller than a standard `u32` to leave
 /// room for niches; currently there is only one niche, so that
 /// `Option<Id>` is the same size as an `Id`.
 ///
 /// As an end-user of `Salsa` you will generally not use `Id` directly,
 /// it is wrapped in new types.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Id {
-    value: NonZeroU32,
+    value: NonZeroU64,
 }
 
 impl Id {
     pub const MAX_U32: u32 = u32::MAX - 0xFF;
     pub const MAX_USIZE: usize = Self::MAX_U32 as usize;
 
-    /// Create a `salsa::Id` from a u32 value. This value should
-    /// be less than [`Self::MAX_U32`].
+    /// Create a `salsa::Id` from a u32 value, without a generation. This
+    /// value should be less than [`Self::MAX_U32`].
     ///
     /// In general, you should not need to create salsa ids yourself,
     /// but it can be useful if you are using the type as a general
@@ -38,23 +39,90 @@ impl Id {
     #[doc(hidden)]
     #[track_caller]
     #[inline]
-    pub const unsafe fn from_u32(v: u32) -> Self {
+    pub const unsafe fn from_index(v: u32) -> Self {
         debug_assert!(v < Self::MAX_U32);
+
+        Id {
+            // SAFETY: Caller obligation.
+            value: unsafe { NonZeroU64::new_unchecked((v + 1) as u64) },
+        }
+    }
+
+    /// Create a `salsa::Id` from a u64 value.
+    ///
+    /// This should only be used to recreate an `Id` together with `Id::as_u64`.
+    ///
+    /// # Safety
+    ///
+    /// The data bits of the supplied value must represent a valid `Id` returned
+    /// by `Id::as_u64`.
+    #[doc(hidden)]
+    #[track_caller]
+    #[inline]
+    pub const unsafe fn from_bits(v: u64) -> Self {
         Id {
-            // SAFETY: Caller obligation
-            value: unsafe { NonZeroU32::new_unchecked(v + 1) },
+            // SAFETY: Caller obligation.
+            value: unsafe { NonZeroU64::new_unchecked(v) },
         }
     }
 
+    /// Returns a new `Id` with same index, but the generation incremented by one.
+    ///
+    /// Returns `None` if the generation would overflow, i.e. the current generation
+    /// is `u32::MAX`.
     #[inline]
-    pub const fn as_u32(self) -> u32 {
-        self.value.get() - 1
+    pub fn next_generation(self) -> Option<Id> {
+        self.generation()
+            .checked_add(1)
+            .map(|generation| self.with_generation(generation))
+    }
+
+    /// Mark the `Id` with a generation.
+    ///
+    /// This `Id` will refer to the same page and slot in the database,
+    /// but will differ from other identifiers of the slot based on the
+    /// provided generation.
+    #[inline]
+    pub const fn with_generation(self, generation: u32) -> Id {
+        let mut value = self.value.get();
+
+        value &= 0xFFFFFFFF;
+        value |= (generation as u64) << 32;
+
+        Id {
+            // SAFETY: The niche of `value` is in the lower bits, which we did not touch.
+            value: unsafe { NonZeroU64::new_unchecked(value) },
+        }
+    }
+
+    /// Return the index portion of this `Id`.
+    #[inline]
+    pub const fn index(self) -> u32 {
+        // Truncate the high-order bits.
+        (self.value.get() as u32) - 1
+    }
+
+    /// Return the generation of this `Id`.
+    #[inline]
+    pub const fn generation(self) -> u32 {
+        // Shift away the low-order bits.
+        (self.value.get() >> 32) as u32
+    }
+
+    /// Return the internal `u64` representation of this `Id`.
+    #[inline]
+    pub const fn as_bits(self) -> u64 {
+        self.value.get()
     }
 }
 
 impl Debug for Id {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "Id({:x})", self.as_u32())
+        if self.generation() == 0 {
+            write!(f, "Id({:x})", self.index())
+        } else {
+            write!(f, "Id({:x}g{:x})", self.index(), self.generation())
+        }
     }
 }
 

src/input/singleton.rs

@@ -1,4 +1,4 @@
-use crate::loom::sync::atomic::{AtomicU32, Ordering};
+use crate::loom::sync::atomic::{AtomicU64, Ordering};
 use crate::Id;
 
 mod sealed {
@@ -11,7 +11,7 @@ pub trait SingletonChoice: sealed::Sealed + Default {
 }
 
 pub struct Singleton {
-    index: AtomicU32,
+    index: AtomicU64,
 }
 impl sealed::Sealed for Singleton {}
 impl SingletonChoice for Singleton {
@@ -22,7 +22,7 @@ impl SingletonChoice for Singleton {
         let id = cb();
         if self
             .index
-            .compare_exchange(0, id.as_u32() + 1, Ordering::AcqRel, Ordering::Acquire)
+            .compare_exchange(0, id.as_bits(), Ordering::AcqRel, Ordering::Acquire)
             .is_err()
         {
             panic!("singleton struct may not be duplicated");
@@ -33,16 +33,17 @@ impl SingletonChoice for Singleton {
     fn index(&self) -> Option<Id> {
         match self.index.load(Ordering::Acquire) {
             0 => None,
-            // SAFETY: Our u32 is derived from an ID and thus safe to convert back.
-            id => Some(unsafe { Id::from_u32(id - 1) }),
+
+            // SAFETY: Our u64 is derived from an ID and thus safe to convert back.
+            id => Some(unsafe { Id::from_bits(id) }),
         }
     }
 }
 
 impl Default for Singleton {
     fn default() -> Self {
         Self {
-            index: AtomicU32::new(0),
+            index: AtomicU64::new(0),
         }
     }
 }

src/key.rs

@@ -12,26 +12,29 @@ use crate::{Database, Id};
 /// only for inserting into maps and the like.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct DatabaseKeyIndex {
+    key_index: u32,
+    key_generation: u32,
     ingredient_index: IngredientIndex,
-    key_index: Id,
 }
 // ANCHOR_END: DatabaseKeyIndex
 
 impl DatabaseKeyIndex {
     #[inline]
     pub(crate) fn new(ingredient_index: IngredientIndex, key_index: Id) -> Self {
         Self {
-            key_index,
+            key_index: key_index.index(),
+            key_generation: key_index.generation(),
             ingredient_index,
         }
     }
 
-    pub fn ingredient_index(self) -> IngredientIndex {
+    pub const fn ingredient_index(self) -> IngredientIndex {
         self.ingredient_index
     }
 
-    pub fn key_index(self) -> Id {
-        self.key_index
+    pub const fn key_index(self) -> Id {
+        // SAFETY: `self.key_index` was returned by `Id::data`.
+        unsafe { Id::from_index(self.key_index) }.with_generation(self.key_generation)
     }
 
     pub(crate) fn maybe_changed_after(
@@ -44,15 +47,15 @@ impl DatabaseKeyIndex {
         // SAFETY: The `db` belongs to the ingredient
         unsafe {
             zalsa
-                .lookup_ingredient(self.ingredient_index)
-                .maybe_changed_after(db, self.key_index, last_verified_at, cycle_heads)
+                .lookup_ingredient(self.ingredient_index())
+                .maybe_changed_after(db, self.key_index(), last_verified_at, cycle_heads)
         }
     }
 
     pub(crate) fn remove_stale_output(&self, zalsa: &Zalsa, executor: DatabaseKeyIndex) {
         zalsa
-            .lookup_ingredient(self.ingredient_index)
-            .remove_stale_output(zalsa, executor, self.key_index)
+            .lookup_ingredient(self.ingredient_index())
+            .remove_stale_output(zalsa, executor, self.key_index())
     }
 
     pub(crate) fn mark_validated_output(
@@ -61,21 +64,21 @@ impl DatabaseKeyIndex {
         database_key_index: DatabaseKeyIndex,
     ) {
         zalsa
-            .lookup_ingredient(self.ingredient_index)
-            .mark_validated_output(zalsa, database_key_index, self.key_index)
+            .lookup_ingredient(self.ingredient_index())
+            .mark_validated_output(zalsa, database_key_index, self.key_index())
     }
 }
 
 impl fmt::Debug for DatabaseKeyIndex {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         crate::attach::with_attached_database(|db| {
-            let ingredient = db.zalsa().lookup_ingredient(self.ingredient_index);
-            ingredient.fmt_index(self.key_index, f)
+            let ingredient = db.zalsa().lookup_ingredient(self.ingredient_index());
+            ingredient.fmt_index(self.key_index(), f)
         })
         .unwrap_or_else(|| {
             f.debug_tuple("DatabaseKeyIndex")
-                .field(&self.ingredient_index)
-                .field(&self.key_index)
+                .field(&self.ingredient_index())
+                .field(&self.key_index())
                 .finish()
         })
     }

src/table.rs

@@ -400,13 +400,13 @@ fn make_id(page: PageIndex, slot: SlotIndex) -> Id {
     let page = page.0 as u32;
     let slot = slot.0 as u32;
     // SAFETY: `slot` is guaranteed to be small enough that the resulting Id won't be bigger than `Id::MAX_U32`
-    unsafe { Id::from_u32((page << PAGE_LEN_BITS) | slot) }
+    unsafe { Id::from_index((page << PAGE_LEN_BITS) | slot) }
 }
 
 #[inline]
 fn split_id(id: Id) -> (PageIndex, SlotIndex) {
-    let id = id.as_u32() as usize;
-    let slot = id & PAGE_LEN_MASK;
-    let page = id >> PAGE_LEN_BITS;
+    let index = id.index() as usize;
+    let slot = index & PAGE_LEN_MASK;
+    let page = index >> PAGE_LEN_BITS;
     (PageIndex::new(page), SlotIndex::new(slot))
 }