Skip to content

Commit

Permalink
Deduplicate Storage and StorageHandle fields
Browse files Browse the repository at this point in the history
  • Loading branch information
Veykril committed Feb 6, 2025
1 parent db30255 commit 25154f7
Showing 1 changed file with 24 additions and 44 deletions.
68 changes: 24 additions & 44 deletions src/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,18 @@ pub struct StorageHandle<Db> {
phantom: PhantomData<fn() -> Db>,
}

impl<Db> Clone for StorageHandle<Db> {
fn clone(&self) -> Self {
*self.coordinate.clones.lock() += 1;

Self {
zalsa_impl: self.zalsa_impl.clone(),
coordinate: CoordinateDrop(Arc::clone(&self.coordinate)),
phantom: PhantomData,
}
}
}

impl<Db: Database> Default for StorageHandle<Db> {
fn default() -> Self {
Self {
Expand All @@ -40,15 +52,8 @@ impl<Db: Database> Default for StorageHandle<Db> {

impl<Db> StorageHandle<Db> {
pub fn into_storage(self) -> Storage<Db> {
let StorageHandle {
zalsa_impl,
coordinate,
phantom,
} = self;
Storage {
zalsa_impl,
coordinate,
phantom,
handle: self,
zalsa_local: ZalsaLocal::new(),
}
}
Expand All @@ -68,20 +73,10 @@ pub unsafe trait HasStorage: Database + Clone + Sized {

/// Concrete implementation of the [`Database`] trait with local state that can be used to drive computations.
pub struct Storage<Db> {
// Note: Drop order is important, zalsa_impl needs to drop before coordinate
/// Reference to the database.
zalsa_impl: Arc<Zalsa>,

// Note: Drop order is important, coordinate needs to drop after zalsa_impl
/// Coordination data for cancellation of other handles when `zalsa_mut` is called.
/// This could be stored in Zalsa but it makes things marginally cleaner to keep it separate.
coordinate: CoordinateDrop,
handle: StorageHandle<Db>,

/// Per-thread state
zalsa_local: zalsa_local::ZalsaLocal,

/// We store references to `Db`
phantom: PhantomData<fn() -> Db>,
}

struct Coordinate {
Expand All @@ -98,13 +93,8 @@ impl RefUnwindSafe for Coordinate {}
impl<Db: Database> Default for Storage<Db> {
fn default() -> Self {
Self {
zalsa_impl: Arc::new(Zalsa::new::<Db>()),
coordinate: CoordinateDrop(Arc::new(Coordinate {
clones: Mutex::new(1),
cvar: Default::default(),
})),
handle: StorageHandle::default(),
zalsa_local: ZalsaLocal::new(),
phantom: PhantomData,
}
}
}
Expand All @@ -114,16 +104,10 @@ impl<Db: Database> Storage<Db> {
/// and [`std::panic::UnwindSafe`].
pub fn into_zalsa_handle(self) -> StorageHandle<Db> {
let Storage {
zalsa_impl,
coordinate,
phantom,
handle,
zalsa_local: _,
} = self;
StorageHandle {
zalsa_impl,
coordinate,
phantom,
}
handle
}

pub fn debug_input_entries<T>(&self) -> impl Iterator<Item = &input::Value<T>>
Expand Down Expand Up @@ -171,7 +155,7 @@ impl<Db: Database> Storage<Db> {
/// possible as `zalsa_impl` only becomes
/// `None` once we are in the `Drop` impl.
fn zalsa_impl(&self) -> &Arc<Zalsa> {
&self.zalsa_impl
&self.handle.zalsa_impl
}

// ANCHOR: cancel_other_workers
Expand All @@ -181,29 +165,29 @@ impl<Db: Database> Storage<Db> {
/// This could deadlock if there is a single worker with two handles to the
/// same database!
fn cancel_others(&self, db: &Db) {
self.zalsa_impl.set_cancellation_flag();
self.handle.zalsa_impl.set_cancellation_flag();

db.salsa_event(&|| Event::new(EventKind::DidSetCancellationFlag));

let mut clones = self.coordinate.clones.lock();
let mut clones = self.handle.coordinate.clones.lock();
while *clones != 1 {
self.coordinate.cvar.wait(&mut clones);
self.handle.coordinate.cvar.wait(&mut clones);
}
}
// ANCHOR_END: cancel_other_workers
}

unsafe impl<T: HasStorage> ZalsaDatabase for T {
fn zalsa(&self) -> &Zalsa {
&self.storage().zalsa_impl
&self.storage().handle.zalsa_impl
}

fn zalsa_mut(&mut self) -> &mut Zalsa {
self.storage().cancel_others(self);

let storage = self.storage_mut();
// The ref count on the `Arc` should now be 1
let zalsa_mut = Arc::get_mut(&mut storage.zalsa_impl).unwrap();
let zalsa_mut = Arc::get_mut(&mut storage.handle.zalsa_impl).unwrap();
zalsa_mut.new_revision();
zalsa_mut
}
Expand All @@ -219,13 +203,9 @@ unsafe impl<T: HasStorage> ZalsaDatabase for T {

impl<Db: Database> Clone for Storage<Db> {
fn clone(&self) -> Self {
*self.coordinate.clones.lock() += 1;

Self {
zalsa_impl: self.zalsa_impl.clone(),
coordinate: CoordinateDrop(Arc::clone(&self.coordinate)),
handle: self.handle.clone(),
zalsa_local: ZalsaLocal::new(),
phantom: PhantomData,
}
}
}
Expand Down

0 comments on commit 25154f7

Please sign in to comment.