use generational identifiers for tracked structs

Author: ibraheemdev

src/id.rs

@@ -1,32 +1,38 @@
 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 16-bit generation counter
+/// as well as an optional 16-bit ingredient index.
 ///
-/// 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 pack a [`PageIndex`](`crate::table::PageIndex`) and
+/// [`SlotIndex`](`crate::table::SlotIndex`) within the page.
+///
+/// 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`].
+    const INGREDIENT_MASK: u64 = 0x0000FFFFFFFFFFFF;
+    const GENERATION_MASK: u64 = 0xFFFF0000FFFFFFFF;
+
+    /// 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 +44,98 @@ impl Id {
     #[doc(hidden)]
     #[track_caller]
     #[inline]
-    pub const unsafe fn from_u32(v: u32) -> Self {
+    pub const unsafe fn from_data(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) },
         }
     }
 
+    /// 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 as_u32(self) -> u32 {
-        self.value.get() - 1
+    pub fn with_generation(self, generation: u16) -> Id {
+        let mut value = self.value.get();
+
+        value &= Id::GENERATION_MASK;
+        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) },
+        }
+    }
+
+    /// Mark the `Id` with an ingredient index.
+    #[inline]
+    pub fn with_ingredient_index(self, ingredient: u16) -> Id {
+        let mut value = self.value.get();
+
+        value &= Id::INGREDIENT_MASK;
+        value |= (ingredient as u64) << 48;
+
+        Id {
+            // SAFETY: The niche of `value` is in the lower bits, which we did not touch.
+            value: unsafe { NonZeroU64::new_unchecked(value) },
+        }
+    }
+
+    /// Return the internal `u64` representation of this `Id`.
+    #[inline]
+    pub const fn as_bits(self) -> u64 {
+        self.value.get()
+    }
+
+    /// Return the data portion of this `Id`.
+    #[inline]
+    pub const fn data(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) -> u16 {
+        ((self.value.get() & !Id::GENERATION_MASK) >> 32) as u16
+    }
+
+    /// Return the ingredient index of this `Id`.
+    #[inline]
+    pub const fn ingredient_index(self) -> u16 {
+        (self.value.get() >> 48) as u16
     }
 }
 
 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.data())
+        } else {
+            write!(f, "Id({:x}g{:x})", self.data(), 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,26 @@ use crate::{Database, Id};
 /// only for inserting into maps and the like.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct DatabaseKeyIndex {
-    ingredient_index: IngredientIndex,
-    key_index: Id,
+    packed_key_index: Id,
 }
 // ANCHOR_END: DatabaseKeyIndex
 
 impl DatabaseKeyIndex {
     #[inline]
     pub(crate) fn new(ingredient_index: IngredientIndex, key_index: Id) -> Self {
         Self {
-            key_index,
-            ingredient_index,
+            // Pack the ingredient index with the key.
+            packed_key_index: key_index.with_ingredient_index(ingredient_index.as_u16()),
         }
     }
 
     pub fn ingredient_index(self) -> IngredientIndex {
-        self.ingredient_index
+        IngredientIndex::from(self.packed_key_index.ingredient_index() as usize)
     }
 
     pub fn key_index(self) -> Id {
-        self.key_index
+        // Clear the ingredient index from the key as it not part of the dependency.
+        self.packed_key_index.with_ingredient_index(0)
     }
 
     pub(crate) fn maybe_changed_after(
@@ -44,15 +44,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 +61,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_data((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 data = id.data() as usize;
+    let slot = data & PAGE_LEN_MASK;
+    let page = data >> PAGE_LEN_BITS;
     (PageIndex::new(page), SlotIndex::new(slot))
 }