Per ingredient sync table

Veykril · Veykril · commit ddc7c9d92d1c · 2025-03-20T15:52:24.000+01:00
diff --git a/src/function.rs b/src/function.rs
@@ -3,11 +3,11 @@ use std::{any::Any, fmt, ptr::NonNull};
 use crate::{
     accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues},
     cycle::{CycleRecoveryAction, CycleRecoveryStrategy},
+    function::sync::{ClaimResult, SyncTable},
     ingredient::fmt_index,
     key::DatabaseKeyIndex,
     plumbing::MemoIngredientMap,
     salsa_struct::SalsaStructInDb,
-    table::sync::ClaimResult,
     table::Table,
     views::DatabaseDownCaster,
     zalsa::{IngredientIndex, MemoIngredientIndex, Zalsa},
@@ -32,6 +32,7 @@ mod lru;
 mod maybe_changed_after;
 mod memo;
 mod specify;
+mod sync;
 
 pub trait Configuration: Any {
     const DEBUG_NAME: &'static str;
@@ -119,6 +120,8 @@ pub struct IngredientImpl<C: Configuration> {
     /// instances that this downcaster was derived from.
     view_caster: DatabaseDownCaster<C::DbView>,
 
+    sync_table: SyncTable,
+
     /// When `fetch` and friends executes, they return a reference to the
     /// value stored in the memo that is extended to live as long as the `&self`
     /// reference we start with. This means that whenever we remove something
@@ -157,6 +160,7 @@ where
             lru: lru::Lru::new(lru),
             deleted_entries: Default::default(),
             view_caster,
+            sync_table: Default::default(),
         }
     }
 
