Deduplicate Storage and StorageHandle fields

Author: Veykril

src/storage.rs

@@ -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 {
@@ -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(),
         }
     }
@@ -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 {
@@ -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,
         }
     }
 }
@@ -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>>
@@ -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
@@ -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
     }
@@ -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,
         }
     }
 }