From a8383e3e4359b09449b9fdb95ef8ba36961e00ee Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Tue, 25 Feb 2025 23:42:24 +0200 Subject: [PATCH 1/5] Reduce memory usage by deduplicating type information We were storing the type information, 3 words wide, for each memo in each slot, while it is always constant wrt. the ingredient (different slots of the same ingredients will always have the same memos in the same order). This introduces some more unsafety, and the result wasn't as fast so I also had to use some lock-free structures, but the result is worth it: this shaves off 230mb from rust-analyzer with new Salsa. --- .../salsa-macro-rules/src/setup_tracked_fn.rs | 44 ++- src/accumulator.rs | 8 + src/function.rs | 17 +- src/function/memo.rs | 6 +- src/ingredient.rs | 4 + src/input.rs | 32 +- src/input/input_field.rs | 8 + src/interned.rs | 34 +- src/lib.rs | 6 +- src/memo_ingredient_indices.rs | 68 +++- src/table.rs | 62 ++- src/table/memo.rs | 372 +++++++++++------- src/tracked_struct.rs | 52 ++- src/tracked_struct/tracked_field.rs | 10 +- src/zalsa.rs | 17 +- src/zalsa_local.rs | 7 +- 16 files changed, 531 insertions(+), 216 deletions(-) diff --git a/components/salsa-macro-rules/src/setup_tracked_fn.rs b/components/salsa-macro-rules/src/setup_tracked_fn.rs index efc27398f..b35e32abe 100644 --- a/components/salsa-macro-rules/src/setup_tracked_fn.rs +++ b/components/salsa-macro-rules/src/setup_tracked_fn.rs @@ -254,21 +254,45 @@ macro_rules! setup_tracked_fn { struct_index } }; - let memo_ingredient_indices = From::from((zalsa, struct_index, first_index)); - let fn_ingredient = <$zalsa::function::IngredientImpl<$Configuration>>::new( - first_index, - memo_ingredient_indices, - $lru, - zalsa.views().downcaster_for::() - ); + $zalsa::macro_if! { $needs_interner => + let intern_ingredient = <$zalsa::interned::IngredientImpl<$Configuration>>::new( + first_index.successor(0) + ); + } + + let intern_ingredient_memo_types = $zalsa::macro_if! { + if $needs_interner { + Some($zalsa::Ingredient::memo_table_types(&intern_ingredient)) + } else { + None + } + }; + // SAFETY: We call with the correct memo types. + let memo_ingredient_indices = unsafe { + $zalsa::NewMemoIngredientIndices::create( + zalsa, + struct_index, + first_index, + $zalsa::function::MemoEntryType::of::<$zalsa::function::Memo<$Configuration>>(), + intern_ingredient_memo_types, + ) + }; + + // SAFETY: We pass the MemoEntryType for this Configuration, and we lookup the memo types table correctly. + let fn_ingredient = unsafe { + <$zalsa::function::IngredientImpl<$Configuration>>::new( + first_index, + memo_ingredient_indices, + $lru, + zalsa.views().downcaster_for::(), + ) + }; $zalsa::macro_if! { if $needs_interner { vec![ Box::new(fn_ingredient), - Box::new(<$zalsa::interned::IngredientImpl<$Configuration>>::new( - first_index.successor(0) - )), + Box::new(intern_ingredient), ] } else { vec![ diff --git a/src/accumulator.rs b/src/accumulator.rs index c01e487de..20260a589 100644 --- a/src/accumulator.rs +++ b/src/accumulator.rs @@ -4,12 +4,14 @@ use std::any::{Any, TypeId}; use std::fmt; use std::marker::PhantomData; use std::panic::UnwindSafe; +use std::sync::Arc; use accumulated::{Accumulated, AnyAccumulated}; use crate::function::VerifyResult; use crate::ingredient::{fmt_index, Ingredient, Jar}; use crate::plumbing::IngredientIndices; +use crate::table::memo::MemoTableTypes; use crate::zalsa::{IngredientIndex, Zalsa}; use crate::{Database, Id, Revision}; @@ -55,6 +57,7 @@ impl Jar for JarImpl { pub struct IngredientImpl { index: IngredientIndex, + memo_table_types: Arc, phantom: PhantomData>, } @@ -73,6 +76,7 @@ impl IngredientImpl { pub fn new(index: IngredientIndex) -> Self { Self { index, + memo_table_types: Arc::new(MemoTableTypes::default()), phantom: PhantomData, } } @@ -110,6 +114,10 @@ impl Ingredient for IngredientImpl { fn debug_name(&self) -> &'static str { A::DEBUG_NAME } + + fn memo_table_types(&self) -> Arc { + self.memo_table_types.clone() + } } impl std::fmt::Debug for IngredientImpl diff --git a/src/function.rs b/src/function.rs index d483323c2..4b3491051 100644 --- a/src/function.rs +++ b/src/function.rs @@ -1,6 +1,7 @@ use std::any::Any; use std::fmt; use std::ptr::NonNull; +use std::sync::Arc; pub(crate) use maybe_changed_after::VerifyResult; @@ -11,6 +12,7 @@ use crate::ingredient::{fmt_index, Ingredient}; use crate::key::DatabaseKeyIndex; use crate::plumbing::MemoIngredientMap; use crate::salsa_struct::SalsaStructInDb; +use crate::table::memo::MemoTableTypes; use crate::table::sync::ClaimResult; use crate::table::Table; use crate::views::DatabaseDownCaster; @@ -30,6 +32,8 @@ mod maybe_changed_after; mod memo; mod specify; +pub type Memo = memo::Memo<::Output<'static>>; + pub trait Configuration: Any { const DEBUG_NAME: &'static str; @@ -129,6 +133,9 @@ pub struct IngredientImpl { /// we don't know that we can trust the database to give us the same runtime /// everytime and so forth. deleted_entries: DeletedEntries, + + /// This is empty, but we need this for the trait and it doesn't consume a lot of memory anyway. + _memo_table_types: Arc, } /// True if `old_value == new_value`. Invoked by the generated @@ -142,7 +149,10 @@ impl IngredientImpl where C: Configuration, { - pub fn new( + /// # Safety + /// + /// `memo_type` and `memo_table_types` must be correct. + pub unsafe fn new( index: IngredientIndex, memo_ingredient_indices: as SalsaStructInDb>::MemoIngredientMap, lru: usize, @@ -153,6 +163,7 @@ where memo_ingredient_indices, lru: lru::Lru::new(lru), deleted_entries: Default::default(), + _memo_table_types: Arc::new(MemoTableTypes::default()), view_caster, } } @@ -314,6 +325,10 @@ where C::DEBUG_NAME } + fn memo_table_types(&self) -> Arc { + self._memo_table_types.clone() + } + fn cycle_recovery_strategy(&self) -> CycleRecoveryStrategy { C::CYCLE_STRATEGY } diff --git a/src/function/memo.rs b/src/function/memo.rs index eaa315cb2..5fe0a5432 100644 --- a/src/function/memo.rs +++ b/src/function/memo.rs @@ -9,7 +9,7 @@ use crate::cycle::{CycleHeads, CycleRecoveryStrategy, EMPTY_CYCLE_HEADS}; use crate::function::{Configuration, IngredientImpl}; use crate::key::DatabaseKeyIndex; use crate::revision::AtomicRevision; -use crate::table::memo::MemoTable; +use crate::table::memo::MemoTableWithTypesMut; use crate::zalsa::{MemoIngredientIndex, Zalsa}; use crate::zalsa_local::{QueryOrigin, QueryRevisions}; use crate::{Event, EventKind, Id, Revision}; @@ -84,7 +84,7 @@ impl IngredientImpl { /// with an equivalent memo that has no value. If the memo is untracked, FixpointInitial, /// or has values assigned as output of another query, this has no effect. pub(super) fn evict_value_from_memo_for( - table: &mut MemoTable, + table: MemoTableWithTypesMut<'_>, memo_ingredient_index: MemoIngredientIndex, ) { let map = |memo: &mut Memo>| { @@ -120,7 +120,7 @@ impl IngredientImpl { } #[derive(Debug)] -pub(super) struct Memo { +pub struct Memo { /// The result of the query, if we decide to memoize it. pub(super) value: Option, diff --git a/src/ingredient.rs b/src/ingredient.rs index a62e660de..418735e26 100644 --- a/src/ingredient.rs +++ b/src/ingredient.rs @@ -1,10 +1,12 @@ use std::any::{Any, TypeId}; use std::fmt; +use std::sync::Arc; use crate::accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues}; use crate::cycle::CycleRecoveryStrategy; use crate::function::VerifyResult; use crate::plumbing::IngredientIndices; +use crate::table::memo::MemoTableTypes; use crate::table::Table; use crate::zalsa::{transmute_data_mut_ptr, transmute_data_ptr, IngredientIndex, Zalsa}; use crate::zalsa_local::QueryOrigin; @@ -132,6 +134,8 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync { ); } + fn memo_table_types(&self) -> Arc; + fn fmt_index(&self, index: crate::Id, fmt: &mut fmt::Formatter<'_>) -> fmt::Result; // Function ingredient methods diff --git a/src/input.rs b/src/input.rs index 4d4405d74..d9261993d 100644 --- a/src/input.rs +++ b/src/input.rs @@ -1,6 +1,7 @@ use std::any::{Any, TypeId}; use std::fmt; use std::ops::DerefMut; +use std::sync::Arc; pub mod input_field; pub mod setter; @@ -14,7 +15,7 @@ use crate::ingredient::{fmt_index, Ingredient}; use crate::input::singleton::{Singleton, SingletonChoice}; use crate::key::DatabaseKeyIndex; use crate::plumbing::{Jar, Stamp}; -use crate::table::memo::MemoTable; +use crate::table::memo::{MemoTable, MemoTableTypes}; use crate::table::sync::SyncTable; use crate::table::{Slot, Table}; use crate::zalsa::{IngredientIndex, Zalsa}; @@ -72,6 +73,7 @@ impl Jar for JarImpl { pub struct IngredientImpl { ingredient_index: IngredientIndex, singleton: C::Singleton, + memo_table_types: Arc, _phantom: std::marker::PhantomData, } @@ -80,6 +82,7 @@ impl IngredientImpl { Self { ingredient_index: index, singleton: Default::default(), + memo_table_types: Arc::new(MemoTableTypes::default()), _phantom: std::marker::PhantomData, } } @@ -100,12 +103,17 @@ impl IngredientImpl { let (zalsa, zalsa_local) = db.zalsas(); let id = self.singleton.with_scope(|| { - zalsa_local.allocate(zalsa.table(), self.ingredient_index, |_| Value:: { - fields, - stamps, - memos: Default::default(), - syncs: Default::default(), - }) + zalsa_local.allocate( + zalsa, + zalsa.table(), + self.ingredient_index, + |_| Value:: { + fields, + stamps, + memos: Default::default(), + syncs: Default::default(), + }, + ) }); FromIdWithDb::from_id(id, db) @@ -219,6 +227,10 @@ impl Ingredient for IngredientImpl { fn debug_name(&self) -> &'static str { C::DEBUG_NAME } + + fn memo_table_types(&self) -> Arc { + self.memo_table_types.clone() + } } impl std::fmt::Debug for IngredientImpl { @@ -286,4 +298,10 @@ where unsafe fn syncs(&self, _current_revision: Revision) -> &SyncTable { &self.syncs } + + unsafe fn drop_memos(&mut self, types: &MemoTableTypes) { + // SAFETY: Our precondition. + let memos = unsafe { types.attach_memos_mut(&mut self.memos) }; + memos.drop(); + } } diff --git a/src/input/input_field.rs b/src/input/input_field.rs index 17aef7044..006942468 100644 --- a/src/input/input_field.rs +++ b/src/input/input_field.rs @@ -1,9 +1,11 @@ use std::fmt; use std::marker::PhantomData; +use std::sync::Arc; use crate::function::VerifyResult; use crate::ingredient::{fmt_index, Ingredient}; use crate::input::{Configuration, IngredientImpl, Value}; +use crate::table::memo::MemoTableTypes; use crate::zalsa::IngredientIndex; use crate::{Database, Id, Revision}; @@ -19,6 +21,7 @@ use crate::{Database, Id, Revision}; pub struct FieldIngredientImpl { index: IngredientIndex, field_index: usize, + memo_table_types: Arc, phantom: PhantomData Value>, } @@ -30,6 +33,7 @@ where Self { index: struct_index.successor(field_index), field_index, + memo_table_types: Arc::new(MemoTableTypes::default()), phantom: PhantomData, } } @@ -69,6 +73,10 @@ where fn debug_name(&self) -> &'static str { C::FIELD_DEBUG_NAMES[self.field_index] } + + fn memo_table_types(&self) -> Arc { + self.memo_table_types.clone() + } } impl std::fmt::Debug for FieldIngredientImpl diff --git a/src/interned.rs b/src/interned.rs index 7be6a3420..8c8f91e2e 100644 --- a/src/interned.rs +++ b/src/interned.rs @@ -16,7 +16,7 @@ use crate::hash::FxDashMap; use crate::ingredient::{fmt_index, Ingredient}; use crate::plumbing::{IngredientIndices, Jar}; use crate::revision::AtomicRevision; -use crate::table::memo::MemoTable; +use crate::table::memo::{MemoTable, MemoTableTypes}; use crate::table::sync::SyncTable; use crate::table::Slot; use crate::zalsa::{IngredientIndex, Zalsa}; @@ -62,6 +62,8 @@ pub struct IngredientImpl { /// /// Deadlock requirement: We access `value_map` while holding lock on `key_map`, but not vice versa. key_map: FxDashMap, Id>, + + memo_table_types: Arc, } /// Struct storing the interned fields. @@ -132,6 +134,7 @@ where Self { ingredient_index, key_map: Default::default(), + memo_table_types: Arc::new(MemoTableTypes::default()), } } @@ -288,15 +291,16 @@ where // If there is no active query this durability does not actually matter. .unwrap_or(Durability::MAX); - let id = zalsa_local.allocate(table, self.ingredient_index, |id| Value:: { - 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, - last_interned_at: AtomicRevision::from(current_revision), - }); + let id = + zalsa_local.allocate(zalsa, table, self.ingredient_index, |id| Value:: { + 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, + last_interned_at: AtomicRevision::from(current_revision), + }); let value = table.get::>(id); @@ -409,6 +413,10 @@ where fn debug_name(&self) -> &'static str { C::DEBUG_NAME } + + fn memo_table_types(&self) -> Arc { + self.memo_table_types.clone() + } } impl std::fmt::Debug for IngredientImpl @@ -440,6 +448,12 @@ where unsafe fn syncs(&self, _current_revision: Revision) -> &crate::table::sync::SyncTable { &self.syncs } + + unsafe fn drop_memos(&mut self, types: &MemoTableTypes) { + // SAFETY: Our precondition. + let memos = unsafe { types.attach_memos_mut(&mut self.memos) }; + memos.drop(); + } } /// A trait for types that hash and compare like `O`. diff --git a/src/lib.rs b/src/lib.rs index 48676b83f..d69c59c17 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -87,6 +87,7 @@ pub mod plumbing { pub use crate::key::DatabaseKeyIndex; pub use crate::memo_ingredient_indices::{ IngredientIndices, MemoIngredientIndices, MemoIngredientMap, MemoIngredientSingletonIndex, + NewMemoIngredientIndices, }; pub use crate::revision::Revision; pub use crate::runtime::{stamp, Runtime, Stamp, StampedValue}; @@ -118,7 +119,10 @@ pub mod plumbing { } pub mod function { - pub use crate::function::{Configuration, IngredientImpl}; + pub use crate::function::Configuration; + pub use crate::function::IngredientImpl; + pub use crate::function::Memo; + pub use crate::table::memo::MemoEntryType; } pub mod tracked_struct { diff --git a/src/memo_ingredient_indices.rs b/src/memo_ingredient_indices.rs index a784b4ea4..4bffcc24a 100644 --- a/src/memo_ingredient_indices.rs +++ b/src/memo_ingredient_indices.rs @@ -1,3 +1,6 @@ +use std::sync::Arc; + +use crate::table::memo::{MemoEntryType, MemoTableTypes}; use crate::zalsa::{MemoIngredientIndex, Zalsa}; use crate::{Id, IngredientIndex}; @@ -42,11 +45,34 @@ impl IngredientIndices { } } -impl From<(&Zalsa, IngredientIndices, IngredientIndex)> for MemoIngredientIndices { - #[inline] - fn from( - (zalsa, struct_indices, ingredient): (&Zalsa, IngredientIndices, IngredientIndex), +pub trait NewMemoIngredientIndices { + /// # Safety + /// + /// The memo types must be correct. + unsafe fn create( + zalsa: &Zalsa, + struct_indices: IngredientIndices, + ingredient: IngredientIndex, + memo_type: MemoEntryType, + intern_ingredient_memo_types: Option>, + ) -> Self; +} + +impl NewMemoIngredientIndices for MemoIngredientIndices { + /// # Safety + /// + /// The memo types must be correct. + unsafe fn create( + zalsa: &Zalsa, + struct_indices: IngredientIndices, + ingredient: IngredientIndex, + memo_type: MemoEntryType, + _intern_ingredient_memo_types: Option>, ) -> Self { + debug_assert!( + _intern_ingredient_memo_types.is_none(), + "intern ingredient can only have a singleton memo ingredient" + ); let Some(&last) = struct_indices.indices.last() else { unreachable!("Attempting to construct struct memo mapping for non tracked function?") }; @@ -56,8 +82,18 @@ impl From<(&Zalsa, IngredientIndices, IngredientIndex)> for MemoIngredientIndice MemoIngredientIndex::from_usize((u32::MAX - 1) as usize), ); for &struct_ingredient in &struct_indices.indices { - indices[struct_ingredient.as_usize()] = - zalsa.next_memo_ingredient_index(struct_ingredient, ingredient); + let memo_types = zalsa + .lookup_ingredient(struct_ingredient) + .memo_table_types(); + // SAFETY: By our precondition the memo types are correct. + indices[struct_ingredient.as_usize()] = unsafe { + zalsa.next_memo_ingredient_index( + struct_ingredient, + ingredient, + memo_type, + &memo_types, + ) + }; } MemoIngredientIndices { indices: indices.into_boxed_slice(), @@ -112,14 +148,28 @@ impl MemoIngredientMap for MemoIngredientSingletonIndex { } } -impl From<(&Zalsa, IngredientIndices, IngredientIndex)> for MemoIngredientSingletonIndex { +impl NewMemoIngredientIndices for MemoIngredientSingletonIndex { #[inline] - fn from((zalsa, indices, ingredient): (&Zalsa, IngredientIndices, IngredientIndex)) -> Self { + unsafe fn create( + zalsa: &Zalsa, + indices: IngredientIndices, + ingredient: IngredientIndex, + memo_type: MemoEntryType, + intern_ingredient_memo_types: Option>, + ) -> Self { let &[struct_ingredient] = &*indices.indices else { unreachable!("Attempting to construct struct memo mapping from enum?") }; - Self(zalsa.next_memo_ingredient_index(struct_ingredient, ingredient)) + let memo_types = intern_ingredient_memo_types.unwrap_or_else(|| { + zalsa + .lookup_ingredient(struct_ingredient) + .memo_table_types() + }); + // SAFETY: By our precondition, we are passed the correct memo types. + Self(unsafe { + zalsa.next_memo_ingredient_index(struct_ingredient, ingredient, memo_type, &memo_types) + }) } } diff --git a/src/table.rs b/src/table.rs index 768746ca5..c3cecd53b 100644 --- a/src/table.rs +++ b/src/table.rs @@ -6,12 +6,15 @@ use std::mem::{self, MaybeUninit}; use std::ptr::{self, NonNull}; use std::slice; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; use memo::MemoTable; use parking_lot::Mutex; use rustc_hash::FxHashMap; use sync::SyncTable; +use crate::table::memo::{MemoTableTypes, MemoTableWithTypes, MemoTableWithTypesMut}; +use crate::zalsa::Zalsa; use crate::{Id, IngredientIndex, Revision}; pub(crate) mod memo; @@ -49,6 +52,11 @@ pub(crate) trait Slot: Any + Send + Sync { /// /// The current revision MUST be the current revision of the database containing this slot. unsafe fn syncs(&self, current_revision: Revision) -> &SyncTable; + + /// # Safety + /// + /// `types` must be correct. + unsafe fn drop_memos(&mut self, types: &MemoTableTypes); } /// [Slot::memos] @@ -71,20 +79,23 @@ struct SlotVTable { 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>` and initialized length - drop_impl: unsafe fn(data: *mut (), initialized: usize), + /// SAFETY: The caller is required to supply a correct data pointer to a `Box>` and initialized length, + /// and correct memo types. + drop_impl: unsafe fn(data: *mut (), initialized: usize, memo_types: &MemoTableTypes), } impl SlotVTable { const fn of() -> &'static Self { const { &Self { - drop_impl: |data, initialized| + drop_impl: |data, initialized, memo_types| // SAFETY: The caller is required to supply a correct data pointer and initialized length unsafe { let data = Box::from_raw(data.cast::>()); for i in 0..initialized { - ptr::drop_in_place(data[i].get().cast::()); + let item = data[i].get().cast::(); + T::drop_memos(&mut *item, memo_types); + ptr::drop_in_place(item); } }, layout: Layout::new::(), @@ -133,6 +144,8 @@ struct Page { /// The type name of what is stored as entries in data. // FIXME: Move this into SlotVTable once const stable slot_type_name: &'static str, + + memo_types: Arc, } // SAFETY: `Page` is `Send` as we make sure to only ever store `Slot` types in it which @@ -215,8 +228,12 @@ impl Table { } /// Allocate a new page for the given ingredient and with slots of type `T` - pub(crate) fn push_page(&self, ingredient: IngredientIndex) -> PageIndex { - PageIndex::new(self.pages.push(Page::new::(ingredient))) + pub(crate) fn push_page( + &self, + ingredient: IngredientIndex, + zalsa: &Zalsa, + ) -> PageIndex { + PageIndex::new(self.pages.push(Page::new::(ingredient, zalsa))) } /// Get the memo table associated with `id` @@ -225,15 +242,21 @@ impl Table { /// /// The parameter `current_revision` MUST be the current revision /// of the owner of database owning this table. - pub(crate) unsafe fn memos(&self, id: Id, current_revision: Revision) -> &MemoTable { + pub(crate) unsafe fn memos( + &self, + id: Id, + current_revision: Revision, + ) -> MemoTableWithTypes<'_> { 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.memos)(page.get(slot), current_revision) } + let memos = unsafe { &*(page.slot_vtable.memos)(page.get(slot), current_revision) }; + // SAFETY: The `Page` keeps the correct memo types. + unsafe { page.memo_types.attach_memos(memos) } } /// Get the memo table associated with `id` - pub(crate) fn memos_mut(&mut self, id: Id) -> &mut MemoTable { + pub(crate) fn memos_mut(&mut self, id: Id) -> MemoTableWithTypesMut<'_> { let (page, slot) = split_id(id); let page_index = page.0; let page = self @@ -241,7 +264,9 @@ impl Table { .get_mut(page_index) .unwrap_or_else(|| panic!("index `{page_index}` is uninitialized")); // SAFETY: We supply a proper slot pointer and the caller is required to pass the `current_revision`. - unsafe { &mut *(page.slot_vtable.memos_mut)(page.get(slot)) } + let memos = unsafe { &mut *(page.slot_vtable.memos_mut)(page.get(slot)) }; + // SAFETY: The `Page` keeps the correct memo types. + unsafe { page.memo_types.attach_memos_mut(memos) } } /// Get the sync table associated with `id` @@ -264,7 +289,11 @@ impl Table { .flat_map(|view| view.data()) } - pub(crate) fn fetch_or_push_page(&self, ingredient: IngredientIndex) -> PageIndex { + pub(crate) fn fetch_or_push_page( + &self, + ingredient: IngredientIndex, + zalsa: &Zalsa, + ) -> PageIndex { if let Some(page) = self .non_full_pages .lock() @@ -273,7 +302,7 @@ impl Table { { return page; } - self.push_page::(ingredient) + self.push_page::(ingredient, zalsa) } pub(crate) fn record_unfilled_page(&self, ingredient: IngredientIndex, page: PageIndex) { @@ -325,7 +354,11 @@ impl<'p, T: Slot> PageView<'p, T> { } impl Page { - fn new(ingredient: IngredientIndex) -> Self { + fn new(ingredient: IngredientIndex, zalsa: &Zalsa) -> Self { + let memo_types = zalsa + .lookup_ingredient(ingredient) + .memo_table_types() + .clone(); let data: Box> = Box::new([const { UnsafeCell::new(MaybeUninit::uninit()) }; PAGE_LEN]); Self { @@ -336,6 +369,7 @@ impl Page { allocated: Default::default(), allocation_lock: Default::default(), data: NonNull::from(Box::leak(data)).cast::<()>(), + memo_types, } } @@ -382,7 +416,7 @@ impl Drop for Page { fn drop(&mut self) { let &mut len = self.allocated.get_mut(); // SAFETY: We supply the data pointer and the initialized length - unsafe { (self.slot_vtable.drop_impl)(self.data.as_ptr(), len) }; + unsafe { (self.slot_vtable.drop_impl)(self.data.as_ptr(), len, &self.memo_types) }; } } diff --git a/src/table/memo.rs b/src/table/memo.rs index 07e3085d7..148585c9b 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -1,12 +1,15 @@ -use std::any::{Any, TypeId}; -use std::ptr::NonNull; -use std::sync::atomic::{AtomicPtr, Ordering}; +use std::{ + any::{Any, TypeId}, + fmt::Debug, + mem, + ptr::{self, NonNull}, + sync::atomic::{AtomicPtr, Ordering}, +}; use parking_lot::RwLock; use thin_vec::ThinVec; -use crate::zalsa::MemoIngredientIndex; -use crate::zalsa_local::QueryOrigin; +use crate::{zalsa::MemoIngredientIndex, zalsa_local::QueryOrigin}; /// The "memo table" stores the memoized results of tracked function calls. /// Every tracked function must take a salsa struct as its first argument @@ -16,19 +19,11 @@ pub(crate) struct MemoTable { memos: RwLock>, } -pub(crate) trait Memo: Any + Send + Sync { +pub trait Memo: Any + Send + Sync { /// Returns the `origin` of this memo fn origin(&self) -> &QueryOrigin; } -/// Wraps the data stored for a memoized entry. -/// This struct has a customized Drop that will -/// ensure that its `data` field is properly freed. -#[derive(Default)] -struct MemoEntry { - data: Option, -} - /// Data for a memoized entry. /// This is a type-erased `Box`, where `M` is the type of memo associated /// with that particular ingredient index. @@ -46,21 +41,22 @@ struct MemoEntry { /// Therefore, we hide the type by transmuting to `DummyMemo`; but we must then be very careful /// when freeing `MemoEntryData` values to transmute things back. See the `Drop` impl for /// [`MemoEntry`][] for details. -struct MemoEntryData { +#[derive(Default)] +struct MemoEntry { + /// An [`AtomicPtr`][] to a `Box` for the erased memo type `M` + atomic_memo: AtomicPtr, +} + +#[derive(Clone, Copy)] +pub struct MemoEntryType { /// The `type_id` of the erased memo type `M` type_id: TypeId, /// A type-coercion function for the erased memo type `M` to_dyn_fn: fn(NonNull) -> NonNull, - - /// An [`AtomicPtr`][] to a `Box` for the erased memo type `M` - atomic_memo: AtomicPtr, } -/// Dummy placeholder type that we use when erasing the memo type `M` in [`MemoEntryData`][]. -struct DummyMemo {} - -impl MemoTable { +impl MemoEntryType { fn to_dummy(memo: NonNull) -> NonNull { memo.cast() } @@ -74,46 +70,122 @@ impl MemoTable { #[allow(clippy::undocumented_unsafe_blocks)] // TODO(#697) document safety unsafe { - std::mem::transmute::< + mem::transmute::< fn(NonNull) -> NonNull, fn(NonNull) -> NonNull, >(f) } } + #[inline] + pub fn of() -> Self { + Self { + type_id: TypeId::of::(), + to_dyn_fn: Self::to_dyn_fn::(), + } + } +} + +/// Dummy placeholder type that we use when erasing the memo type `M` in [`MemoEntryData`][]. +#[derive(Debug)] +struct DummyMemo {} + +impl Memo for DummyMemo { + fn origin(&self) -> &QueryOrigin { + unreachable!("should not get here") + } +} + +#[derive(Default)] +pub struct MemoTableTypes { + types: RwLock>>, +} + +impl MemoTableTypes { + pub(crate) fn set(&self, memo_ingredient_index: MemoIngredientIndex, memo_type: MemoEntryType) { + let mut types = self.types.write(); + let memo_ingredient_index = memo_ingredient_index.as_usize(); + if memo_ingredient_index >= types.len() { + types.resize_with(memo_ingredient_index + 1, Default::default); + } + match &mut types[memo_ingredient_index] { + Some(existing) => { + assert_eq!( + existing.type_id, memo_type.type_id, + "inconsistent type-id for `{memo_ingredient_index:?}`" + ); + } + entry @ None => { + *entry = Some(memo_type); + } + } + } + /// # Safety /// - /// The caller needs to make sure to not free the returned value until no more references into - /// the database exist as there may be outstanding borrows into the pointer contents. + /// The types table must be the correct one of `memos`. + #[inline] + pub(crate) unsafe fn attach_memos<'a>( + &'a self, + memos: &'a MemoTable, + ) -> MemoTableWithTypes<'a> { + MemoTableWithTypes { types: self, memos } + } + + /// # Safety + /// + /// The types table must be the correct one of `memos`. + #[inline] + pub(crate) unsafe fn attach_memos_mut<'a>( + &'a self, + memos: &'a mut MemoTable, + ) -> MemoTableWithTypesMut<'a> { + MemoTableWithTypesMut { types: self, memos } + } +} + +pub(crate) struct MemoTableWithTypes<'a> { + types: &'a MemoTableTypes, + memos: &'a MemoTable, +} + +impl<'a> MemoTableWithTypes<'a> { + /// # Safety + /// + /// The caller needs to make sure to not drop the returned value until no more references into + /// the database exist as there may be outstanding borrows into the `Arc` contents. pub(crate) unsafe fn insert( - &self, + self, memo_ingredient_index: MemoIngredientIndex, memo: NonNull, ) -> Option> { + // The type must already exist, we insert it when creating the memo ingredient. + let types = self.types.types.read(); + assert_eq!( + types[memo_ingredient_index.as_usize()] + .as_ref() + .expect("memo type should be available in insert") + .type_id, + TypeId::of::(), + "inconsistent type-id for `{memo_ingredient_index:?}`" + ); + drop(types); + // If the memo slot is already occupied, it must already have the // right type info etc, and we only need the read-lock. - if let Some(MemoEntry { - data: - Some(MemoEntryData { - type_id, - to_dyn_fn: _, - atomic_memo, - }), - }) = self.memos.read().get(memo_ingredient_index.as_usize()) + if let Some(MemoEntry { atomic_memo }) = self + .memos + .memos + .read() + .get(memo_ingredient_index.as_usize()) { - assert_eq!( - *type_id, - TypeId::of::(), - "inconsistent type-id for `{memo_ingredient_index:?}`" - ); + let old_memo = + atomic_memo.swap(MemoEntryType::to_dummy(memo).as_ptr(), Ordering::AcqRel); - let old_memo = atomic_memo.swap(Self::to_dummy(memo).as_ptr(), Ordering::AcqRel); - - // SAFETY: The `atomic_memo` field is never null. - let old_memo = unsafe { NonNull::new_unchecked(old_memo) }; + let old_memo = NonNull::new(old_memo); // SAFETY: `type_id` check asserted above - return Some(unsafe { Self::from_dummy(old_memo) }); + return old_memo.map(|old_memo| unsafe { MemoEntryType::from_dummy(old_memo) }); } // Otherwise we need the write lock. @@ -123,142 +195,162 @@ impl MemoTable { /// # Safety /// - /// The caller needs to make sure to not free the returned value until no more references into - /// the database exist as there may be outstanding borrows into the pointer contents. + /// The caller needs to make sure to not drop the returned value until no more references into + /// the database exist as there may be outstanding borrows into the `Arc` contents. unsafe fn insert_cold( - &self, + self, memo_ingredient_index: MemoIngredientIndex, memo: NonNull, ) -> Option> { - let mut memos = self.memos.write(); let memo_ingredient_index = memo_ingredient_index.as_usize(); + let mut memos = self.memos.memos.write(); + let additional_len = memo_ingredient_index - memos.len() + 1; + memos.reserve(additional_len); while memos.len() < memo_ingredient_index + 1 { - memos.push(MemoEntry { data: None }); + memos.push(MemoEntry::default()); } - let old_entry = memos[memo_ingredient_index].data.replace(MemoEntryData { - type_id: TypeId::of::(), - to_dyn_fn: Self::to_dyn_fn::(), - atomic_memo: AtomicPtr::new(Self::to_dummy(memo).as_ptr()), - }); - old_entry.map( - |MemoEntryData { - type_id: _, - to_dyn_fn: _, - atomic_memo, - }| - // SAFETY: The `atomic_memo` field is never null. - unsafe { Self::from_dummy(NonNull::new_unchecked(atomic_memo.into_inner())) }, - ) + let old_entry = mem::replace( + memos[memo_ingredient_index].atomic_memo.get_mut(), + MemoEntryType::to_dummy(memo).as_ptr(), + ); + let old_entry = NonNull::new(old_entry); + // SAFETY: The `TypeId` is asserted in `insert()`. + old_entry.map(|memo| unsafe { MemoEntryType::from_dummy(memo) }) } - pub(crate) fn get(&self, memo_ingredient_index: MemoIngredientIndex) -> Option<&M> { - let memos = self.memos.read(); - - let Some(MemoEntry { - data: - Some(MemoEntryData { - type_id, - to_dyn_fn: _, - atomic_memo, - }), - }) = memos.get(memo_ingredient_index.as_usize()) - else { - return None; - }; + pub(crate) fn get(self, memo_ingredient_index: MemoIngredientIndex) -> Option<&'a M> { + if let Some(MemoEntry { atomic_memo }) = self + .memos + .memos + .read() + .get(memo_ingredient_index.as_usize()) + { + if let Some(Some(MemoEntryType { + type_id, + to_dyn_fn: _, + })) = self + .types + .types + .read() + .get(memo_ingredient_index.as_usize()) + { + assert_eq!( + *type_id, + TypeId::of::(), + "inconsistent type-id for `{memo_ingredient_index:?}`" + ); + let memo = NonNull::new(atomic_memo.load(Ordering::Acquire)); + // SAFETY: `type_id` check asserted above + return memo.map(|memo| unsafe { MemoEntryType::from_dummy(memo).as_ref() }); + } + } - assert_eq!( - *type_id, - TypeId::of::(), - "inconsistent type-id for `{memo_ingredient_index:?}`" - ); + None + } +} - // SAFETY: The `atomic_memo` field is never null. - let memo = unsafe { NonNull::new_unchecked(atomic_memo.load(Ordering::Acquire)) }; +pub(crate) struct MemoTableWithTypesMut<'a> { + types: &'a MemoTableTypes, + memos: &'a mut MemoTable, +} - // SAFETY: `type_id` check asserted above - unsafe { Some(Self::from_dummy(memo).as_ref()) } +impl<'a> MemoTableWithTypesMut<'a> { + #[inline] + pub(crate) fn reborrow<'b>(&'b mut self) -> MemoTableWithTypesMut<'b> + where + 'a: 'b, + { + MemoTableWithTypesMut { + types: self.types, + memos: self.memos, + } } /// Calls `f` on the memo at `memo_ingredient_index`. /// /// If the memo is not present, `f` is not called. pub(crate) fn map_memo( - &mut self, + self, memo_ingredient_index: MemoIngredientIndex, f: impl FnOnce(&mut M), ) { - let memos = self.memos.get_mut(); - let Some(MemoEntry { - data: - Some(MemoEntryData { - type_id, - to_dyn_fn: _, - atomic_memo, - }), - }) = memos.get_mut(memo_ingredient_index.as_usize()) - else { + let types = self.types.types.read(); + let Some(Some(memo_type)) = types.get(memo_ingredient_index.as_usize()) else { return; }; - assert_eq!( - *type_id, + memo_type.type_id, TypeId::of::(), "inconsistent type-id for `{memo_ingredient_index:?}`" ); + drop(types); - // SAFETY: The `atomic_memo` field is never null. - let memo = unsafe { NonNull::new_unchecked(*atomic_memo.get_mut()) }; + // If the memo slot is already occupied, it must already have the + // right type info etc, and we only need the read-lock. + let memos = self.memos.memos.get_mut(); + let Some(MemoEntry { atomic_memo }) = memos.get_mut(memo_ingredient_index.as_usize()) + else { + return; + }; + let Some(memo) = NonNull::new(*atomic_memo.get_mut()) else { + return; + }; // SAFETY: `type_id` check asserted above - f(unsafe { Self::from_dummy(memo).as_mut() }); + f(unsafe { MemoEntryType::from_dummy(memo).as_mut() }); + } + + /// To drop an entry, we need its type, so we don't implement `Drop`, and instead have this method. + #[inline] + pub fn drop(self) { + let types = self.types.types.read(); + let types = types.iter(); + for (type_, memo) in std::iter::zip(types, self.memos.memos.get_mut()) { + if let Some(type_) = type_ { + // SAFETY: The types match because this is an invariant of `MemoTableWithTypesMut`. + unsafe { + memo.drop(type_); + } + } + } } /// # Safety /// /// The caller needs to make sure to not call this function until no more references into /// the database exist as there may be outstanding borrows into the pointer contents. - pub(crate) unsafe fn into_memos( - self, - ) -> impl Iterator)> { - self.memos - .into_inner() - .into_iter() + pub(crate) unsafe fn with_memos(self, mut f: impl FnMut(MemoIngredientIndex, Box)) { + let memos = self.memos.memos.get_mut(); + let types = self.types.types.read(); + memos + .iter_mut() + .zip(types.iter()) .zip(0..) - .filter_map(|(mut memo, index)| memo.data.take().map(|d| (d, index))) - .map( - |( - MemoEntryData { - type_id: _, - to_dyn_fn, - atomic_memo, - }, - index, - )| { - // SAFETY: The `atomic_memo` field is never null. - let memo = - unsafe { to_dyn_fn(NonNull::new_unchecked(atomic_memo.into_inner())) }; - // SAFETY: The caller guarantees that there are no outstanding borrows into the `Box` contents. - let memo = unsafe { Box::from_raw(memo.as_ptr()) }; - - (MemoIngredientIndex::from_usize(index), memo) - }, - ) + .filter_map(|((memo, type_), index)| { + let memo = mem::replace(memo.atomic_memo.get_mut(), ptr::null_mut()); + let memo = NonNull::new(memo)?; + Some((memo, type_.as_ref()?, index)) + }) + .map(|(memo, type_, index)| { + // SAFETY: We took ownership of the memo, and converted it to the correct type. + // The caller guarantees that there are no outstanding borrows into the `Box` contents. + let memo = unsafe { Box::from_raw((type_.to_dyn_fn)(memo).as_ptr()) }; + (MemoIngredientIndex::from_usize(index), memo) + }) + .for_each(|(index, memo)| f(index, memo)); } } -impl Drop for MemoEntry { - fn drop(&mut self) { - if let Some(MemoEntryData { - type_id: _, - to_dyn_fn, - atomic_memo, - }) = self.data.take() +impl MemoEntry { + /// # Safety + /// + /// The type must match. + #[inline] + unsafe fn drop(&mut self, type_: &MemoEntryType) { + if let Some(memo) = NonNull::new(mem::replace(self.atomic_memo.get_mut(), ptr::null_mut())) { - // SAFETY: The `atomic_memo` field is never null. - let memo = unsafe { to_dyn_fn(NonNull::new_unchecked(atomic_memo.into_inner())) }; - // SAFETY: We have `&mut self`, so there are no outstanding borrows into the `Box` contents. - let memo = unsafe { Box::from_raw(memo.as_ptr()) }; - std::mem::drop(memo); + // SAFETY: Our preconditions. + mem::drop(unsafe { Box::from_raw((type_.to_dyn_fn)(memo).as_ptr()) }); } } } @@ -271,6 +363,6 @@ impl Drop for DummyMemo { impl std::fmt::Debug for MemoTable { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("MemoTable").finish() + f.debug_struct("MemoTable").finish_non_exhaustive() } } diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs index dc835daad..41ab93ec7 100644 --- a/src/tracked_struct.rs +++ b/src/tracked_struct.rs @@ -5,6 +5,7 @@ use std::fmt; use std::hash::Hash; use std::marker::PhantomData; use std::ops::DerefMut; +use std::sync::Arc; use crossbeam_queue::SegQueue; use tracked_field::FieldIngredientImpl; @@ -16,7 +17,7 @@ use crate::plumbing::ZalsaLocal; use crate::revision::OptionalAtomicRevision; use crate::runtime::StampedValue; use crate::salsa_struct::SalsaStructInDb; -use crate::table::memo::MemoTable; +use crate::table::memo::{MemoTable, MemoTableTypes, MemoTableWithTypesMut}; use crate::table::sync::SyncTable; use crate::table::{Slot, Table}; use crate::zalsa::{IngredientIndex, Zalsa}; @@ -168,6 +169,8 @@ where /// Store freed ids free_list: SegQueue, + + memo_table_types: Arc, } /// Defines the identity of a tracked struct. @@ -385,6 +388,7 @@ where ingredient_index: index, phantom: PhantomData, free_list: Default::default(), + memo_table_types: Arc::new(MemoTableTypes::default()), } } @@ -466,7 +470,7 @@ where id } else { - zalsa_local.allocate::>(zalsa.table(), self.ingredient_index, value) + zalsa_local.allocate::>(zalsa, zalsa.table(), self.ingredient_index, value) } } @@ -629,21 +633,35 @@ where // Take the memo table. This is safe because we have modified `data_ref.updated_at` to `None` // and the code that references the memo-table has a read-lock. - let memo_table = unsafe { (*data).take_memo_table() }; - + struct MemoTableWithTypes<'a>(MemoTableWithTypesMut<'a>); + impl Drop for MemoTableWithTypes<'_> { + fn drop(&mut self) { + self.0.reborrow().drop(); + } + } + // Keep those statements in this order, so that if the first panics we won't leak the table. + let mut memo_table = unsafe { (*data).take_memo_table() }; + // SAFETY: We use the correct types table. + let mut memo_table = + MemoTableWithTypes(unsafe { self.memo_table_types.attach_memos_mut(&mut memo_table) }); // SAFETY: We have verified that no more references to these memos exist and so we are good // to drop them. - for (memo_ingredient_index, memo) in unsafe { memo_table.into_memos() } { - let ingredient_index = - zalsa.ingredient_index_for_memo(self.ingredient_index, memo_ingredient_index); + unsafe { + memo_table + .0 + .reborrow() + .with_memos(|memo_ingredient_index, memo| { + let ingredient_index = zalsa + .ingredient_index_for_memo(self.ingredient_index, memo_ingredient_index); - let executor = DatabaseKeyIndex::new(ingredient_index, id); + let executor = DatabaseKeyIndex::new(ingredient_index, id); - db.salsa_event(&|| Event::new(EventKind::DidDiscard { key: executor })); + db.salsa_event(&|| Event::new(EventKind::DidDiscard { key: executor })); - for stale_output in memo.origin().outputs() { - stale_output.remove_stale_output(zalsa, db, executor, provisional); - } + for stale_output in memo.origin().outputs() { + stale_output.remove_stale_output(zalsa, db, executor, provisional); + } + }); } // now that all cleanup has occurred, make available for re-use @@ -790,6 +808,10 @@ where fn debug_name(&self) -> &'static str { C::DEBUG_NAME } + + fn memo_table_types(&self) -> Arc { + self.memo_table_types.clone() + } } impl std::fmt::Debug for IngredientImpl @@ -875,6 +897,12 @@ where self.read_lock(current_revision); &self.syncs } + + unsafe fn drop_memos(&mut self, types: &MemoTableTypes) { + // SAFETY: Our precondition. + let memos = unsafe { types.attach_memos_mut(&mut self.memos) }; + memos.drop(); + } } #[cfg(test)] diff --git a/src/tracked_struct/tracked_field.rs b/src/tracked_struct/tracked_field.rs index 03e435a88..bb0d4137c 100644 --- a/src/tracked_struct/tracked_field.rs +++ b/src/tracked_struct/tracked_field.rs @@ -1,7 +1,8 @@ -use std::marker::PhantomData; +use std::{marker::PhantomData, sync::Arc}; use crate::function::VerifyResult; use crate::ingredient::Ingredient; +use crate::table::memo::MemoTableTypes; use crate::tracked_struct::{Configuration, Value}; use crate::zalsa::IngredientIndex; use crate::{Database, Id}; @@ -24,7 +25,7 @@ where /// The absolute index of this field on the tracked struct. field_index: usize, - + memo_table_types: Arc, phantom: PhantomData Value>, } @@ -35,6 +36,7 @@ where pub(super) fn new(field_index: usize, ingredient_index: IngredientIndex) -> Self { Self { field_index, + memo_table_types: Arc::new(MemoTableTypes::default()), ingredient_index, phantom: PhantomData, } @@ -74,6 +76,10 @@ where fn debug_name(&self) -> &'static str { C::FIELD_DEBUG_NAMES[self.field_index] } + + fn memo_table_types(&self) -> Arc { + self.memo_table_types.clone() + } } impl std::fmt::Debug for FieldIngredientImpl diff --git a/src/zalsa.rs b/src/zalsa.rs index bb5716343..2a1a96f8c 100644 --- a/src/zalsa.rs +++ b/src/zalsa.rs @@ -13,7 +13,7 @@ use rustc_hash::FxHashMap; use crate::ingredient::{Ingredient, Jar}; use crate::nonce::{Nonce, NonceGenerator}; use crate::runtime::Runtime; -use crate::table::memo::MemoTable; +use crate::table::memo::{MemoEntryType, MemoTableTypes, MemoTableWithTypes}; use crate::table::sync::SyncTable; use crate::table::Table; use crate::views::Views; @@ -189,9 +189,10 @@ impl Zalsa { } /// Returns the [`MemoTable`][] for the salsa struct with the given id - pub(crate) fn memo_table_for(&self, id: Id) -> &MemoTable { + pub(crate) fn memo_table_for(&self, id: Id) -> MemoTableWithTypes<'_> { + let table = self.table(); // SAFETY: We are supplying the correct current revision - unsafe { self.table().memos(id, self.current_revision()) } + unsafe { table.memos(id, self.current_revision()) } } /// Returns the [`SyncTable`][] for the salsa struct with the given id @@ -233,10 +234,15 @@ impl Zalsa { } } - pub(crate) fn next_memo_ingredient_index( + /// # Safety + /// + /// The memo types must be correct. + pub(crate) unsafe fn next_memo_ingredient_index( &self, struct_ingredient_index: IngredientIndex, ingredient_index: IngredientIndex, + memo_type: MemoEntryType, + memo_types: &MemoTableTypes, ) -> MemoIngredientIndex { let mut memo_ingredients = self.memo_ingredient_indices.write(); let idx = struct_ingredient_index.as_usize(); @@ -248,6 +254,9 @@ impl Zalsa { }; let mi = MemoIngredientIndex(u32::try_from(memo_ingredients.len()).unwrap()); memo_ingredients.push(ingredient_index); + + memo_types.set(mi, memo_type); + mi } } diff --git a/src/zalsa_local.rs b/src/zalsa_local.rs index 70d48e506..e0e6f27b3 100644 --- a/src/zalsa_local.rs +++ b/src/zalsa_local.rs @@ -13,7 +13,7 @@ use crate::key::DatabaseKeyIndex; use crate::runtime::Stamp; use crate::table::{PageIndex, Slot, Table}; use crate::tracked_struct::{Disambiguator, Identity, IdentityHash, IdentityMap}; -use crate::zalsa::IngredientIndex; +use crate::zalsa::{IngredientIndex, Zalsa}; use crate::{Accumulator, Cancelled, Id, Revision}; /// State that is specific to a single execution thread. @@ -54,6 +54,7 @@ impl ZalsaLocal { /// thread and attempts to reuse it. pub(crate) fn allocate( &self, + zalsa: &Zalsa, table: &Table, ingredient: IngredientIndex, mut value: impl FnOnce(Id) -> T, @@ -63,7 +64,7 @@ impl ZalsaLocal { .most_recent_pages .borrow_mut() .entry(ingredient) - .or_insert_with(|| table.fetch_or_push_page::(ingredient)); + .or_insert_with(|| table.fetch_or_push_page::(ingredient, zalsa)); loop { // Try to allocate an entry on that page @@ -77,7 +78,7 @@ impl ZalsaLocal { // it is unlikely that there is a non-full one available. Err(v) => { value = v; - page = table.push_page::(ingredient); + page = table.push_page::(ingredient, zalsa); self.most_recent_pages.borrow_mut().insert(ingredient, page); } } From c96ccfb1e0241463dd012b83a129aa08158bb14f Mon Sep 17 00:00:00 2001 From: Lukas Wirth Date: Thu, 17 Apr 2025 15:32:20 +0200 Subject: [PATCH 2/5] Simplify --- src/accumulator.rs | 4 +--- src/function.rs | 6 +----- src/input.rs | 23 ++++++--------------- src/input/input_field.rs | 4 +--- src/interned.rs | 31 ++++++++++------------------- src/memo_ingredient_indices.rs | 22 +++++++++----------- src/runtime.rs | 1 + src/table.rs | 25 +++++++++-------------- src/table/memo.rs | 17 ++-------------- src/tracked_struct.rs | 31 ++++++++++------------------- src/tracked_struct/tracked_field.rs | 4 +--- src/zalsa.rs | 12 +++-------- src/zalsa_local.rs | 17 ++++++++++++---- 13 files changed, 69 insertions(+), 128 deletions(-) diff --git a/src/accumulator.rs b/src/accumulator.rs index 20260a589..819767bd3 100644 --- a/src/accumulator.rs +++ b/src/accumulator.rs @@ -57,7 +57,6 @@ impl Jar for JarImpl { pub struct IngredientImpl { index: IngredientIndex, - memo_table_types: Arc, phantom: PhantomData>, } @@ -76,7 +75,6 @@ impl IngredientImpl { pub fn new(index: IngredientIndex) -> Self { Self { index, - memo_table_types: Arc::new(MemoTableTypes::default()), phantom: PhantomData, } } @@ -116,7 +114,7 @@ impl Ingredient for IngredientImpl { } fn memo_table_types(&self) -> Arc { - self.memo_table_types.clone() + unreachable!("accumulator does not allocate pages") } } diff --git a/src/function.rs b/src/function.rs index 4b3491051..49d6e7f5d 100644 --- a/src/function.rs +++ b/src/function.rs @@ -133,9 +133,6 @@ pub struct IngredientImpl { /// we don't know that we can trust the database to give us the same runtime /// everytime and so forth. deleted_entries: DeletedEntries, - - /// This is empty, but we need this for the trait and it doesn't consume a lot of memory anyway. - _memo_table_types: Arc, } /// True if `old_value == new_value`. Invoked by the generated @@ -163,7 +160,6 @@ where memo_ingredient_indices, lru: lru::Lru::new(lru), deleted_entries: Default::default(), - _memo_table_types: Arc::new(MemoTableTypes::default()), view_caster, } } @@ -326,7 +322,7 @@ where } fn memo_table_types(&self) -> Arc { - self._memo_table_types.clone() + unreachable!("function does not allocate pages") } fn cycle_recovery_strategy(&self) -> CycleRecoveryStrategy { diff --git a/src/input.rs b/src/input.rs index d9261993d..4802860e5 100644 --- a/src/input.rs +++ b/src/input.rs @@ -103,17 +103,12 @@ impl IngredientImpl { let (zalsa, zalsa_local) = db.zalsas(); let id = self.singleton.with_scope(|| { - zalsa_local.allocate( - zalsa, - zalsa.table(), - self.ingredient_index, - |_| Value:: { - fields, - stamps, - memos: Default::default(), - syncs: Default::default(), - }, - ) + zalsa_local.allocate(zalsa, self.ingredient_index, |_| Value:: { + fields, + stamps, + memos: Default::default(), + syncs: Default::default(), + }) }); FromIdWithDb::from_id(id, db) @@ -298,10 +293,4 @@ where unsafe fn syncs(&self, _current_revision: Revision) -> &SyncTable { &self.syncs } - - unsafe fn drop_memos(&mut self, types: &MemoTableTypes) { - // SAFETY: Our precondition. - let memos = unsafe { types.attach_memos_mut(&mut self.memos) }; - memos.drop(); - } } diff --git a/src/input/input_field.rs b/src/input/input_field.rs index 006942468..a987eb909 100644 --- a/src/input/input_field.rs +++ b/src/input/input_field.rs @@ -21,7 +21,6 @@ use crate::{Database, Id, Revision}; pub struct FieldIngredientImpl { index: IngredientIndex, field_index: usize, - memo_table_types: Arc, phantom: PhantomData Value>, } @@ -33,7 +32,6 @@ where Self { index: struct_index.successor(field_index), field_index, - memo_table_types: Arc::new(MemoTableTypes::default()), phantom: PhantomData, } } @@ -75,7 +73,7 @@ where } fn memo_table_types(&self) -> Arc { - self.memo_table_types.clone() + unreachable!("input fields do not allocate pages") } } diff --git a/src/interned.rs b/src/interned.rs index 8c8f91e2e..0b0cb808d 100644 --- a/src/interned.rs +++ b/src/interned.rs @@ -282,8 +282,6 @@ where // We won any races so should intern the data Err(slot) => { - let table = zalsa.table(); - // Record the durability of the current query on the interned value. let durability = zalsa_local .active_query() @@ -291,18 +289,17 @@ where // If there is no active query this durability does not actually matter. .unwrap_or(Durability::MAX); - let id = - zalsa_local.allocate(zalsa, table, self.ingredient_index, |id| Value:: { - 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, - last_interned_at: AtomicRevision::from(current_revision), - }); + let id = zalsa_local.allocate(zalsa, self.ingredient_index, |id| Value:: { + 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, + last_interned_at: AtomicRevision::from(current_revision), + }); - let value = table.get::>(id); + let value = zalsa.table().get::>(id); let slot_value = (value.fields.clone(), SharedValue::new(id)); unsafe { lock.insert_in_slot(data_hash, slot, slot_value) }; @@ -311,7 +308,7 @@ where data_hash, self.key_map .hasher() - .hash_one(table.get::>(id).fields.clone()) + .hash_one(zalsa.table().get::>(id).fields.clone()) ); // Record a dependency on this value. @@ -448,12 +445,6 @@ where unsafe fn syncs(&self, _current_revision: Revision) -> &crate::table::sync::SyncTable { &self.syncs } - - unsafe fn drop_memos(&mut self, types: &MemoTableTypes) { - // SAFETY: Our precondition. - let memos = unsafe { types.attach_memos_mut(&mut self.memos) }; - memos.drop(); - } } /// A trait for types that hash and compare like `O`. diff --git a/src/memo_ingredient_indices.rs b/src/memo_ingredient_indices.rs index 4bffcc24a..c1eba9da2 100644 --- a/src/memo_ingredient_indices.rs +++ b/src/memo_ingredient_indices.rs @@ -85,15 +85,11 @@ impl NewMemoIngredientIndices for MemoIngredientIndices { let memo_types = zalsa .lookup_ingredient(struct_ingredient) .memo_table_types(); - // SAFETY: By our precondition the memo types are correct. - indices[struct_ingredient.as_usize()] = unsafe { - zalsa.next_memo_ingredient_index( - struct_ingredient, - ingredient, - memo_type, - &memo_types, - ) - }; + + let mi = zalsa.next_memo_ingredient_index(struct_ingredient, ingredient); + memo_types.set(mi, memo_type); + + indices[struct_ingredient.as_usize()] = mi; } MemoIngredientIndices { indices: indices.into_boxed_slice(), @@ -166,10 +162,10 @@ impl NewMemoIngredientIndices for MemoIngredientSingletonIndex { .lookup_ingredient(struct_ingredient) .memo_table_types() }); - // SAFETY: By our precondition, we are passed the correct memo types. - Self(unsafe { - zalsa.next_memo_ingredient_index(struct_ingredient, ingredient, memo_type, &memo_types) - }) + + let mi = zalsa.next_memo_ingredient_index(struct_ingredient, ingredient); + memo_types.set(mi, memo_type); + Self(mi) } } diff --git a/src/runtime.rs b/src/runtime.rs index c1c31184a..e9450840b 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -133,6 +133,7 @@ impl Runtime { } /// Returns the [`Table`] used to store the value of salsa structs + #[inline] pub(crate) fn table(&self) -> &Table { &self.table } diff --git a/src/table.rs b/src/table.rs index c3cecd53b..113b3146e 100644 --- a/src/table.rs +++ b/src/table.rs @@ -14,7 +14,6 @@ use rustc_hash::FxHashMap; use sync::SyncTable; use crate::table::memo::{MemoTableTypes, MemoTableWithTypes, MemoTableWithTypesMut}; -use crate::zalsa::Zalsa; use crate::{Id, IngredientIndex, Revision}; pub(crate) mod memo; @@ -52,11 +51,6 @@ pub(crate) trait Slot: Any + Send + Sync { /// /// The current revision MUST be the current revision of the database containing this slot. unsafe fn syncs(&self, current_revision: Revision) -> &SyncTable; - - /// # Safety - /// - /// `types` must be correct. - unsafe fn drop_memos(&mut self, types: &MemoTableTypes); } /// [Slot::memos] @@ -94,7 +88,7 @@ impl SlotVTable { let data = Box::from_raw(data.cast::>()); for i in 0..initialized { let item = data[i].get().cast::(); - T::drop_memos(&mut *item, memo_types); + memo_types.attach_memos_mut((*item).memos_mut()).drop(); ptr::drop_in_place(item); } }, @@ -159,6 +153,7 @@ unsafe impl Sync for Page {} pub struct PageIndex(usize); impl PageIndex { + #[inline] fn new(idx: usize) -> Self { debug_assert!(idx < MAX_PAGES); Self(idx) @@ -228,12 +223,13 @@ impl Table { } /// Allocate a new page for the given ingredient and with slots of type `T` + #[inline] pub(crate) fn push_page( &self, ingredient: IngredientIndex, - zalsa: &Zalsa, + memo_types: Arc, ) -> PageIndex { - PageIndex::new(self.pages.push(Page::new::(ingredient, zalsa))) + PageIndex::new(self.pages.push(Page::new::(ingredient, memo_types))) } /// Get the memo table associated with `id` @@ -292,7 +288,7 @@ impl Table { pub(crate) fn fetch_or_push_page( &self, ingredient: IngredientIndex, - zalsa: &Zalsa, + memo_types: impl FnOnce() -> Arc, ) -> PageIndex { if let Some(page) = self .non_full_pages @@ -302,7 +298,7 @@ impl Table { { return page; } - self.push_page::(ingredient, zalsa) + self.push_page::(ingredient, memo_types()) } pub(crate) fn record_unfilled_page(&self, ingredient: IngredientIndex, page: PageIndex) { @@ -354,11 +350,8 @@ impl<'p, T: Slot> PageView<'p, T> { } impl Page { - fn new(ingredient: IngredientIndex, zalsa: &Zalsa) -> Self { - let memo_types = zalsa - .lookup_ingredient(ingredient) - .memo_table_types() - .clone(); + #[inline] + fn new(ingredient: IngredientIndex, memo_types: Arc) -> Self { let data: Box> = Box::new([const { UnsafeCell::new(MaybeUninit::uninit()) }; PAGE_LEN]); Self { diff --git a/src/table/memo.rs b/src/table/memo.rs index 148585c9b..123eab8a9 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -254,18 +254,7 @@ pub(crate) struct MemoTableWithTypesMut<'a> { memos: &'a mut MemoTable, } -impl<'a> MemoTableWithTypesMut<'a> { - #[inline] - pub(crate) fn reborrow<'b>(&'b mut self) -> MemoTableWithTypesMut<'b> - where - 'a: 'b, - { - MemoTableWithTypesMut { - types: self.types, - memos: self.memos, - } - } - +impl MemoTableWithTypesMut<'_> { /// Calls `f` on the memo at `memo_ingredient_index`. /// /// If the memo is not present, `f` is not called. @@ -308,9 +297,7 @@ impl<'a> MemoTableWithTypesMut<'a> { for (type_, memo) in std::iter::zip(types, self.memos.memos.get_mut()) { if let Some(type_) = type_ { // SAFETY: The types match because this is an invariant of `MemoTableWithTypesMut`. - unsafe { - memo.drop(type_); - } + unsafe { memo.drop(type_) }; } } } diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs index 41ab93ec7..ec389ecbd 100644 --- a/src/tracked_struct.rs +++ b/src/tracked_struct.rs @@ -17,7 +17,7 @@ use crate::plumbing::ZalsaLocal; use crate::revision::OptionalAtomicRevision; use crate::runtime::StampedValue; use crate::salsa_struct::SalsaStructInDb; -use crate::table::memo::{MemoTable, MemoTableTypes, MemoTableWithTypesMut}; +use crate::table::memo::{MemoTable, MemoTableTypes}; use crate::table::sync::SyncTable; use crate::table::{Slot, Table}; use crate::zalsa::{IngredientIndex, Zalsa}; @@ -470,7 +470,7 @@ where id } else { - zalsa_local.allocate::>(zalsa, zalsa.table(), self.ingredient_index, value) + zalsa_local.allocate::>(zalsa, self.ingredient_index, value) } } @@ -633,24 +633,20 @@ where // Take the memo table. This is safe because we have modified `data_ref.updated_at` to `None` // and the code that references the memo-table has a read-lock. - struct MemoTableWithTypes<'a>(MemoTableWithTypesMut<'a>); + struct MemoTableWithTypes<'a>(MemoTable, &'a MemoTableTypes); impl Drop for MemoTableWithTypes<'_> { fn drop(&mut self) { - self.0.reborrow().drop(); + // SAFETY: We use the correct types table. + unsafe { self.1.attach_memos_mut(&mut self.0) }.drop(); } } - // Keep those statements in this order, so that if the first panics we won't leak the table. - let mut memo_table = unsafe { (*data).take_memo_table() }; - // SAFETY: We use the correct types table. let mut memo_table = - MemoTableWithTypes(unsafe { self.memo_table_types.attach_memos_mut(&mut memo_table) }); + MemoTableWithTypes(unsafe { (*data).take_memo_table() }, &self.memo_table_types); // SAFETY: We have verified that no more references to these memos exist and so we are good // to drop them. unsafe { - memo_table - .0 - .reborrow() - .with_memos(|memo_ingredient_index, memo| { + memo_table.1.attach_memos_mut(&mut memo_table.0).with_memos( + |memo_ingredient_index, memo| { let ingredient_index = zalsa .ingredient_index_for_memo(self.ingredient_index, memo_ingredient_index); @@ -661,8 +657,9 @@ where for stale_output in memo.origin().outputs() { stale_output.remove_stale_output(zalsa, db, executor, provisional); } - }); - } + }, + ) + }; // now that all cleanup has occurred, make available for re-use self.free_list.push(id); @@ -897,12 +894,6 @@ where self.read_lock(current_revision); &self.syncs } - - unsafe fn drop_memos(&mut self, types: &MemoTableTypes) { - // SAFETY: Our precondition. - let memos = unsafe { types.attach_memos_mut(&mut self.memos) }; - memos.drop(); - } } #[cfg(test)] diff --git a/src/tracked_struct/tracked_field.rs b/src/tracked_struct/tracked_field.rs index bb0d4137c..8c09cbefa 100644 --- a/src/tracked_struct/tracked_field.rs +++ b/src/tracked_struct/tracked_field.rs @@ -25,7 +25,6 @@ where /// The absolute index of this field on the tracked struct. field_index: usize, - memo_table_types: Arc, phantom: PhantomData Value>, } @@ -36,7 +35,6 @@ where pub(super) fn new(field_index: usize, ingredient_index: IngredientIndex) -> Self { Self { field_index, - memo_table_types: Arc::new(MemoTableTypes::default()), ingredient_index, phantom: PhantomData, } @@ -78,7 +76,7 @@ where } fn memo_table_types(&self) -> Arc { - self.memo_table_types.clone() + unreachable!("tracked field does not allocate pages") } } diff --git a/src/zalsa.rs b/src/zalsa.rs index 2a1a96f8c..67ca5b7f3 100644 --- a/src/zalsa.rs +++ b/src/zalsa.rs @@ -13,7 +13,7 @@ use rustc_hash::FxHashMap; use crate::ingredient::{Ingredient, Jar}; use crate::nonce::{Nonce, NonceGenerator}; use crate::runtime::Runtime; -use crate::table::memo::{MemoEntryType, MemoTableTypes, MemoTableWithTypes}; +use crate::table::memo::MemoTableWithTypes; use crate::table::sync::SyncTable; use crate::table::Table; use crate::views::Views; @@ -184,6 +184,7 @@ impl Zalsa { } /// Returns the [`Table`] used to store the value of salsa structs + #[inline] pub(crate) fn table(&self) -> &Table { self.runtime.table() } @@ -234,15 +235,10 @@ impl Zalsa { } } - /// # Safety - /// - /// The memo types must be correct. - pub(crate) unsafe fn next_memo_ingredient_index( + pub(crate) fn next_memo_ingredient_index( &self, struct_ingredient_index: IngredientIndex, ingredient_index: IngredientIndex, - memo_type: MemoEntryType, - memo_types: &MemoTableTypes, ) -> MemoIngredientIndex { let mut memo_ingredients = self.memo_ingredient_indices.write(); let idx = struct_ingredient_index.as_usize(); @@ -255,8 +251,6 @@ impl Zalsa { let mi = MemoIngredientIndex(u32::try_from(memo_ingredients.len()).unwrap()); memo_ingredients.push(ingredient_index); - memo_types.set(mi, memo_type); - mi } } diff --git a/src/zalsa_local.rs b/src/zalsa_local.rs index e0e6f27b3..b3096a03d 100644 --- a/src/zalsa_local.rs +++ b/src/zalsa_local.rs @@ -55,20 +55,29 @@ impl ZalsaLocal { pub(crate) fn allocate( &self, zalsa: &Zalsa, - table: &Table, ingredient: IngredientIndex, mut value: impl FnOnce(Id) -> T, ) -> Id { + let memo_types = || { + zalsa + .lookup_ingredient(ingredient) + .memo_table_types() + .clone() + }; // Find the most recent page, pushing a page if needed let mut page = *self .most_recent_pages .borrow_mut() .entry(ingredient) - .or_insert_with(|| table.fetch_or_push_page::(ingredient, zalsa)); + .or_insert_with(|| { + zalsa + .table() + .fetch_or_push_page::(ingredient, memo_types) + }); loop { // Try to allocate an entry on that page - let page_ref = table.page::(page); + let page_ref = zalsa.table().page::(page); match page_ref.allocate(page, value) { // If successful, return Ok(id) => return id, @@ -78,7 +87,7 @@ impl ZalsaLocal { // it is unlikely that there is a non-full one available. Err(v) => { value = v; - page = table.push_page::(ingredient, zalsa); + page = zalsa.table().push_page::(ingredient, memo_types()); self.most_recent_pages.borrow_mut().insert(ingredient, page); } } From 7e1b66f451679edd7880ba4063522d66cfe15062 Mon Sep 17 00:00:00 2001 From: Lukas Wirth Date: Thu, 17 Apr 2025 16:14:18 +0200 Subject: [PATCH 3/5] Replace `RwLock` with boxcar + `AtomicPtr` --- src/memo_ingredient_indices.rs | 4 +- src/table.rs | 1 + src/table/const_type_id.rs | 105 +++++++++++++++++++++++++ src/table/memo.rs | 137 ++++++++++++++++++--------------- src/zalsa.rs | 1 + 5 files changed, 184 insertions(+), 64 deletions(-) create mode 100644 src/table/const_type_id.rs diff --git a/src/memo_ingredient_indices.rs b/src/memo_ingredient_indices.rs index c1eba9da2..6b2fc9c23 100644 --- a/src/memo_ingredient_indices.rs +++ b/src/memo_ingredient_indices.rs @@ -87,7 +87,7 @@ impl NewMemoIngredientIndices for MemoIngredientIndices { .memo_table_types(); let mi = zalsa.next_memo_ingredient_index(struct_ingredient, ingredient); - memo_types.set(mi, memo_type); + memo_types.set(mi, &memo_type); indices[struct_ingredient.as_usize()] = mi; } @@ -164,7 +164,7 @@ impl NewMemoIngredientIndices for MemoIngredientSingletonIndex { }); let mi = zalsa.next_memo_ingredient_index(struct_ingredient, ingredient); - memo_types.set(mi, memo_type); + memo_types.set(mi, &memo_type); Self(mi) } } diff --git a/src/table.rs b/src/table.rs index 113b3146e..0e4ae1db8 100644 --- a/src/table.rs +++ b/src/table.rs @@ -16,6 +16,7 @@ use sync::SyncTable; use crate::table::memo::{MemoTableTypes, MemoTableWithTypes, MemoTableWithTypesMut}; use crate::{Id, IngredientIndex, Revision}; +mod const_type_id; pub(crate) mod memo; pub(crate) mod sync; mod util; diff --git a/src/table/const_type_id.rs b/src/table/const_type_id.rs new file mode 100644 index 000000000..d56a34449 --- /dev/null +++ b/src/table/const_type_id.rs @@ -0,0 +1,105 @@ +//! Forked from + +use core::any::TypeId; +use core::cmp::Ordering; +use core::fmt::{self, Debug}; +use core::hash::{Hash, Hasher}; +use core::marker::PhantomData; +use core::mem; + +/// TypeId equivalent usable in const contexts. +#[derive(Copy, Clone)] +#[repr(C)] +pub struct ConstTypeId { + type_id_fn: fn() -> TypeId, +} + +impl ConstTypeId { + /// Create a [`ConstTypeId`] for a type. + #[must_use] + pub const fn of() -> Self + where + T: ?Sized, + { + ConstTypeId { + type_id_fn: of::, + } + } + + #[inline] + fn get(self) -> TypeId { + (self.type_id_fn)() + } +} + +impl Debug for ConstTypeId { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + Debug::fmt(&self.get(), formatter) + } +} + +impl PartialEq for ConstTypeId { + #[inline] + fn eq(&self, other: &Self) -> bool { + self.get() == other.get() + } +} + +impl PartialEq for ConstTypeId { + fn eq(&self, other: &TypeId) -> bool { + self.get() == *other + } +} + +impl Eq for ConstTypeId {} + +impl PartialOrd for ConstTypeId { + #[inline] + fn partial_cmp(&self, other: &Self) -> Option { + Some(Ord::cmp(self, other)) + } +} + +impl Ord for ConstTypeId { + #[inline] + fn cmp(&self, other: &Self) -> Ordering { + Ord::cmp(&self.get(), &other.get()) + } +} + +impl Hash for ConstTypeId { + fn hash(&self, state: &mut H) { + self.get().hash(state); + } +} + +/// Create a [`TypeId`] for a type. +#[must_use] +#[inline(always)] +pub fn of() -> TypeId +where + T: ?Sized, +{ + trait NonStaticAny { + fn get_type_id(&self) -> TypeId + where + Self: 'static; + } + + impl NonStaticAny for PhantomData { + #[inline(always)] + fn get_type_id(&self) -> TypeId + where + Self: 'static, + { + TypeId::of::() + } + } + + let phantom_data = PhantomData::; + // SAFETY: We only punt the lifetime to fetch the type id, we don't actually use the lifetime + // in any meaningful way. + NonStaticAny::get_type_id(unsafe { + mem::transmute::<&dyn NonStaticAny, &(dyn NonStaticAny + 'static)>(&phantom_data) + }) +} diff --git a/src/table/memo.rs b/src/table/memo.rs index 123eab8a9..d65395f6b 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -9,7 +9,9 @@ use std::{ use parking_lot::RwLock; use thin_vec::ThinVec; -use crate::{zalsa::MemoIngredientIndex, zalsa_local::QueryOrigin}; +use crate::{ + table::const_type_id::ConstTypeId, zalsa::MemoIngredientIndex, zalsa_local::QueryOrigin, +}; /// The "memo table" stores the memoized results of tracked function calls. /// Every tracked function must take a salsa struct as its first argument @@ -47,10 +49,14 @@ struct MemoEntry { atomic_memo: AtomicPtr, } -#[derive(Clone, Copy)] pub struct MemoEntryType { + data: AtomicPtr, +} + +#[derive(Clone, Copy)] +struct MemoEntryTypeData { /// The `type_id` of the erased memo type `M` - type_id: TypeId, + type_id: ConstTypeId, /// A type-coercion function for the erased memo type `M` to_dyn_fn: fn(NonNull) -> NonNull, @@ -65,7 +71,7 @@ impl MemoEntryType { memo.cast() } - fn to_dyn_fn() -> fn(NonNull) -> NonNull { + const fn to_dyn_fn() -> fn(NonNull) -> NonNull { let f: fn(NonNull) -> NonNull = |x| x; #[allow(clippy::undocumented_unsafe_blocks)] // TODO(#697) document safety @@ -80,10 +86,24 @@ impl MemoEntryType { #[inline] pub fn of() -> Self { Self { - type_id: TypeId::of::(), - to_dyn_fn: Self::to_dyn_fn::(), + data: AtomicPtr::new( + (&const { + MemoEntryTypeData { + type_id: ConstTypeId::of::(), + to_dyn_fn: Self::to_dyn_fn::(), + } + } as *const MemoEntryTypeData) + .cast_mut(), + ), } } + + #[inline] + fn load(&self) -> Option<&MemoEntryTypeData> { + // Note: Relaxed is fine as we only set the pointer once, right when we initialize it. + // SAFETY: The pointer is either null or initialized properly. + unsafe { self.data.load(Ordering::Relaxed).as_ref() } + } } /// Dummy placeholder type that we use when erasing the memo type `M` in [`MemoEntryData`][]. @@ -98,27 +118,26 @@ impl Memo for DummyMemo { #[derive(Default)] pub struct MemoTableTypes { - types: RwLock>>, + types: boxcar::Vec, } impl MemoTableTypes { - pub(crate) fn set(&self, memo_ingredient_index: MemoIngredientIndex, memo_type: MemoEntryType) { - let mut types = self.types.write(); + pub(crate) fn set( + &self, + memo_ingredient_index: MemoIngredientIndex, + memo_type: &MemoEntryType, + ) { let memo_ingredient_index = memo_ingredient_index.as_usize(); - if memo_ingredient_index >= types.len() { - types.resize_with(memo_ingredient_index + 1, Default::default); - } - match &mut types[memo_ingredient_index] { - Some(existing) => { - assert_eq!( - existing.type_id, memo_type.type_id, - "inconsistent type-id for `{memo_ingredient_index:?}`" - ); - } - entry @ None => { - *entry = Some(memo_type); - } + while memo_ingredient_index >= self.types.count() { + self.types.push(MemoEntryType { + data: AtomicPtr::default(), + }); } + let memo_entry_type = self.types.get(memo_ingredient_index).unwrap(); + let old = memo_entry_type + .data + .swap(memo_type.data.load(Ordering::Relaxed), Ordering::AcqRel); + debug_assert!(old.is_null(), "memo type should only be set once"); } /// # Safety @@ -160,16 +179,15 @@ impl<'a> MemoTableWithTypes<'a> { memo: NonNull, ) -> Option> { // The type must already exist, we insert it when creating the memo ingredient. - let types = self.types.types.read(); assert_eq!( - types[memo_ingredient_index.as_usize()] - .as_ref() - .expect("memo type should be available in insert") + self.types + .types + .get(memo_ingredient_index.as_usize()) + .and_then(MemoEntryType::load)? .type_id, TypeId::of::(), "inconsistent type-id for `{memo_ingredient_index:?}`" ); - drop(types); // If the memo slot is already occupied, it must already have the // right type info etc, and we only need the read-lock. @@ -225,24 +243,18 @@ impl<'a> MemoTableWithTypes<'a> { .read() .get(memo_ingredient_index.as_usize()) { - if let Some(Some(MemoEntryType { - type_id, - to_dyn_fn: _, - })) = self - .types - .types - .read() - .get(memo_ingredient_index.as_usize()) - { - assert_eq!( - *type_id, - TypeId::of::(), - "inconsistent type-id for `{memo_ingredient_index:?}`" - ); - let memo = NonNull::new(atomic_memo.load(Ordering::Acquire)); - // SAFETY: `type_id` check asserted above - return memo.map(|memo| unsafe { MemoEntryType::from_dummy(memo).as_ref() }); - } + assert_eq!( + self.types + .types + .get(memo_ingredient_index.as_usize()) + .and_then(MemoEntryType::load)? + .type_id, + TypeId::of::(), + "inconsistent type-id for `{memo_ingredient_index:?}`" + ); + let memo = NonNull::new(atomic_memo.load(Ordering::Acquire)); + // SAFETY: `type_id` check asserted above + return memo.map(|memo| unsafe { MemoEntryType::from_dummy(memo).as_ref() }); } None @@ -263,16 +275,19 @@ impl MemoTableWithTypesMut<'_> { memo_ingredient_index: MemoIngredientIndex, f: impl FnOnce(&mut M), ) { - let types = self.types.types.read(); - let Some(Some(memo_type)) = types.get(memo_ingredient_index.as_usize()) else { + let Some(type_) = self + .types + .types + .get(memo_ingredient_index.as_usize()) + .and_then(MemoEntryType::load) + else { return; }; assert_eq!( - memo_type.type_id, + type_.type_id, TypeId::of::(), "inconsistent type-id for `{memo_ingredient_index:?}`" ); - drop(types); // If the memo slot is already occupied, it must already have the // right type info etc, and we only need the read-lock. @@ -292,13 +307,10 @@ impl MemoTableWithTypesMut<'_> { /// To drop an entry, we need its type, so we don't implement `Drop`, and instead have this method. #[inline] pub fn drop(self) { - let types = self.types.types.read(); - let types = types.iter(); - for (type_, memo) in std::iter::zip(types, self.memos.memos.get_mut()) { - if let Some(type_) = type_ { - // SAFETY: The types match because this is an invariant of `MemoTableWithTypesMut`. - unsafe { memo.drop(type_) }; - } + let types = self.types.types.iter(); + for ((_, type_), memo) in std::iter::zip(types, self.memos.memos.get_mut()) { + // SAFETY: The types match because this is an invariant of `MemoTableWithTypesMut`. + unsafe { memo.drop(type_) }; } } @@ -308,15 +320,14 @@ impl MemoTableWithTypesMut<'_> { /// the database exist as there may be outstanding borrows into the pointer contents. pub(crate) unsafe fn with_memos(self, mut f: impl FnMut(MemoIngredientIndex, Box)) { let memos = self.memos.memos.get_mut(); - let types = self.types.types.read(); memos .iter_mut() - .zip(types.iter()) + .zip(self.types.types.iter()) .zip(0..) - .filter_map(|((memo, type_), index)| { + .filter_map(|((memo, (_, type_)), index)| { let memo = mem::replace(memo.atomic_memo.get_mut(), ptr::null_mut()); let memo = NonNull::new(memo)?; - Some((memo, type_.as_ref()?, index)) + Some((memo, type_.load()?, index)) }) .map(|(memo, type_, index)| { // SAFETY: We took ownership of the memo, and converted it to the correct type. @@ -336,8 +347,10 @@ impl MemoEntry { unsafe fn drop(&mut self, type_: &MemoEntryType) { if let Some(memo) = NonNull::new(mem::replace(self.atomic_memo.get_mut(), ptr::null_mut())) { - // SAFETY: Our preconditions. - mem::drop(unsafe { Box::from_raw((type_.to_dyn_fn)(memo).as_ptr()) }); + if let Some(type_) = type_.load() { + // SAFETY: Our preconditions. + mem::drop(unsafe { Box::from_raw((type_.to_dyn_fn)(memo).as_ptr()) }); + } } } } diff --git a/src/zalsa.rs b/src/zalsa.rs index 67ca5b7f3..c21a713d3 100644 --- a/src/zalsa.rs +++ b/src/zalsa.rs @@ -106,6 +106,7 @@ impl MemoIngredientIndex { MemoIngredientIndex(u as u32) } + #[inline] pub(crate) fn as_usize(self) -> usize { self.0 as usize } From 7c68655b109f9e26606c45dcd7f1c214e199a9fa Mon Sep 17 00:00:00 2001 From: Lukas Wirth Date: Sun, 20 Apr 2025 17:58:14 +0200 Subject: [PATCH 4/5] Use TypeId and allocate instead --- src/table.rs | 1 - src/table/const_type_id.rs | 105 ------------------------------------- src/table/memo.rs | 51 ++++++++---------- 3 files changed, 21 insertions(+), 136 deletions(-) delete mode 100644 src/table/const_type_id.rs diff --git a/src/table.rs b/src/table.rs index 0e4ae1db8..113b3146e 100644 --- a/src/table.rs +++ b/src/table.rs @@ -16,7 +16,6 @@ use sync::SyncTable; use crate::table::memo::{MemoTableTypes, MemoTableWithTypes, MemoTableWithTypesMut}; use crate::{Id, IngredientIndex, Revision}; -mod const_type_id; pub(crate) mod memo; pub(crate) mod sync; mod util; diff --git a/src/table/const_type_id.rs b/src/table/const_type_id.rs deleted file mode 100644 index d56a34449..000000000 --- a/src/table/const_type_id.rs +++ /dev/null @@ -1,105 +0,0 @@ -//! Forked from - -use core::any::TypeId; -use core::cmp::Ordering; -use core::fmt::{self, Debug}; -use core::hash::{Hash, Hasher}; -use core::marker::PhantomData; -use core::mem; - -/// TypeId equivalent usable in const contexts. -#[derive(Copy, Clone)] -#[repr(C)] -pub struct ConstTypeId { - type_id_fn: fn() -> TypeId, -} - -impl ConstTypeId { - /// Create a [`ConstTypeId`] for a type. - #[must_use] - pub const fn of() -> Self - where - T: ?Sized, - { - ConstTypeId { - type_id_fn: of::, - } - } - - #[inline] - fn get(self) -> TypeId { - (self.type_id_fn)() - } -} - -impl Debug for ConstTypeId { - fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { - Debug::fmt(&self.get(), formatter) - } -} - -impl PartialEq for ConstTypeId { - #[inline] - fn eq(&self, other: &Self) -> bool { - self.get() == other.get() - } -} - -impl PartialEq for ConstTypeId { - fn eq(&self, other: &TypeId) -> bool { - self.get() == *other - } -} - -impl Eq for ConstTypeId {} - -impl PartialOrd for ConstTypeId { - #[inline] - fn partial_cmp(&self, other: &Self) -> Option { - Some(Ord::cmp(self, other)) - } -} - -impl Ord for ConstTypeId { - #[inline] - fn cmp(&self, other: &Self) -> Ordering { - Ord::cmp(&self.get(), &other.get()) - } -} - -impl Hash for ConstTypeId { - fn hash(&self, state: &mut H) { - self.get().hash(state); - } -} - -/// Create a [`TypeId`] for a type. -#[must_use] -#[inline(always)] -pub fn of() -> TypeId -where - T: ?Sized, -{ - trait NonStaticAny { - fn get_type_id(&self) -> TypeId - where - Self: 'static; - } - - impl NonStaticAny for PhantomData { - #[inline(always)] - fn get_type_id(&self) -> TypeId - where - Self: 'static, - { - TypeId::of::() - } - } - - let phantom_data = PhantomData::; - // SAFETY: We only punt the lifetime to fetch the type id, we don't actually use the lifetime - // in any meaningful way. - NonStaticAny::get_type_id(unsafe { - mem::transmute::<&dyn NonStaticAny, &(dyn NonStaticAny + 'static)>(&phantom_data) - }) -} diff --git a/src/table/memo.rs b/src/table/memo.rs index d65395f6b..9e7c11abb 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -9,9 +9,7 @@ use std::{ use parking_lot::RwLock; use thin_vec::ThinVec; -use crate::{ - table::const_type_id::ConstTypeId, zalsa::MemoIngredientIndex, zalsa_local::QueryOrigin, -}; +use crate::{zalsa::MemoIngredientIndex, zalsa_local::QueryOrigin}; /// The "memo table" stores the memoized results of tracked function calls. /// Every tracked function must take a salsa struct as its first argument @@ -56,7 +54,7 @@ pub struct MemoEntryType { #[derive(Clone, Copy)] struct MemoEntryTypeData { /// The `type_id` of the erased memo type `M` - type_id: ConstTypeId, + type_id: TypeId, /// A type-coercion function for the erased memo type `M` to_dyn_fn: fn(NonNull) -> NonNull, @@ -86,15 +84,10 @@ impl MemoEntryType { #[inline] pub fn of() -> Self { Self { - data: AtomicPtr::new( - (&const { - MemoEntryTypeData { - type_id: ConstTypeId::of::(), - to_dyn_fn: Self::to_dyn_fn::(), - } - } as *const MemoEntryTypeData) - .cast_mut(), - ), + data: AtomicPtr::new(Box::into_raw(Box::new(MemoEntryTypeData { + type_id: TypeId::of::(), + to_dyn_fn: Self::to_dyn_fn::(), + }))), } } @@ -236,28 +229,26 @@ impl<'a> MemoTableWithTypes<'a> { old_entry.map(|memo| unsafe { MemoEntryType::from_dummy(memo) }) } + #[inline] pub(crate) fn get(self, memo_ingredient_index: MemoIngredientIndex) -> Option<&'a M> { - if let Some(MemoEntry { atomic_memo }) = self + let memo = self .memos .memos .read() + .get(memo_ingredient_index.as_usize())?; + let type_ = self + .types + .types .get(memo_ingredient_index.as_usize()) - { - assert_eq!( - self.types - .types - .get(memo_ingredient_index.as_usize()) - .and_then(MemoEntryType::load)? - .type_id, - TypeId::of::(), - "inconsistent type-id for `{memo_ingredient_index:?}`" - ); - let memo = NonNull::new(atomic_memo.load(Ordering::Acquire)); - // SAFETY: `type_id` check asserted above - return memo.map(|memo| unsafe { MemoEntryType::from_dummy(memo).as_ref() }); - } - - None + .and_then(MemoEntryType::load)?; + assert_eq!( + type_.type_id, + TypeId::of::(), + "inconsistent type-id for `{memo_ingredient_index:?}`" + ); + let memo = NonNull::new(memo.atomic_memo.load(Ordering::Acquire)); + // SAFETY: `type_id` check asserted above + memo.map(|memo| unsafe { MemoEntryType::from_dummy(memo).as_ref() }) } } From ea12b6eb69022d528506283d22ecf1e8bd758716 Mon Sep 17 00:00:00 2001 From: Lukas Wirth Date: Sun, 20 Apr 2025 18:44:41 +0200 Subject: [PATCH 5/5] Use `OnceLock` instead of atomic-ptr --- src/table/memo.rs | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/src/table/memo.rs b/src/table/memo.rs index 9e7c11abb..5dbe3c9e6 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -3,7 +3,10 @@ use std::{ fmt::Debug, mem, ptr::{self, NonNull}, - sync::atomic::{AtomicPtr, Ordering}, + sync::{ + atomic::{AtomicPtr, Ordering}, + OnceLock, + }, }; use parking_lot::RwLock; @@ -48,7 +51,7 @@ struct MemoEntry { } pub struct MemoEntryType { - data: AtomicPtr, + data: OnceLock, } #[derive(Clone, Copy)] @@ -84,18 +87,16 @@ impl MemoEntryType { #[inline] pub fn of() -> Self { Self { - data: AtomicPtr::new(Box::into_raw(Box::new(MemoEntryTypeData { + data: OnceLock::from(MemoEntryTypeData { type_id: TypeId::of::(), to_dyn_fn: Self::to_dyn_fn::(), - }))), + }), } } #[inline] fn load(&self) -> Option<&MemoEntryTypeData> { - // Note: Relaxed is fine as we only set the pointer once, right when we initialize it. - // SAFETY: The pointer is either null or initialized properly. - unsafe { self.data.load(Ordering::Relaxed).as_ref() } + self.data.get() } } @@ -123,14 +124,20 @@ impl MemoTableTypes { let memo_ingredient_index = memo_ingredient_index.as_usize(); while memo_ingredient_index >= self.types.count() { self.types.push(MemoEntryType { - data: AtomicPtr::default(), + data: OnceLock::new(), }); } let memo_entry_type = self.types.get(memo_ingredient_index).unwrap(); - let old = memo_entry_type + memo_entry_type .data - .swap(memo_type.data.load(Ordering::Relaxed), Ordering::AcqRel); - debug_assert!(old.is_null(), "memo type should only be set once"); + .set( + *memo_type + .data + .get() + .expect("cannot provide an empty `MemoEntryType` for `MemoEntryType::set()`"), + ) + .ok() + .expect("memo type should only be set once"); } /// # Safety @@ -231,11 +238,8 @@ impl<'a> MemoTableWithTypes<'a> { #[inline] pub(crate) fn get(self, memo_ingredient_index: MemoIngredientIndex) -> Option<&'a M> { - let memo = self - .memos - .memos - .read() - .get(memo_ingredient_index.as_usize())?; + let read = self.memos.memos.read(); + let memo = read.get(memo_ingredient_index.as_usize())?; let type_ = self .types .types