From 2466ebf1729f8bec8f8ed23b3c9af86591102061 Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Thu, 25 Dec 2025 13:12:20 +0200 Subject: [PATCH 1/2] Shrink `DatabaseKeyIndex` to 8 bytes to save memory Via the observation that: - The `IngredientIndex` can be retrieved from the page, except for tracked fns - The generation is useless, except for interneds And since they don't overlap, we can store only what needed for each. --- .../src/setup_tracked_struct.rs | 4 +- src/active_query.rs | 29 +++-- src/function.rs | 24 ++-- src/function/accumulated.rs | 2 +- src/function/diff_outputs.rs | 2 +- src/function/execute.rs | 28 +++-- src/function/maybe_changed_after.rs | 6 +- src/function/memo.rs | 4 +- src/function/specify.rs | 13 +- src/function/sync.rs | 18 ++- src/id.rs | 35 ++++-- src/input.rs | 4 +- src/interned.rs | 4 +- src/key.rs | 115 +++++++++++++++--- src/tracked_struct.rs | 21 ++-- src/zalsa.rs | 18 +-- src/zalsa_local.rs | 48 ++++---- tests/cycle_output.rs | 6 +- tests/interned-revisions.rs | 2 +- tests/memory-usage.rs | 6 +- tests/persistence.rs | 68 +++++------ 21 files changed, 287 insertions(+), 170 deletions(-) diff --git a/components/salsa-macro-rules/src/setup_tracked_struct.rs b/components/salsa-macro-rules/src/setup_tracked_struct.rs index 970cb31d6..9b94f298c 100644 --- a/components/salsa-macro-rules/src/setup_tracked_struct.rs +++ b/components/salsa-macro-rules/src/setup_tracked_struct.rs @@ -300,8 +300,8 @@ macro_rules! setup_tracked_struct { } impl $zalsa::TrackedStructInDb for $Struct<'_> { - fn database_key_index(zalsa: &$zalsa::Zalsa, id: $zalsa::Id) -> $zalsa::DatabaseKeyIndex { - $Configuration::ingredient_(zalsa).database_key_index(id) + fn ingredient_index(zalsa: &$zalsa::Zalsa, id: $zalsa::Id) -> $zalsa::IngredientIndex { + $Configuration::ingredient_(zalsa).ingredient_index() } } diff --git a/src/active_query.rs b/src/active_query.rs index c80cded3b..affa4f54d 100644 --- a/src/active_query.rs +++ b/src/active_query.rs @@ -5,7 +5,6 @@ use crate::accumulator::{ accumulated_map::{AccumulatedMap, AtomicInputAccumulatedValues, InputAccumulatedValues}, Accumulator, }; -use crate::hash::FxIndexSet; use crate::key::DatabaseKeyIndex; use crate::runtime::Stamp; use crate::sync::atomic::AtomicBool; @@ -17,6 +16,7 @@ use crate::{ Id, }; use crate::{durability::Durability, tracked_struct::Identity}; +use crate::{hash::FxIndexSet, IngredientIndex}; #[derive(Debug)] pub(crate) struct ActiveQuery { @@ -146,8 +146,9 @@ impl ActiveQuery { } /// Adds a key to our list of outputs. - pub(super) fn add_output(&mut self, key: DatabaseKeyIndex) { - self.input_outputs.insert(QueryEdge::output(key)); + pub(super) fn add_output(&mut self, key: Id, ingredient_index: IngredientIndex) { + self.input_outputs + .insert(QueryEdge::output(key, ingredient_index)); } /// True if the given key was output by this query. @@ -417,7 +418,8 @@ pub(crate) struct CompletedQuery { } struct CapturedQuery { - database_key_index: DatabaseKeyIndex, + id: Id, + ingredient_index: IngredientIndex, durability: Durability, changed_at: Revision, cycle_heads: CycleHeads, @@ -428,7 +430,8 @@ impl fmt::Debug for CapturedQuery { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut debug_struct = f.debug_struct("CapturedQuery"); debug_struct - .field("database_key_index", &self.database_key_index) + .field("id", &self.id) + .field("ingredient_index", &self.ingredient_index) .field("durability", &self.durability) .field("changed_at", &self.changed_at); if !self.cycle_heads.is_empty() { @@ -445,13 +448,17 @@ pub struct Backtrace(Box<[CapturedQuery]>); impl Backtrace { pub fn capture() -> Option { crate::with_attached_database(|db| { - db.zalsa_local().try_with_query_stack(|stack| { + let (zalsa, zalsa_local) = db.zalsas(); + zalsa_local.try_with_query_stack(|stack| { Backtrace( stack .iter() .rev() .map(|query| CapturedQuery { - database_key_index: query.database_key_index, + id: query.database_key_index.key_index(), + ingredient_index: query + .database_key_index + .ingredient_index_with_zalsa(zalsa), durability: query.durability, changed_at: query.changed_at, cycle_heads: query.cycle_heads.clone(), @@ -486,7 +493,8 @@ impl fmt::Display for Backtrace { for ( idx, &CapturedQuery { - database_key_index, + id, + ingredient_index, durability, changed_at, ref cycle_heads, @@ -494,6 +502,7 @@ impl fmt::Display for Backtrace { }, ) in self.0.iter().enumerate() { + let database_key_index = DatabaseKeyIndex::new_non_interned(ingredient_index, id); write!(fmt, "{idx:>4}: {database_key_index:?}")?; if full { write!(fmt, " -> ({changed_at:?}, {durability:#?}")?; @@ -504,9 +513,7 @@ impl fmt::Display for Backtrace { } writeln!(fmt)?; crate::attach::with_attached_database(|db| { - let ingredient = db - .zalsa() - .lookup_ingredient(database_key_index.ingredient_index()); + let ingredient = db.zalsa().lookup_ingredient(ingredient_index); let loc = ingredient.location(); writeln!(fmt, "{indent}at {}:{}", loc.file, loc.line)?; if !cycle_heads.is_empty() { diff --git a/src/function.rs b/src/function.rs index f7f302727..b27b3cb70 100644 --- a/src/function.rs +++ b/src/function.rs @@ -230,7 +230,7 @@ where #[inline] pub fn database_key_index(&self, key: Id) -> DatabaseKeyIndex { - DatabaseKeyIndex::new(self.index, key) + DatabaseKeyIndex::new_non_interned(self.index, key) } pub fn set_capacity(&mut self, capacity: usize) { @@ -350,7 +350,7 @@ where continue; } - let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index()); + let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index_with_zalsa(zalsa)); dependency.collect_minimum_serialized_edges( zalsa, *edge, @@ -527,7 +527,9 @@ where // We only serialize the query if there are any memos associated with it. for entry in as SalsaStructInDb>::entries(zalsa) { - let memo_ingredient_index = self.memo_ingredient_indices.get(entry.ingredient_index()); + let memo_ingredient_index = self + .memo_ingredient_indices + .get(entry.ingredient_index_with_zalsa(zalsa)); let memo = self.get_memo_from_table_for(zalsa, entry.key_index(), memo_ingredient_index); @@ -614,7 +616,7 @@ mod persistence { .filter(|entry| { let memo_ingredient_index = ingredient .memo_ingredient_indices - .get(entry.ingredient_index()); + .get(entry.ingredient_index_with_zalsa(zalsa)); let memo = ingredient.get_memo_from_table_for( zalsa, @@ -634,7 +636,7 @@ mod persistence { for entry in as SalsaStructInDb>::entries(zalsa) { let memo_ingredient_index = ingredient .memo_ingredient_indices - .get(entry.ingredient_index()); + .get(entry.ingredient_index_with_zalsa(zalsa)); let memo = ingredient.get_memo_from_table_for( zalsa, @@ -666,14 +668,18 @@ mod persistence { QueryOrigin::derived_untracked(flattened_edges.drain(..).collect()) } QueryOriginRef::Assigned(key) => { - let dependency = zalsa.lookup_ingredient(key.ingredient_index()); + let dependency = + zalsa.lookup_ingredient(key.ingredient_index_with_zalsa(zalsa)); assert!( dependency.is_persistable(), "specified query `{}` must be persistable", dependency.debug_name() ); - QueryOrigin::assigned(key) + QueryOrigin::assigned( + key.key_index(), + key.ingredient_index_with_zalsa(zalsa), + ) } QueryOriginRef::FixpointInitial => unreachable!( "`should_serialize` returns `false` for provisional queries" @@ -685,7 +691,7 @@ mod persistence { // TODO: Group structs by ingredient index into a nested map. let key = format!( "{}:{}", - entry.ingredient_index().as_u32(), + entry.ingredient_index_with_zalsa(self.zalsa).as_u32(), entry.key_index().as_bits() ); @@ -707,7 +713,7 @@ mod persistence { flattened_edges: &mut FxIndexSet, ) { for &edge in edges { - let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index()); + let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index_with_zalsa(zalsa)); if dependency.is_persistable() { // If the dependency will be serialized, we can serialize the edge directly. diff --git a/src/function/accumulated.rs b/src/function/accumulated.rs index a65804e64..13b234885 100644 --- a/src/function/accumulated.rs +++ b/src/function/accumulated.rs @@ -51,7 +51,7 @@ where continue; } - let ingredient = zalsa.lookup_ingredient(k.ingredient_index()); + let ingredient = zalsa.lookup_ingredient(k.ingredient_index_with_zalsa(zalsa)); // Extend `output` with any values accumulated by `k`. // SAFETY: `db` owns the `ingredient` let (accumulated_map, input) = diff --git a/src/function/diff_outputs.rs b/src/function/diff_outputs.rs index 37d71d41d..a2dc41f10 100644 --- a/src/function/diff_outputs.rs +++ b/src/function/diff_outputs.rs @@ -38,7 +38,7 @@ fn diff_outputs_on_revision( // Note that tracked structs are not stored as direct query outputs, but they are still outputs // that need to be reported as stale. for (identity, id) in &completed_query.stale_tracked_structs { - let output = DatabaseKeyIndex::new(identity.ingredient_index(), *id); + let output = DatabaseKeyIndex::new_non_interned(identity.ingredient_index(), *id); report_stale_output(zalsa, key, output); } diff --git a/src/function/execute.rs b/src/function/execute.rs index a4dbe4986..fb8623118 100644 --- a/src/function/execute.rs +++ b/src/function/execute.rs @@ -44,7 +44,7 @@ where let database_key_index = claim_guard.database_key_index(); let zalsa = claim_guard.zalsa(); - let id = database_key_index.key_index(); + let id = claim_guard.key_index(); let memo_ingredient_index = self.memo_ingredient_index(zalsa, id); crate::tracing::info!("{:?}: executing query", database_key_index); @@ -61,6 +61,7 @@ where db, zalsa, zalsa_local.push_query(database_key_index, IterationCount::initial()), + id, opt_old_memo, ); (new_value, active_query.pop()) @@ -70,6 +71,7 @@ where db, zalsa, zalsa_local.push_query(database_key_index, IterationCount::initial()), + id, opt_old_memo, ); @@ -166,7 +168,7 @@ where let database_key_index = claim_guard.database_key_index(); let zalsa = claim_guard.zalsa(); - let id = database_key_index.key_index(); + let id = claim_guard.key_index(); // Our provisional value from the previous iteration, when doing fixpoint iteration. // This is different from `opt_old_memo` which might be from a different revision. @@ -219,6 +221,7 @@ where db, zalsa, active_query, + id, last_provisional_memo_opt.or(opt_old_memo), ); @@ -377,6 +380,7 @@ where db: &'db C::DbView, zalsa: &'db Zalsa, active_query: ActiveQueryGuard<'db>, + key_index: Id, opt_old_memo: Option<&Memo<'db, C>>, ) -> (C::Output<'db>, ActiveQueryGuard<'db>) { if let Some(old_memo) = opt_old_memo { @@ -398,10 +402,7 @@ where // Query was not previously executed, or value is potentially // stale, or value is absent. Let's execute! - let new_value = C::execute( - db, - C::id_to_input(zalsa, active_query.database_key_index.key_index()), - ); + let new_value = C::execute(db, C::id_to_input(zalsa, key_index)); (new_value, active_query) } @@ -492,7 +493,8 @@ fn outer_cycle( cycle_heads .iter_not_eq(current_key) .rfind(|head| { - let ingredient = zalsa.lookup_ingredient(head.database_key_index.ingredient_index()); + let ingredient = + zalsa.lookup_ingredient(head.database_key_index.ingredient_index_with_zalsa(zalsa)); matches!( ingredient.wait_for(zalsa, head.database_key_index.key_index()), @@ -531,7 +533,8 @@ fn collect_all_cycle_heads( max_iteration_count = max_iteration_count.max(head.iteration_count.load()); depends_on_self |= head.database_key_index == database_key_index; - let ingredient = zalsa.lookup_ingredient(head.database_key_index.ingredient_index()); + let ingredient = + zalsa.lookup_ingredient(head.database_key_index.ingredient_index_with_zalsa(zalsa)); let provisional_status = ingredient .provisional_status(zalsa, head.database_key_index.key_index()) @@ -654,7 +657,8 @@ fn try_complete_cycle_head( let converged = this_converged && cycle_heads.iter_not_eq(me).all(|head| { let database_key_index = head.database_key_index; - let ingredient = zalsa.lookup_ingredient(database_key_index.ingredient_index()); + let ingredient = + zalsa.lookup_ingredient(database_key_index.ingredient_index_with_zalsa(zalsa)); let converged = ingredient.cycle_converged(zalsa, database_key_index.key_index()); @@ -673,7 +677,8 @@ fn try_complete_cycle_head( // Set the nested cycles as verified. This is necessary because // `validate_provisional` doesn't follow cycle heads recursively (and the memos now depend on all cycle heads). for head in cycle_heads.iter_not_eq(me) { - let ingredient = zalsa.lookup_ingredient(head.database_key_index.ingredient_index()); + let ingredient = + zalsa.lookup_ingredient(head.database_key_index.ingredient_index_with_zalsa(zalsa)); ingredient.finalize_cycle_head(zalsa, head.database_key_index.key_index()); } @@ -700,7 +705,8 @@ fn try_complete_cycle_head( // Update the iteration count of nested cycles. for head in cycle_heads.iter_not_eq(me) { - let ingredient = zalsa.lookup_ingredient(head.database_key_index.ingredient_index()); + let ingredient = + zalsa.lookup_ingredient(head.database_key_index.ingredient_index_with_zalsa(zalsa)); ingredient.set_cycle_iteration_count( zalsa, diff --git a/src/function/maybe_changed_after.rs b/src/function/maybe_changed_after.rs index 6ea17b13f..aa4101937 100644 --- a/src/function/maybe_changed_after.rs +++ b/src/function/maybe_changed_after.rs @@ -631,7 +631,11 @@ fn validate_provisional( for cycle_head in cycle_heads { // Test if our cycle heads (with the same revision) are now finalized. let Some(kind) = zalsa - .lookup_ingredient(cycle_head.database_key_index.ingredient_index()) + .lookup_ingredient( + cycle_head + .database_key_index + .ingredient_index_with_zalsa(zalsa), + ) .provisional_status(zalsa, cycle_head.database_key_index.key_index()) else { return false; diff --git a/src/function/memo.rs b/src/function/memo.rs index 234829cb1..f2dc4b7d7 100644 --- a/src/function/memo.rs +++ b/src/function/memo.rs @@ -260,7 +260,7 @@ where } for (identity, id) in self.revisions.tracked_struct_ids() { - let key = DatabaseKeyIndex::new(identity.ingredient_index(), *id); + let key = DatabaseKeyIndex::new_non_interned(identity.ingredient_index(), *id); key.remove_stale_output(zalsa, executor); } } @@ -447,7 +447,7 @@ impl<'me> Iterator for TryClaimCycleHeadsIter<'me> { let head_key_index = head_database_key.key_index(); let ingredient = self .zalsa - .lookup_ingredient(head_database_key.ingredient_index()); + .lookup_ingredient(head_database_key.ingredient_index_with_zalsa(self.zalsa)); match ingredient.wait_for(self.zalsa, head_key_index) { WaitForResult::Cycle { .. } => { diff --git a/src/function/specify.rs b/src/function/specify.rs index 99539d375..6f1490e80 100644 --- a/src/function/specify.rs +++ b/src/function/specify.rs @@ -39,10 +39,11 @@ where // * Q4 invokes Q2 and then Q1 // // Now, if We invoke Q3 first, We get one result for Q2, but if We invoke Q4 first, We get a different value. That's no good. - let database_key_index = >::database_key_index(zalsa, key); - if !zalsa_local.is_tracked_struct_of_active_query(database_key_index) { + let ingredient_index = >::ingredient_index(zalsa, key); + if !zalsa_local.is_tracked_struct_of_active_query(key, ingredient_index) { panic!("can only use `specify` on salsa structs created during the current tracked fn"); } + let database_key_index = DatabaseKeyIndex::new_non_interned(ingredient_index, key); // Subtle: we treat the "input" to a set query as if it were // volatile. @@ -68,7 +69,10 @@ where revisions: QueryRevisions { changed_at: current_deps.changed_at, durability: current_deps.durability, - origin: QueryOrigin::assigned(active_query_key), + origin: QueryOrigin::assigned( + active_query_key.key_index(), + active_query_key.ingredient_index_with_zalsa(zalsa), + ), #[cfg(feature = "accumulator")] accumulated_inputs: Default::default(), verified_final: AtomicBool::new(true), @@ -102,8 +106,7 @@ where self.insert_memo(zalsa, key, memo, memo_ingredient_index); // Record that the current query *specified* a value for this cell. - let database_key_index = self.database_key_index(key); - zalsa_local.add_output(database_key_index); + zalsa_local.add_output(key, self.index); } /// Invoked when the query `executor` has been validated as having green inputs diff --git a/src/function/sync.rs b/src/function/sync.rs index c9a74a307..b0d1596ec 100644 --- a/src/function/sync.rs +++ b/src/function/sync.rs @@ -97,7 +97,7 @@ impl SyncTable { // not to gate future atomic reads. *anyone_waiting = true; match zalsa.runtime().block( - DatabaseKeyIndex::new(self.ingredient, key_index), + DatabaseKeyIndex::new_non_interned(self.ingredient, key_index), id, write, ) { @@ -155,7 +155,7 @@ impl SyncTable { // not to gate future atomic reads. *anyone_waiting = true; match zalsa.runtime().block( - DatabaseKeyIndex::new(self.ingredient, key_index), + DatabaseKeyIndex::new_non_interned(self.ingredient, key_index), id, write, ) { @@ -176,7 +176,7 @@ impl SyncTable { reentrant: Reentrancy, ) -> Result, Box>> { let key_index = *entry.key(); - let database_key_index = DatabaseKeyIndex::new(self.ingredient, key_index); + let database_key_index = DatabaseKeyIndex::new_non_interned(self.ingredient, key_index); let thread_id = thread::current().id(); match zalsa @@ -230,7 +230,7 @@ impl SyncTable { reentrant: Reentrancy, ) -> Result, Box>> { let key_index = *entry.key(); - let database_key_index = DatabaseKeyIndex::new(self.ingredient, key_index); + let database_key_index = DatabaseKeyIndex::new_non_interned(self.ingredient, key_index); let thread_id = thread::current().id(); match zalsa @@ -303,7 +303,11 @@ impl<'me> ClaimGuard<'me> { } pub(crate) const fn database_key_index(&self) -> DatabaseKeyIndex { - DatabaseKeyIndex::new(self.sync_table.ingredient, self.key_index) + DatabaseKeyIndex::new_non_interned(self.sync_table.ingredient, self.key_index) + } + + pub(crate) fn key_index(&self) -> Id { + self.key_index } pub(crate) fn set_release_mode(&mut self, mode: ReleaseMode) { @@ -370,7 +374,9 @@ impl<'me> ClaimGuard<'me> { #[cold] #[inline(never)] pub(crate) fn transfer(&self, new_owner: DatabaseKeyIndex) -> bool { - let owner_ingredient = self.zalsa.lookup_ingredient(new_owner.ingredient_index()); + let owner_ingredient = self + .zalsa + .lookup_ingredient(new_owner.ingredient_index_with_zalsa(self.zalsa)); // Get the owning thread of `new_owner`. // The thread id is guaranteed to not be stale because `new_owner` must be blocked on `self_key` diff --git a/src/id.rs b/src/id.rs index bc2565410..e38863760 100644 --- a/src/id.rs +++ b/src/id.rs @@ -20,6 +20,11 @@ use crate::zalsa::Zalsa; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct Id { index: NonZeroU32, + /// `generation` is shifted by one bit, to allow bit tagging for `DatabaseKeyIndex`. + /// The LSB is always 1. + /// + /// We chose to tag the LSB and not the MSB because then the generation can be treated + /// as a continuous integer, without special treatment. generation: u32, } @@ -46,7 +51,7 @@ impl Id { Id { // SAFETY: Caller obligation. index: unsafe { NonZeroU32::new_unchecked(index + 1) }, - generation: 0, + generation: 0b1, } } @@ -65,6 +70,7 @@ impl Id { // SAFETY: Caller obligation. let index = unsafe { NonZeroU32::new_unchecked(bits as u32) }; let generation = (bits >> 32) as u32; + debug_assert!((generation & 0b1) == 0b1); Id { index, generation } } @@ -78,6 +84,7 @@ impl Id { pub const fn from_bits(bits: u64) -> Self { let index = NonZeroU32::new(bits as u32).expect("attempted to create invalid `Id`"); let generation = (bits >> 32) as u32; + debug_assert!((generation & 0b1) == 0b1); Id { index, generation } } @@ -94,9 +101,10 @@ impl Id { /// is `u32::MAX`. #[inline] pub fn next_generation(self) -> Option { - self.generation() - .checked_add(1) - .map(|generation| self.with_generation(generation)) + self.generation.checked_add(1).map(|generation| Id { + index: self.index, + generation: generation | 0b1, + }) } /// Mark the `Id` with a generation. @@ -108,10 +116,22 @@ impl Id { pub const fn with_generation(self, generation: u32) -> Id { Id { index: self.index, - generation, + generation: (generation << 1) | 0b1, } } + #[inline] + pub(crate) const fn index_nonzero(self) -> NonZeroU32 { + self.index + } + + #[inline] + pub(crate) fn from_raw_parts(index: NonZeroU32, generation: u32) -> Self { + debug_assert!(index.get() <= Self::MAX_U32); + debug_assert!((generation & 0b1) == 0b1); + Self { index, generation } + } + /// Return the index portion of this `Id`. #[inline] pub const fn index(self) -> u32 { @@ -154,10 +174,11 @@ impl<'de> serde::Deserialize<'de> for Id { impl Debug for Id { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - if self.generation() == 0 { + let generation = self.generation() >> 1; + if generation == 0 { write!(f, "Id({:x})", self.index()) } else { - write!(f, "Id({:x}g{:x})", self.index(), self.generation()) + write!(f, "Id({:x}g{:x})", self.index(), generation) } } } diff --git a/src/input.rs b/src/input.rs index 0ce0f5c87..45adae9e8 100644 --- a/src/input.rs +++ b/src/input.rs @@ -140,7 +140,7 @@ impl IngredientImpl { } pub fn database_key_index(&self, id: Id) -> DatabaseKeyIndex { - DatabaseKeyIndex::new(self.ingredient_index, id) + DatabaseKeyIndex::new_non_interned(self.ingredient_index, id) } pub fn new_input( @@ -230,7 +230,7 @@ impl IngredientImpl { let durability = value.durabilities[field_index]; let revision = value.revisions[field_index]; zalsa_local.report_tracked_read_simple( - DatabaseKeyIndex::new(field_ingredient_index, id), + DatabaseKeyIndex::new_non_interned(field_ingredient_index, id), durability, revision, ); diff --git a/src/interned.rs b/src/interned.rs index 544a8d0ee..3b7e237c2 100644 --- a/src/interned.rs +++ b/src/interned.rs @@ -712,7 +712,7 @@ where 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_non_interned(ingredient_index, id); zalsa.event(&|| Event::new(EventKind::DidDiscard { key: executor })); @@ -766,7 +766,7 @@ where /// Returns the database key index for an interned value with the given id. #[inline] pub fn database_key_index(&self, id: Id) -> DatabaseKeyIndex { - DatabaseKeyIndex::new(self.ingredient_index, id) + DatabaseKeyIndex::new_interned(self.ingredient_index, id) } /// Lookup the data for an interned value based on its ID. diff --git a/src/key.rs b/src/key.rs index 364015756..4f9cc1186 100644 --- a/src/key.rs +++ b/src/key.rs @@ -1,8 +1,9 @@ use std::fmt; +use std::num::NonZeroU32; use crate::function::{VerifyCycleHeads, VerifyResult}; use crate::zalsa::{IngredientIndex, Zalsa}; -use crate::Id; +use crate::{Database, Id}; // ANCHOR: DatabaseKeyIndex /// An integer that uniquely identifies a particular query instance within the @@ -11,26 +12,95 @@ use crate::Id; /// only for inserting into maps and the like. #[derive(Copy, Clone, PartialEq, Eq, Hash)] pub struct DatabaseKeyIndex { - key_index: Id, - ingredient_index: IngredientIndex, + index: NonZeroU32, + /// `DatabaseKeyIndex` is stored *a lot*, as every query dependency stores one. Therefore + /// we want to make it compact. + /// + /// The ingredient index is technically needed only for tracked fns - other things can + /// grab it from the page, whose index is stored in `index`. On the other hand, + /// only interned structs need a generation - for GC. So we store like the following: + /// + /// If this is an interned, this stores the generation (left-shifted by a bit and ORed with 0b1, + /// as every generation is). If not, the LSB is 0, and the rest of the bits store the + /// `IngredientIndex` shifted left by 1 bit. + generation_or_ingredient_index: u32, } // ANCHOR_END: DatabaseKeyIndex impl DatabaseKeyIndex { + const INGREDIENT_INDEX_SHIFT: u32 = 2; + + #[inline] + pub(crate) const fn new_interned(_ingredient_index: IngredientIndex, key_index: Id) -> Self { + Self { + index: key_index.index_nonzero(), + generation_or_ingredient_index: key_index.generation(), + } + } + #[inline] - pub(crate) const fn new(ingredient_index: IngredientIndex, key_index: Id) -> Self { + pub(crate) const fn new_non_interned(ingredient_index: IngredientIndex, key_index: Id) -> Self { Self { - key_index, - ingredient_index, + index: key_index.index_nonzero(), + generation_or_ingredient_index: ingredient_index.as_u32() + << Self::INGREDIENT_INDEX_SHIFT, + } + } + + #[inline] + pub(crate) fn new_non_interned_with_tag( + ingredient_index: IngredientIndex, + key_index: Id, + ) -> Self { + let mut result = Self::new_non_interned(ingredient_index, key_index); + result.generation_or_ingredient_index |= 0b10; + result + } + + #[inline] + pub(crate) fn has_tag(self) -> bool { + // The LSB is 0, meaning a non-interned, and the second LSB is 1, meaning the tag is on. + (self.generation_or_ingredient_index & 0b11) == 0b10 + } + + #[inline] + pub(crate) fn ingredient_index_with_zalsa(self, zalsa: &Zalsa) -> IngredientIndex { + if self.is_interned() { + zalsa.ingredient_index(self.key_index()) + } else { + IngredientIndex::new( + self.generation_or_ingredient_index >> Self::INGREDIENT_INDEX_SHIFT, + ) } } - pub const fn ingredient_index(self) -> IngredientIndex { - self.ingredient_index + #[inline] + pub fn ingredient_index(self, db: &dyn Database) -> IngredientIndex { + self.ingredient_index_with_zalsa(db.zalsa()) } - pub const fn key_index(self) -> Id { - self.key_index + #[track_caller] + #[cfg(feature = "persistence")] + pub(crate) fn ingredient_index_assert_non_interned(self) -> IngredientIndex { + assert!(!self.is_interned()); + IngredientIndex::new(self.generation_or_ingredient_index >> Self::INGREDIENT_INDEX_SHIFT) + } + + /// **Warning:** For tracked fns on interned structs, the generation is *permanently lost* once stored + /// in a `DatabaseKeyIndex`. Usually it is not a problem to bring an accurate [`Id`] from a different place, + /// and where it is hard, it is not needed. However you need to be careful when handling this. + pub fn key_index(self) -> Id { + let generation = if self.is_interned() { + self.generation_or_ingredient_index + } else { + 1 + }; + Id::from_raw_parts(self.index, generation) + } + + #[inline] + fn is_interned(self) -> bool { + (self.generation_or_ingredient_index & 0b1) == 0b1 } pub(crate) fn maybe_changed_after( @@ -45,14 +115,14 @@ impl DatabaseKeyIndex { // here, `db` has to be either the correct type already, or a subtype (as far as trait // hierarchy is concerned) zalsa - .lookup_ingredient(self.ingredient_index()) + .lookup_ingredient(self.ingredient_index_with_zalsa(zalsa)) .maybe_changed_after(zalsa, db, self.key_index(), last_verified_at, cycle_heads) } } pub(crate) fn remove_stale_output(&self, zalsa: &Zalsa, executor: DatabaseKeyIndex) { zalsa - .lookup_ingredient(self.ingredient_index()) + .lookup_ingredient(self.ingredient_index_with_zalsa(zalsa)) .remove_stale_output(zalsa, executor, self.key_index()) } @@ -62,7 +132,7 @@ impl DatabaseKeyIndex { database_key_index: DatabaseKeyIndex, ) { zalsa - .lookup_ingredient(self.ingredient_index()) + .lookup_ingredient(self.ingredient_index_with_zalsa(zalsa)) .mark_validated_output(zalsa, database_key_index, self.key_index()) } } @@ -73,7 +143,10 @@ impl serde::Serialize for DatabaseKeyIndex { where S: serde::Serializer, { - serde::Serialize::serialize(&(self.key_index, self.ingredient_index), serializer) + serde::Serialize::serialize( + &(self.index, self.generation_or_ingredient_index), + serializer, + ) } } @@ -83,11 +156,12 @@ impl<'de> serde::Deserialize<'de> for DatabaseKeyIndex { where D: serde::Deserializer<'de>, { - let (key_index, ingredient_index) = serde::Deserialize::deserialize(deserializer)?; + let (index, generation_or_ingredient_index) = + serde::Deserialize::deserialize(deserializer)?; Ok(DatabaseKeyIndex { - key_index, - ingredient_index, + index, + generation_or_ingredient_index, }) } } @@ -95,13 +169,14 @@ impl<'de> serde::Deserialize<'de> for DatabaseKeyIndex { impl fmt::Debug for DatabaseKeyIndex { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { crate::attach::with_attached_database(|db| { - let ingredient = db.zalsa().lookup_ingredient(self.ingredient_index()); + let zalsa = db.zalsa(); + let ingredient = zalsa.lookup_ingredient(self.ingredient_index_with_zalsa(zalsa)); ingredient.fmt_index(self.key_index(), f) }) .unwrap_or_else(|| { f.debug_tuple("DatabaseKeyIndex") - .field(&self.ingredient_index()) - .field(&self.key_index()) + .field(&self.index) + .field(&self.generation_or_ingredient_index) .finish() }) } diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs index 78ebef5d9..fefc7b4a3 100644 --- a/src/tracked_struct.rs +++ b/src/tracked_struct.rs @@ -169,8 +169,7 @@ impl Jar for JarImpl { } pub trait TrackedStructInDb: SalsaStructInDb { - /// Converts the identifier for this tracked struct into a `DatabaseKeyIndex`. - fn database_key_index(zalsa: &Zalsa, id: Id) -> DatabaseKeyIndex; + fn ingredient_index(zalsa: &Zalsa, id: Id) -> IngredientIndex; } /// Created for each tracked struct. @@ -311,13 +310,10 @@ impl IdentityMap { } /// Returns `true` if the given tracked struct key was created in the current query execution. - pub(crate) fn is_active(&self, key: DatabaseKeyIndex) -> bool { + pub(crate) fn is_active(&self, id: Id, ingredient: IngredientIndex) -> bool { self.table .iter() - .find(|entry| { - entry.id == key.key_index() - && entry.identity.ingredient_index() == key.ingredient_index() - }) + .find(|entry| entry.id == id && entry.identity.ingredient_index() == ingredient) .is_some_and(|entry| entry.active) } @@ -480,9 +476,14 @@ where } } + #[inline] + pub fn ingredient_index(&self) -> IngredientIndex { + self.ingredient_index + } + /// Returns the database key index for a tracked struct with the given id. pub fn database_key_index(&self, id: Id) -> DatabaseKeyIndex { - DatabaseKeyIndex::new(self.ingredient_index, id) + DatabaseKeyIndex::new_non_interned(self.ingredient_index, id) } pub fn new_struct<'db>( @@ -825,7 +826,7 @@ where 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_non_interned(ingredient_index, id); zalsa.event(&|| Event::new(EventKind::DidDiscard { key: executor })); @@ -871,7 +872,7 @@ where let field_changed_at = data.revisions[relative_tracked_index]; zalsa_local.report_tracked_read_simple( - DatabaseKeyIndex::new(field_ingredient_index, id), + DatabaseKeyIndex::new_non_interned(field_ingredient_index, id), data.durability, field_changed_at, ); diff --git a/src/zalsa.rs b/src/zalsa.rs index f2f520b49..67bdd4d3c 100644 --- a/src/zalsa.rs +++ b/src/zalsa.rs @@ -74,9 +74,7 @@ pub struct IngredientIndex(u32); impl IngredientIndex { /// The maximum supported ingredient index. - /// - /// This reserves one bit for an optional tag. - const MAX_INDEX: u32 = 0x7FFF_FFFF; + const MAX_INDEX: u32 = 0x3FFF_FFFF; /// Create an ingredient index from a `u32`. pub(crate) fn new(v: u32) -> Self { @@ -95,25 +93,13 @@ impl IngredientIndex { } /// Convert the ingredient index back into a `u32`. - pub(crate) fn as_u32(self) -> u32 { + pub(crate) const fn as_u32(self) -> u32 { self.0 } pub fn successor(self, index: usize) -> Self { IngredientIndex(self.0 + 1 + index as u32) } - - /// Returns a new `IngredientIndex` with the tag bit set to the provided value. - pub(crate) fn with_tag(mut self, tag: bool) -> IngredientIndex { - self.0 &= Self::MAX_INDEX; - self.0 |= (tag as u32) << 31; - self - } - - /// Returns the value of the tag bit. - pub(crate) fn tag(self) -> bool { - self.0 & !Self::MAX_INDEX != 0 - } } /// A special secondary index *just* for ingredients that attach diff --git a/src/zalsa_local.rs b/src/zalsa_local.rs index 8f0239e56..89fd11708 100644 --- a/src/zalsa_local.rs +++ b/src/zalsa_local.rs @@ -234,25 +234,29 @@ impl ZalsaLocal { } /// Add an output to the current query's list of dependencies - pub(crate) fn add_output(&self, entity: DatabaseKeyIndex) { + pub(crate) fn add_output(&self, key: Id, ingredient_index: IngredientIndex) { // SAFETY: We do not access the query stack reentrantly. unsafe { self.with_query_stack_unchecked_mut(|stack| { if let Some(top_query) = stack.last_mut() { - top_query.add_output(entity) + top_query.add_output(key, ingredient_index) } }) } } /// Check whether `entity` is a tracked struct that was created by the currently active query (if any) - pub(crate) fn is_tracked_struct_of_active_query(&self, entity: DatabaseKeyIndex) -> bool { + pub(crate) fn is_tracked_struct_of_active_query( + &self, + id: Id, + ingredient: IngredientIndex, + ) -> bool { // SAFETY: We do not access the query stack reentrantly. unsafe { self.with_query_stack_unchecked_mut(|stack| { - stack - .last_mut() - .is_some_and(|top_query| top_query.tracked_struct_ids().is_active(entity)) + stack.last_mut().is_some_and(|top_query| { + top_query.tracked_struct_ids().is_active(id, ingredient) + }) }) } } @@ -963,13 +967,11 @@ impl QueryOrigin { } /// Create a query origin of type `QueryOriginKind::Assigned`, with the given key. - pub fn assigned(key: DatabaseKeyIndex) -> QueryOrigin { + pub fn assigned(key: Id, ingredient_index: IngredientIndex) -> QueryOrigin { QueryOrigin { kind: QueryOriginKind::Assigned, - metadata: key.ingredient_index().as_u32(), - data: QueryOriginData { - index: key.key_index(), - }, + metadata: ingredient_index.as_u32(), + data: QueryOriginData { index: key }, } } @@ -984,7 +986,10 @@ impl QueryOrigin { // is `QueryOriginKind::Assigned`. let ingredient_index = unsafe { IngredientIndex::new_unchecked(self.metadata) }; - QueryOriginRef::Assigned(DatabaseKeyIndex::new(ingredient_index, index)) + QueryOriginRef::Assigned(DatabaseKeyIndex::new_non_interned( + ingredient_index, + index, + )) } QueryOriginKind::Derived => { @@ -1046,7 +1051,10 @@ impl<'de> serde::Deserialize<'de> for QueryOrigin { } Ok(match QueryOriginOwned::deserialize(deserializer)? { - QueryOriginOwned::Assigned(key) => QueryOrigin::assigned(key), + QueryOriginOwned::Assigned(key) => { + // It's safe to assert_non_interned since only specify can produce `Assigned`, and it will not contain an interned. + QueryOrigin::assigned(key.key_index(), key.ingredient_index_assert_non_interned()) + } QueryOriginOwned::Derived(edges) => QueryOrigin::derived(edges), QueryOriginOwned::DerivedUntracked(edges) => QueryOrigin::derived_untracked(edges), QueryOriginOwned::FixpointInitial => QueryOrigin::fixpoint_initial(), @@ -1106,26 +1114,20 @@ impl QueryEdge { } /// Create an output query edge with the given index. - pub fn output(key: DatabaseKeyIndex) -> QueryEdge { - let ingredient_index = key.ingredient_index().with_tag(true); - + pub fn output(key: Id, ingredient_index: IngredientIndex) -> QueryEdge { Self { - key: DatabaseKeyIndex::new(ingredient_index, key.key_index()), + key: DatabaseKeyIndex::new_non_interned_with_tag(ingredient_index, key), } } /// Return the key of this query edge. pub fn key(self) -> DatabaseKeyIndex { - // Clear the tag to restore the original index. - DatabaseKeyIndex::new( - self.key.ingredient_index().with_tag(false), - self.key.key_index(), - ) + self.key } /// Returns the kind of this query edge. pub fn kind(self) -> QueryEdgeKind { - if self.key.ingredient_index().tag() { + if self.key.has_tag() { QueryEdgeKind::Output(self.key()) } else { QueryEdgeKind::Input(self.key()) diff --git a/tests/cycle_output.rs b/tests/cycle_output.rs index c4a9384e0..bfe9ee1a3 100644 --- a/tests/cycle_output.rs +++ b/tests/cycle_output.rs @@ -188,12 +188,12 @@ fn revalidate_with_change_after_output_read() { "salsa_event(DidDiscard { key: read_value(Id(403)) })", "salsa_event(WillIterateCycle { database_key: query_b(Id(0)), iteration_count: IterationCount(1) })", "salsa_event(WillExecute { database_key: query_a(Id(0)) })", - "salsa_event(WillExecute { database_key: read_value(Id(401g1)) })", + "salsa_event(WillExecute { database_key: read_value(Id(401)) })", "salsa_event(WillIterateCycle { database_key: query_b(Id(0)), iteration_count: IterationCount(2) })", "salsa_event(WillExecute { database_key: query_a(Id(0)) })", - "salsa_event(WillExecute { database_key: read_value(Id(402g1)) })", + "salsa_event(WillExecute { database_key: read_value(Id(402)) })", "salsa_event(WillIterateCycle { database_key: query_b(Id(0)), iteration_count: IterationCount(3) })", "salsa_event(WillExecute { database_key: query_a(Id(0)) })", - "salsa_event(WillExecute { database_key: read_value(Id(403g1)) })", + "salsa_event(WillExecute { database_key: read_value(Id(403)) })", ]"#]]); } diff --git a/tests/interned-revisions.rs b/tests/interned-revisions.rs index bef1db61c..d4e83f8a6 100644 --- a/tests/interned-revisions.rs +++ b/tests/interned-revisions.rs @@ -160,7 +160,7 @@ fn test_immortal() { input.set_field1(&mut db).to(i); let result = function(&db, input); assert_eq!(result.field1(&db).0, i); - assert_eq!(salsa::plumbing::AsId::as_id(&result).generation(), 0); + assert_eq!(salsa::plumbing::AsId::as_id(&result).generation(), 1); } } diff --git a/tests/memory-usage.rs b/tests/memory-usage.rs index a3dbac80f..b64f35297 100644 --- a/tests/memory-usage.rs +++ b/tests/memory-usage.rs @@ -136,7 +136,7 @@ fn test() { IngredientInfo { debug_name: "memory_usage::MyInterned<'_>", count: 3, - size_of_metadata: 192, + size_of_metadata: 168, size_of_fields: 24, heap_size_of_fields: None, }, @@ -168,7 +168,7 @@ fn test() { IngredientInfo { debug_name: "memory_usage::MyTracked<'_>", count: 2, - size_of_metadata: 168, + size_of_metadata: 160, size_of_fields: 16, heap_size_of_fields: None, }, @@ -178,7 +178,7 @@ fn test() { IngredientInfo { debug_name: "(memory_usage::MyTracked<'_>, memory_usage::MyTracked<'_>)", count: 1, - size_of_metadata: 108, + size_of_metadata: 104, size_of_fields: 16, heap_size_of_fields: None, }, diff --git a/tests/persistence.rs b/tests/persistence.rs index b476bd192..e036e748d 100644 --- a/tests/persistence.rs +++ b/tests/persistence.rs @@ -63,7 +63,7 @@ fn everything() { }, "ingredients": { "0": { - "1": { + "4294967297": { "durabilities": [ 0 ], @@ -74,7 +74,7 @@ fn everything() { 1 ] }, - "2": { + "4294967298": { "durabilities": [ 0 ], @@ -113,7 +113,7 @@ fn everything() { }, "ingredients": { "0": { - "1": { + "4294967297": { "durabilities": [ 0 ], @@ -124,7 +124,7 @@ fn everything() { 1 ] }, - "2": { + "4294967298": { "durabilities": [ 0 ], @@ -135,7 +135,7 @@ fn everything() { 2 ] }, - "3": { + "4294967299": { "durabilities": [ 0 ], @@ -146,7 +146,7 @@ fn everything() { 1 ] }, - "4": { + "4294967300": { "durabilities": [ 0 ], @@ -159,7 +159,7 @@ fn everything() { } }, "2": { - "1025": { + "4294968321": { "durabilities": [ 0 ], @@ -172,7 +172,7 @@ fn everything() { } }, "4": { - "3073": { + "4294970369": { "durability": 2, "last_interned_at": 1, "fields": [ @@ -181,7 +181,7 @@ fn everything() { } }, "5": { - "4097": { + "4294971393": { "durability": 0, "updated_at": 1, "revisions": [], @@ -191,24 +191,24 @@ fn everything() { } }, "7": { - "5121": { + "4294972417": { "durability": 2, "last_interned_at": 18446744073709551615, "fields": [ - 3, - 4 + 4294967299, + 4294967300 ] } }, "19": { - "2049": { + "4294969345": { "durability": 2, "last_interned_at": 18446744073709551615, "fields": null } }, "6": { - "7:5121": { + "7:4294972417": { "value": "aaa", "verified_at": 1, "revisions": { @@ -218,11 +218,11 @@ fn everything() { "Derived": [ [ 3, - 1 + 4 ], [ 4, - 1 + 4 ] ] }, @@ -232,8 +232,8 @@ fn everything() { } }, "8": { - "0:3": { - "value": 4097, + "0:4294967299": { + "value": 4294971393, "verified_at": 1, "revisions": { "changed_at": 1, @@ -242,7 +242,7 @@ fn everything() { "Derived": [ [ 3, - 1 + 4 ] ] }, @@ -255,7 +255,7 @@ fn everything() { "hash": 6073466998405137972, "disambiguator": 0 }, - 4097 + 4294971393 ] ], "cycle_heads": [], @@ -265,8 +265,8 @@ fn everything() { } }, "18": { - "19:2049": { - "value": 3073, + "19:4294969345": { + "value": 4294970369, "verified_at": 1, "revisions": { "changed_at": 1, @@ -275,7 +275,7 @@ fn everything() { "Derived": [ [ 3073, - 4 + 1 ] ] }, @@ -347,7 +347,7 @@ fn partial_query() { }, "ingredients": { "0": { - "1": { + "4294967297": { "durabilities": [ 0 ], @@ -360,7 +360,7 @@ fn partial_query() { } }, "13": { - "0:1": { + "0:4294967297": { "value": 1, "verified_at": 1, "revisions": { @@ -370,7 +370,7 @@ fn partial_query() { "Derived": [ [ 1, - 1 + 4 ] ] }, @@ -462,7 +462,7 @@ fn partial_query_interned() { }, "ingredients": { "0": { - "1": { + "4294967297": { "durabilities": [ 0 ], @@ -475,7 +475,7 @@ fn partial_query_interned() { } }, "4": { - "3073": { + "4294970369": { "durability": 0, "last_interned_at": 1, "fields": [ @@ -484,18 +484,18 @@ fn partial_query_interned() { } }, "17": { - "1025": { + "4294968321": { "durability": 2, "last_interned_at": 18446744073709551615, "fields": [ - 1, + 4294967297, 0 ] } }, "16": { - "17:1025": { - "value": 3073, + "17:4294968321": { + "value": 4294970369, "verified_at": 1, "revisions": { "changed_at": 1, @@ -504,11 +504,11 @@ fn partial_query_interned() { "Derived": [ [ 1, - 1 + 4 ], [ 3073, - 4 + 1 ] ] }, From 635405d94080601f1ae8c0e7b4345afa85c77ab9 Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Sun, 28 Dec 2025 08:38:45 +0200 Subject: [PATCH 2/2] Use a generationless ID in the sync table Because we sometimes get id from a `DatabaseKeyIndex` that doesn't have the generation, so it prevent conflicts. And we don't need it either, as explained in a comment. --- src/function/sync.rs | 43 +++++++++++++++-------- src/id.rs | 81 +++++++++++++++++++++++++++++++++++++------- src/key.rs | 13 +++++++ 3 files changed, 110 insertions(+), 27 deletions(-) diff --git a/src/function/sync.rs b/src/function/sync.rs index b0d1596ec..a0419814f 100644 --- a/src/function/sync.rs +++ b/src/function/sync.rs @@ -1,6 +1,7 @@ use rustc_hash::FxHashMap; use std::collections::hash_map::OccupiedEntry; +use crate::id::GenerationlessId; use crate::key::DatabaseKeyIndex; use crate::runtime::{ BlockOnTransferredOwner, BlockResult, BlockTransferredResult, Running, WaitResult, @@ -11,12 +12,15 @@ use crate::tracing; use crate::zalsa::Zalsa; use crate::{Id, IngredientIndex}; -pub(crate) type SyncGuard<'me> = crate::sync::MutexGuard<'me, FxHashMap>; +pub(crate) type SyncGuard<'me> = + crate::sync::MutexGuard<'me, FxHashMap>; /// Tracks the keys that are currently being processed; used to coordinate between /// worker threads. pub(crate) struct SyncTable { - syncs: Mutex>, + // It's okay to use `GenerationlessId` here as `SyncTable` only holds queries while they're executing, and + // therefore if pointing to an interned struct it must be alive, therefore we don't need the generation. + syncs: Mutex>, ingredient: IngredientIndex, } @@ -72,12 +76,17 @@ impl SyncTable { reentrant: Reentrancy, ) -> ClaimResult<'me> { let mut write = self.syncs.lock(); - match write.entry(key_index) { + match write.entry(key_index.generationless()) { std::collections::hash_map::Entry::Occupied(occupied_entry) => { let id = match occupied_entry.get().id { SyncOwner::Thread(id) => id, SyncOwner::Transferred => { - return match self.try_claim_transferred(zalsa, occupied_entry, reentrant) { + return match self.try_claim_transferred( + zalsa, + occupied_entry, + reentrant, + key_index, + ) { Ok(claimed) => claimed, Err(other_thread) => match other_thread.block(write) { BlockResult::Cycle => ClaimResult::Cycle { inner: false }, @@ -130,7 +139,7 @@ impl SyncTable { reentrant: Reentrancy, ) -> ClaimResult<'me, ()> { let mut write = self.syncs.lock(); - match write.entry(key_index) { + match write.entry(key_index.generationless()) { std::collections::hash_map::Entry::Occupied(occupied_entry) => { let id = match occupied_entry.get().id { SyncOwner::Thread(id) => id, @@ -172,10 +181,10 @@ impl SyncTable { fn try_claim_transferred<'me>( &'me self, zalsa: &'me Zalsa, - mut entry: OccupiedEntry, + mut entry: OccupiedEntry, reentrant: Reentrancy, + key_index: Id, ) -> Result, Box>> { - let key_index = *entry.key(); let database_key_index = DatabaseKeyIndex::new_non_interned(self.ingredient, key_index); let thread_id = thread::current().id(); @@ -226,11 +235,12 @@ impl SyncTable { fn peek_claim_transferred<'me>( &'me self, zalsa: &'me Zalsa, - mut entry: OccupiedEntry, + mut entry: OccupiedEntry, reentrant: Reentrancy, ) -> Result, Box>> { let key_index = *entry.key(); - let database_key_index = DatabaseKeyIndex::new_non_interned(self.ingredient, key_index); + let database_key_index = + DatabaseKeyIndex::new_non_interned_generationless(self.ingredient, key_index); let thread_id = thread::current().id(); match zalsa @@ -257,7 +267,7 @@ impl SyncTable { /// is currently blocked on this thread (claiming `key_index` from this thread results in a cycle). pub(super) fn mark_as_transfer_target(&self, key_index: Id) -> Option { let mut syncs = self.syncs.lock(); - syncs.get_mut(&key_index).map(|state| { + syncs.get_mut(&key_index.generationless()).map(|state| { // We set `anyone_waiting` to true because it is used in `ClaimGuard::release` // to exit early if the query doesn't need to release any locks. // However, there are now dependent queries that need to be released, that's why we set `anyone_waiting` to true, @@ -318,7 +328,9 @@ impl<'me> ClaimGuard<'me> { #[inline(never)] fn release_panicking(&self) { let mut syncs = self.sync_table.syncs.lock(); - let state = syncs.remove(&self.key_index).expect("key claimed twice?"); + let state = syncs + .remove(&self.key_index.generationless()) + .expect("key claimed twice?"); tracing::debug!( "Release claim on {:?} due to panic", self.database_key_index() @@ -358,7 +370,8 @@ impl<'me> ClaimGuard<'me> { #[inline(never)] fn release_self(&self) { let mut syncs = self.sync_table.syncs.lock(); - let std::collections::hash_map::Entry::Occupied(mut state) = syncs.entry(self.key_index) + let std::collections::hash_map::Entry::Occupied(mut state) = + syncs.entry(self.key_index.generationless()) else { panic!("key should only be claimed/released once"); }; @@ -388,7 +401,7 @@ impl<'me> ClaimGuard<'me> { self.sync_table .syncs .lock() - .remove(&self.key_index) + .remove(&self.key_index.generationless()) .expect("key should only be claimed/released once"), WaitResult::Panicked, ); @@ -406,7 +419,7 @@ impl<'me> ClaimGuard<'me> { let SyncState { id, claimed_twice, .. } = syncs - .get_mut(&self.key_index) + .get_mut(&self.key_index.generationless()) .expect("key should only be claimed/released once"); *id = SyncOwner::Transferred; @@ -434,7 +447,7 @@ impl<'me> ClaimGuard<'me> { ReleaseMode::Default => { let mut syncs = self.sync_table.syncs.lock(); let state = syncs - .remove(&self.key_index) + .remove(&self.key_index.generationless()) .expect("key should only be claimed/released once"); self.release(state, WaitResult::Completed); diff --git a/src/id.rs b/src/id.rs index e38863760..2493b4d35 100644 --- a/src/id.rs +++ b/src/id.rs @@ -4,6 +4,51 @@ use std::num::NonZeroU32; use crate::zalsa::Zalsa; +/// An [`Id`] without a generation (for interned structs). +/// +/// You need to be very careful with this, and we deliberately do not expose it outside Salsa: +/// it is required in some places because of the packing shenanigans we do for `DatabaseKeyIndex`, +/// but using it in an incorrect place can mean you'll get the wrong value back. +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(crate) struct GenerationlessId { + index: NonZeroU32, +} + +impl GenerationlessId { + /// Create a `salsa::Id` from a u32 value, without a generation. This + /// value should be less than [`Self::MAX_U32`]. + /// + /// In general, you should not need to create salsa ids yourself, + /// but it can be useful if you are using the type as a general + /// purpose "identifier" internally. + /// + /// # Safety + /// + /// The supplied value must be less than [`Self::MAX_U32`]. + #[doc(hidden)] + #[track_caller] + #[inline] + pub const unsafe fn from_index(index: u32) -> Self { + debug_assert!(index < Id::MAX_U32); + + GenerationlessId { + // SAFETY: Caller obligation. + index: unsafe { NonZeroU32::new_unchecked(index + 1) }, + } + } + + #[inline] + pub(crate) const fn index_nonzero(self) -> NonZeroU32 { + self.index + } + + /// Return the index portion of this `Id`. + #[inline] + pub const fn index(self) -> u32 { + self.index.get() - 1 + } +} + /// The `Id` of a salsa struct in the database [`Table`](`crate::table::Table`). /// /// The high-order bits of an `Id` store a 32-bit generation counter, while @@ -19,7 +64,7 @@ use crate::zalsa::Zalsa; /// it is wrapped in new types. #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct Id { - index: NonZeroU32, + generationless: GenerationlessId, /// `generation` is shifted by one bit, to allow bit tagging for `DatabaseKeyIndex`. /// The LSB is always 1. /// @@ -46,11 +91,9 @@ impl Id { #[track_caller] #[inline] pub const unsafe fn from_index(index: u32) -> Self { - debug_assert!(index < Self::MAX_U32); - Id { // SAFETY: Caller obligation. - index: unsafe { NonZeroU32::new_unchecked(index + 1) }, + generationless: unsafe { GenerationlessId::from_index(index) }, generation: 0b1, } } @@ -72,7 +115,10 @@ impl Id { let generation = (bits >> 32) as u32; debug_assert!((generation & 0b1) == 0b1); - Id { index, generation } + Id { + generationless: GenerationlessId { index }, + generation, + } } /// Create a `salsa::Id` from a `u64` value. @@ -86,13 +132,16 @@ impl Id { let generation = (bits >> 32) as u32; debug_assert!((generation & 0b1) == 0b1); - Id { index, generation } + Id { + generationless: GenerationlessId { index }, + generation, + } } /// Return a `u64` representation of this `Id`. #[inline] pub fn as_bits(self) -> u64 { - u64::from(self.index.get()) | (u64::from(self.generation) << 32) + u64::from(self.generationless.index.get()) | (u64::from(self.generation) << 32) } /// Returns a new `Id` with same index, but the generation incremented by one. @@ -102,7 +151,7 @@ impl Id { #[inline] pub fn next_generation(self) -> Option { self.generation.checked_add(1).map(|generation| Id { - index: self.index, + generationless: self.generationless, generation: generation | 0b1, }) } @@ -115,27 +164,30 @@ impl Id { #[inline] pub const fn with_generation(self, generation: u32) -> Id { Id { - index: self.index, + generationless: self.generationless, generation: (generation << 1) | 0b1, } } #[inline] pub(crate) const fn index_nonzero(self) -> NonZeroU32 { - self.index + self.generationless.index_nonzero() } #[inline] pub(crate) fn from_raw_parts(index: NonZeroU32, generation: u32) -> Self { debug_assert!(index.get() <= Self::MAX_U32); debug_assert!((generation & 0b1) == 0b1); - Self { index, generation } + Self { + generationless: GenerationlessId { index }, + generation, + } } /// Return the index portion of this `Id`. #[inline] pub const fn index(self) -> u32 { - self.index.get() - 1 + self.generationless.index() } /// Return the generation of this `Id`. @@ -143,6 +195,11 @@ impl Id { pub const fn generation(self) -> u32 { self.generation } + + #[inline] + pub(crate) fn generationless(self) -> GenerationlessId { + self.generationless + } } impl Hash for Id { diff --git a/src/key.rs b/src/key.rs index 4f9cc1186..0875a9844 100644 --- a/src/key.rs +++ b/src/key.rs @@ -2,6 +2,7 @@ use std::fmt; use std::num::NonZeroU32; use crate::function::{VerifyCycleHeads, VerifyResult}; +use crate::id::GenerationlessId; use crate::zalsa::{IngredientIndex, Zalsa}; use crate::{Database, Id}; @@ -47,6 +48,18 @@ impl DatabaseKeyIndex { } } + #[inline] + pub(crate) const fn new_non_interned_generationless( + ingredient_index: IngredientIndex, + key_index: GenerationlessId, + ) -> Self { + Self { + index: key_index.index_nonzero(), + generation_or_ingredient_index: ingredient_index.as_u32() + << Self::INGREDIENT_INDEX_SHIFT, + } + } + #[inline] pub(crate) fn new_non_interned_with_tag( ingredient_index: IngredientIndex,