@@ -248,12 +252,10 @@ where
     /// Attempts to claim `key_index`, returning `false` if a cycle occurs.
     fn wait_for(&self, db: &dyn Database, key_index: Id) -> bool {
         let zalsa = db.zalsa();
-        match zalsa.sync_table_for(key_index).claim(
-            db,
-            zalsa,
-            self.database_key_index(key_index),
-            self.memo_ingredient_index(zalsa, key_index),
-        ) {
+        match self
+            .sync_table
+            .try_claim(db, zalsa, self.database_key_index(key_index))
+        {
             ClaimResult::Retry | ClaimResult::Claimed(_) => true,
             ClaimResult::Cycle => false,
         }
diff --git a/src/function/fetch.rs b/src/function/fetch.rs
@@ -1,9 +1,9 @@
 use super::{memo::Memo, Configuration, IngredientImpl, VerifyResult};
+use crate::function::sync::ClaimResult;
 use crate::zalsa::MemoIngredientIndex;
 use crate::{
     accumulator::accumulated_map::InputAccumulatedValues,
     runtime::StampedValue,
-    table::sync::ClaimResult,
     zalsa::{Zalsa, ZalsaDatabase},
     zalsa_local::QueryRevisions,
     AsDynDatabase as _, Id,
@@ -107,12 +107,7 @@ where
         let database_key_index = self.database_key_index(id);
 
         // Try to claim this query: if someone else has claimed it already, go back and start again.
-        let _claim_guard = match zalsa.sync_table_for(id).claim(
-            db,
-            zalsa,
-            database_key_index,
-            memo_ingredient_index,
-        ) {
+        let _claim_guard = match self.sync_table.try_claim(db, zalsa, database_key_index) {
             ClaimResult::Retry => return None,
             ClaimResult::Cycle => {
                 // check if there's a provisional value for this query
diff --git a/src/function/maybe_changed_after.rs b/src/function/maybe_changed_after.rs
@@ -1,8 +1,8 @@
 use crate::{
     accumulator::accumulated_map::InputAccumulatedValues,
     cycle::{CycleHeads, CycleRecoveryStrategy},
+    function::sync::ClaimResult,
     key::DatabaseKeyIndex,
-    table::sync::ClaimResult,
     zalsa::{MemoIngredientIndex, Zalsa, ZalsaDatabase},
     zalsa_local::{ActiveQueryGuard, QueryEdge, QueryOrigin},
     AsDynDatabase as _, Id, Revision,
@@ -98,12 +98,7 @@ where
     ) -> Option<VerifyResult> {
         let database_key_index = self.database_key_index(key_index);
 
-        let _claim_guard = match zalsa.sync_table_for(key_index).claim(
-            db,
-            zalsa,
-            database_key_index,
-            memo_ingredient_index,
-        ) {
+        let _claim_guard = match self.sync_table.try_claim(db, zalsa, database_key_index) {
             ClaimResult::Retry => return None,
             ClaimResult::Cycle => match C::CYCLE_STRATEGY {
                 CycleRecoveryStrategy::Panic => panic!(
diff --git a/src/function/sync.rs b/src/function/sync.rs
@@ -1,21 +1,26 @@
 use std::thread::ThreadId;
 
 use parking_lot::Mutex;
+use rustc_hash::FxHashMap;
 
 use crate::{
     key::DatabaseKeyIndex,
     runtime::{BlockResult, WaitResult},
-    zalsa::{MemoIngredientIndex, Zalsa},
-    Database,
+    zalsa::Zalsa,
+    Database, Id,
 };
 
-use super::util;
-
 /// Tracks the keys that are currently being processed; used to coordinate between
 /// worker threads.
 #[derive(Default)]
 pub(crate) struct SyncTable {
-    syncs: Mutex<Vec<Option<SyncState>>>,
+    syncs: Mutex<FxHashMap<Id, SyncState>>,
+}
+
+pub(crate) enum ClaimResult<'a> {
+    Retry,
+    Cycle,
+    Claimed(ClaimGuard<'a>),
 }
 
 struct SyncState {
@@ -26,44 +31,20 @@ struct SyncState {
     anyone_waiting: bool,
 }
 
-pub(crate) enum ClaimResult<'a> {
-    Retry,
-    Cycle,
-    Claimed(ClaimGuard<'a>),
-}
-
 impl SyncTable {
-    #[inline]
-    pub(crate) fn claim<'me>(
+    pub(crate) fn try_claim<'me>(
         &'me self,
         db: &'me (impl ?Sized + Database),
         zalsa: &'me Zalsa,
         database_key_index: DatabaseKeyIndex,
-        memo_ingredient_index: MemoIngredientIndex,
     ) -> ClaimResult<'me> {
-        let mut syncs = self.syncs.lock();
-        let thread_id = std::thread::current().id();
-
-        util::ensure_vec_len(&mut syncs, memo_ingredient_index.as_usize() + 1);
-
-        match &mut syncs[memo_ingredient_index.as_usize()] {
-            None => {
-                syncs[memo_ingredient_index.as_usize()] = Some(SyncState {
-                    id: thread_id,
-                    anyone_waiting: false,
-                });
-                ClaimResult::Claimed(ClaimGuard {
-                    database_key_index,
-                    memo_ingredient_index,
-                    zalsa,
-                    sync_table: self,
-                    _padding: false,
-                })
-            }
-            Some(SyncState {
-                id: other_id,
-                anyone_waiting,
-            }) => {
+        let mut write = self.syncs.lock();
+        match write.entry(database_key_index.key_index()) {
+            std::collections::hash_map::Entry::Occupied(occupied_entry) => {
+                let &mut SyncState {
+                    id,
+                    ref mut anyone_waiting,
+                } = occupied_entry.into_mut();
                 // NB: `Ordering::Relaxed` is sufficient here,
                 // as there are no loads that are "gated" on this
                 // value. Everything that is written is also protected
@@ -75,13 +56,25 @@ impl SyncTable {
                     db.as_dyn_database(),
                     db.zalsa_local(),
                     database_key_index,
-                    *other_id,
-                    syncs,
+                    id,
+                    write,
                 ) {
                     BlockResult::Completed => ClaimResult::Retry,
                     BlockResult::Cycle => ClaimResult::Cycle,
                 }
             }
+            std::collections::hash_map::Entry::Vacant(vacant_entry) => {
+                vacant_entry.insert(SyncState {
+                    id: std::thread::current().id(),
+                    anyone_waiting: false,
+                });
+                ClaimResult::Claimed(ClaimGuard {
+                    database_key_index,
+                    zalsa,
+                    sync_table: self,
+                    _padding: false,
+                })
+            }
         }
     }
 }
@@ -91,7 +84,6 @@ impl SyncTable {
 #[must_use]
 pub(crate) struct ClaimGuard<'me> {
     database_key_index: DatabaseKeyIndex,
-    memo_ingredient_index: MemoIngredientIndex,
     zalsa: &'me Zalsa,
     sync_table: &'me SyncTable,
     // Reduce the size of ClaimResult by making more niches available in ClaimGuard; this fits into
@@ -104,7 +96,9 @@ impl ClaimGuard<'_> {
         let mut syncs = self.sync_table.syncs.lock();
 
         let SyncState { anyone_waiting, .. } =
-            syncs[self.memo_ingredient_index.as_usize()].take().unwrap();
+            syncs.remove(&self.database_key_index.key_index()).unwrap();
+
+        drop(syncs);
 
         if anyone_waiting {
             self.zalsa.runtime().unblock_queries_blocked_on(
diff --git a/src/input.rs b/src/input.rs
@@ -17,7 +17,7 @@ use crate::{
     input::singleton::{Singleton, SingletonChoice},
     key::DatabaseKeyIndex,
     plumbing::{Jar, Stamp},
-    table::{memo::MemoTable, sync::SyncTable, Slot, Table},
+    table::{memo::MemoTable, Slot, Table},
     zalsa::{IngredientIndex, Zalsa},
     zalsa_local::QueryOrigin,
     Database, Durability, Id, Revision, Runtime,
@@ -107,7 +107,6 @@ impl<C: Configuration> IngredientImpl<C> {
                 fields,
                 stamps,
                 memos: Default::default(),
-                syncs: Default::default(),
             })
         });
 
@@ -286,9 +285,6 @@ where
 
     /// Memos
     memos: MemoTable,
-
-    /// Syncs
-    syncs: SyncTable,
 }
 
 impl<C> Value<C>
@@ -322,9 +318,4 @@ where
     fn memos_mut(&mut self) -> &mut crate::table::memo::MemoTable {
         &mut self.memos
     }
-
-    #[inline]
-    unsafe fn syncs(&self, _current_revision: Revision) -> &SyncTable {
-        &self.syncs
-    }
 }
diff --git a/src/interned.rs b/src/interned.rs
@@ -6,7 +6,6 @@ use crate::ingredient::fmt_index;
 use crate::plumbing::{IngredientIndices, Jar};
 use crate::revision::AtomicRevision;
 use crate::table::memo::MemoTable;
-use crate::table::sync::SyncTable;
 use crate::table::Slot;
 use crate::zalsa::{IngredientIndex, Zalsa};
 use crate::zalsa_local::QueryOrigin;
@@ -71,7 +70,6 @@ where
 {
     fields: C::Fields<'static>,
     memos: MemoTable,
-    syncs: SyncTable,
 
     /// The revision the value was first interned in.
     first_interned_at: Revision,
@@ -291,7 +289,6 @@ where
                 let id = zalsa_local.allocate(table, self.ingredient_index, |id| Value::<C> {
                     fields: unsafe { self.to_internal_data(assemble(id, key)) },
                     memos: Default::default(),
-                    syncs: Default::default(),
                     durability: AtomicU8::new(durability.as_u8()),
                     // Record the revision we are interning in.
                     first_interned_at: current_revision,
@@ -479,11 +476,6 @@ where
     fn memos_mut(&mut self) -> &mut MemoTable {
         &mut self.memos
     }
-
-    #[inline]
-    unsafe fn syncs(&self, _current_revision: Revision) -> &crate::table::sync::SyncTable {
-        &self.syncs
-    }
 }
 
 /// A trait for types that hash and compare like `O`.
diff --git a/src/table.rs b/src/table.rs
@@ -12,13 +12,10 @@ use std::{
 use memo::MemoTable;
 use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
-use sync::SyncTable;
 
 use crate::{Id, IngredientIndex, Revision};
 
 pub(crate) mod memo;
-pub(crate) mod sync;
-mod util;
 
 const PAGE_LEN_BITS: usize = 10;
 const PAGE_LEN_MASK: usize = PAGE_LEN - 1;
@@ -44,13 +41,6 @@ pub(crate) trait Slot: Any + Send + Sync {
 
     /// Mutably access the [`MemoTable`] for this slot.
     fn memos_mut(&mut self) -> &mut MemoTable;
-
-    /// Access the [`SyncTable`][] for this slot.
-    ///
-    /// # Safety condition
-    ///
-    /// The current revision MUST be the current revision of the database containing this slot.
-    unsafe fn syncs(&self, current_revision: Revision) -> &SyncTable;
 }
 
 /// [Slot::memos]
@@ -61,17 +51,12 @@ type SlotMemosFn<T> = unsafe fn(&T, current_revision: Revision) -> &MemoTable;
 type SlotMemosMutFnRaw = unsafe fn(*mut ()) -> *mut MemoTable;
 /// [Slot::memos_mut]
 type SlotMemosMutFn<T> = unsafe fn(&mut T) -> &mut MemoTable;
-/// [Slot::syncs]
-type SlotSyncsFnRaw = unsafe fn(*const (), current_revision: Revision) -> *const SyncTable;
-/// [Slot::syncs]
-type SlotSyncsFn<T> = unsafe fn(&T, current_revision: Revision) -> &SyncTable;
 
 struct SlotVTable {
     layout: Layout,
     /// [`Slot`] methods
     memos: SlotMemosFnRaw,
     memos_mut: SlotMemosMutFnRaw,
-    syncs: SlotSyncsFnRaw,
     /// A drop impl to call when the own page drops
     /// SAFETY: The caller is required to supply a correct data pointer to a `Box<PageDataEntry<T>>` and initialized length
     drop_impl: unsafe fn(data: *mut (), initialized: usize),
@@ -95,8 +80,6 @@ impl SlotVTable {
                 memos_mut: unsafe {
                     mem::transmute::<SlotMemosMutFn<T>, SlotMemosMutFnRaw>(T::memos_mut)
                 },
-                // SAFETY: The signatures are compatible
-                syncs: unsafe { mem::transmute::<SlotSyncsFn<T>, SlotSyncsFnRaw>(T::syncs) },
             }
         }
     }
@@ -245,19 +228,6 @@ impl Table {
         unsafe { &mut *(page.slot_vtable.memos_mut)(page.get(slot)) }
     }
 
-    /// Get the sync table associated with `id`
-    ///
-    /// # Safety condition
-    ///
-    /// The parameter `current_revision` MUST be the current revision
-    /// of the owner of database owning this table.
-    pub(crate) unsafe fn syncs(&self, id: Id, current_revision: Revision) -> &SyncTable {
-        let (page, slot) = split_id(id);
-        let page = &self.pages[page.0];
-        // SAFETY: We supply a proper slot pointer and the caller is required to pass the `current_revision`.
-        unsafe { &*(page.slot_vtable.syncs)(page.get(slot), current_revision) }
-    }
-
     pub(crate) fn slots_of<T: Slot>(&self) -> impl Iterator<Item = &T> + '_ {
         self.pages
             .iter()
diff --git a/src/table/util.rs b/src/table/util.rs
diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs
diff --git a/src/zalsa.rs b/src/zalsa.rs
