optimize allocation fast-path

Author: ibraheemdev

src/input.rs

@@ -111,14 +111,16 @@ impl<C: Configuration> IngredientImpl<C> {
         durabilities: C::Durabilities,
     ) -> C::Struct {
         let id = self.singleton.with_scope(|| {
-            zalsa_local.allocate(zalsa, self.ingredient_index, |_| Value::<C> {
+            let (id, _) = zalsa_local.allocate(zalsa, self.ingredient_index, |_| Value::<C> {
                 fields,
                 revisions,
                 durabilities,
                 // SAFETY: We only ever access the memos of a value that we allocated through
                 // our `MemoTableTypes`.
                 memos: unsafe { MemoTable::new(self.memo_table_types()) },
-            })
+            });
+
+            id
         });
 
         FromIdWithDb::from_id(id, zalsa)

src/interned.rs

@@ -583,7 +583,7 @@ where
             .unwrap_or((Durability::MAX, Revision::max()));
 
         // Allocate the value slot.
-        let id = zalsa_local.allocate(zalsa, self.ingredient_index, |id| Value::<C> {
+        let (id, value) = zalsa_local.allocate(zalsa, self.ingredient_index, |id| Value::<C> {
             shard: shard_index as u16,
             link: LinkedListLink::new(),
             // SAFETY: We only ever access the memos of a value that we allocated through
@@ -598,7 +598,6 @@ where
             }),
         });
 
-        let value = zalsa.table().get::<Value<C>>(id);
         // SAFETY: We hold the lock for the shard containing the value.
         let value_shared = unsafe { &mut *value.shared.get() };
 

src/table.rs

@@ -207,7 +207,17 @@ impl Table {
     /// If `page` is out of bounds or the type `T` is incorrect.
     #[inline]
     pub(crate) fn page<T: Slot>(&self, page: PageIndex) -> PageView<'_, T> {
-        self.pages[page.0].assert_type::<T>()
+        #[cold]
+        #[inline(never)]
+        fn assert_failed(index: usize) -> ! {
+            panic!("index {index} is uninitialized");
+        }
+
+        if let Some(page) = self.pages.get(page.0) {
+            return page.assert_type();
+        }
+
+        assert_failed(page.0)
     }
 
     /// Allocate a new page for the given ingredient and with slots of type `T`
@@ -286,6 +296,8 @@ impl Table {
             .flat_map(|view| view.data())
     }
 
+    #[cold]
+    #[inline(never)]
     pub(crate) fn fetch_or_push_page<T: Slot>(
         &self,
         ingredient: IngredientIndex,
@@ -299,6 +311,7 @@ impl Table {
         {
             return page;
         }
+
         self.push_page::<T>(ingredient, memo_types())
     }
 
@@ -311,22 +324,23 @@ impl Table {
     }
 }
 
-impl<'p, T: Slot> PageView<'p, T> {
+impl<'db, T: Slot> PageView<'db, T> {
     #[inline]
-    fn page_data(&self) -> &'p [PageDataEntry<T>] {
+    fn page_data(&self) -> &'db [PageDataEntry<T>] {
         let len = self.0.allocated.load(Ordering::Acquire);
         // SAFETY: `len` is the initialized length of the page
         unsafe { slice::from_raw_parts(self.0.data.cast::<PageDataEntry<T>>().as_ptr(), len) }
     }
 
     #[inline]
-    fn data(&self) -> &'p [T] {
+    fn data(&self) -> &'db [T] {
         let len = self.0.allocated.load(Ordering::Acquire);
         // SAFETY: `len` is the initialized length of the page
         unsafe { slice::from_raw_parts(self.0.data.cast::<T>().as_ptr(), len) }
     }
 
-    pub(crate) fn allocate<V>(&self, page: PageIndex, value: V) -> Result<Id, V>
+    #[inline]
+    pub(crate) fn allocate<V>(&self, page: PageIndex, value: V) -> Result<(Id, &'db T), V>
     where
         V: FnOnce(Id) -> T,
     {
@@ -347,11 +361,14 @@ impl<'p, T: Slot> PageView<'p, T> {
         // interior
         unsafe { (*entry.get()).write(value(id)) };
 
+        // SAFETY: We just initialized the value above.
+        let value = unsafe { (*entry.get()).assume_init_ref() };
+
         // Update the length (this must be done after initialization as otherwise an uninitialized
         // read could occur!)
         self.0.allocated.store(index + 1, Ordering::Release);
 
-        Ok(id)
+        Ok((id, value))
     }
 }
 

src/tracked_struct.rs

@@ -477,7 +477,9 @@ where
             return id;
         }
 
-        zalsa_local.allocate::<Value<C>>(zalsa, self.ingredient_index, value)
+        let (id, _) = zalsa_local.allocate::<Value<C>>(zalsa, self.ingredient_index, value);
+
+        id
     }
 
     /// Get mutable access to the data for `id` -- this holds a write lock for the duration

src/zalsa_local.rs

@@ -53,12 +53,36 @@ impl ZalsaLocal {
     /// Allocate a new id in `table` for the given ingredient
     /// storing `value`. Remembers the most recent page from this
     /// thread and attempts to reuse it.
-    pub(crate) fn allocate<T: Slot>(
+    pub(crate) fn allocate<'db, T: Slot>(
         &self,
-        zalsa: &Zalsa,
+        zalsa: &'db Zalsa,
         ingredient: IngredientIndex,
         mut value: impl FnOnce(Id) -> T,
-    ) -> Id {
+    ) -> (Id, &'db T) {
+        // SAFETY: `ZalsaLocal` is `!Sync`, and we never expose a reference to this field,
+        // so we have exclusive access.
+        let most_recent_pages = unsafe { &mut *self.most_recent_pages.get() };
+
+        // Fast-path, we already have an unfilled page available.
+        if let Some(&page) = most_recent_pages.get(&ingredient) {
+            let page_ref = zalsa.table().page::<T>(page);
+            match page_ref.allocate(page, value) {
+                Ok((id, value)) => return (id, value),
+                Err(v) => value = v,
+            }
+        }
+
+        self.allocate_cold(zalsa, ingredient, value)
+    }
+
+    #[cold]
+    #[inline(never)]
+    pub(crate) fn allocate_cold<'db, T: Slot>(
+        &self,
+        zalsa: &'db Zalsa,
+        ingredient: IngredientIndex,
+        mut value: impl FnOnce(Id) -> T,
+    ) -> (Id, &'db T) {
         let memo_types = || {
             zalsa
                 .lookup_ingredient(ingredient)
@@ -82,7 +106,7 @@ impl ZalsaLocal {
             let page_ref = zalsa.table().page::<T>(page);
             match page_ref.allocate(page, value) {
                 // If successful, return
-                Ok(id) => return id,
+                Ok((id, value)) => return (id, value),
 
                 // Otherwise, create a new page and try again
                 // Note that we could try fetching a page again, but as we just filled one up