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 11, 2025
1 parent 1887089 commit 4e44751
Showing 1 changed file with 23 additions and 43 deletions.
66 changes: 23 additions & 43 deletions src/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,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 @@ -39,15 +51,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 @@ -67,20 +72,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 @@ -97,13 +92,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 @@ -113,16 +103,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
}

// ANCHOR: cancel_other_workers
Expand All @@ -132,29 +116,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 @@ -170,13 +154,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 4e44751

Please sign in to comment.