Simplify WaitGroup implementation (#958)

* Simplify `WaitGroup` implementation

* Slightly cheaper `get_mut`

Co-authored-by: Ibraheem Ahmed <[email protected]>

---------

Co-authored-by: Ibraheem Ahmed <[email protected]>

Author: Veykril

src/storage.rs

@@ -25,8 +25,6 @@ pub struct StorageHandle<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)),
@@ -53,7 +51,7 @@ impl<Db: Database> StorageHandle<Db> {
         Self {
             zalsa_impl: Arc::new(Zalsa::new::<Db>(event_callback, jars)),
             coordinate: CoordinateDrop(Arc::new(Coordinate {
-                clones: Mutex::new(1),
+                coordinate_lock: Mutex::default(),
                 cvar: Default::default(),
             })),
             phantom: PhantomData,
@@ -95,17 +93,6 @@ impl<Db> Drop for Storage<Db> {
     }
 }
 
-struct Coordinate {
-    /// Counter of the number of clones of actor. Begins at 1.
-    /// Incremented when cloned, decremented when dropped.
-    clones: Mutex<usize>,
-    cvar: Condvar,
-}
-
-// We cannot panic while holding a lock to `clones: Mutex<usize>` and therefore we cannot enter an
-// inconsistent state.
-impl RefUnwindSafe for Coordinate {}
-
 impl<Db: Database> Default for Storage<Db> {
     fn default() -> Self {
         Self::new(None)
@@ -168,12 +155,15 @@ impl<Db: Database> Storage<Db> {
             .zalsa_impl
             .event(&|| Event::new(EventKind::DidSetCancellationFlag));
 
-        let mut clones = self.handle.coordinate.clones.lock();
-        while *clones != 1 {
-            clones = self.handle.coordinate.cvar.wait(clones);
-        }
-        // The ref count on the `Arc` should now be 1
-        let zalsa = Arc::get_mut(&mut self.handle.zalsa_impl).unwrap();
+        let mut coordinate_lock = self.handle.coordinate.coordinate_lock.lock();
+        let zalsa = loop {
+            if Arc::strong_count(&self.handle.zalsa_impl) == 1 {
+                // SAFETY: The strong count is 1, and we never create any weak pointers,
+                // so we have a unique reference.
+                break unsafe { &mut *(Arc::as_ptr(&self.handle.zalsa_impl).cast_mut()) };
+            }
+            coordinate_lock = self.handle.coordinate.cvar.wait(coordinate_lock);
+        };
         // cancellation is done, so reset the flag
         zalsa.runtime_mut().reset_cancellation_flag();
         zalsa
@@ -260,6 +250,16 @@ impl<Db: Database> Clone for Storage<Db> {
     }
 }
 
+/// A simplified `WaitGroup`, this is used together with `Arc<Zalsa>` as the actual counter
+struct Coordinate {
+    coordinate_lock: Mutex<()>,
+    cvar: Condvar,
+}
+
+// We cannot panic while holding a lock to `clones: Mutex<usize>` and therefore we cannot enter an
+// inconsistent state.
+impl RefUnwindSafe for Coordinate {}
+
 struct CoordinateDrop(Arc<Coordinate>);
 
 impl std::ops::Deref for CoordinateDrop {
@@ -272,7 +272,6 @@ impl std::ops::Deref for CoordinateDrop {
 
 impl Drop for CoordinateDrop {
     fn drop(&mut self) {
-        *self.0.clones.lock() -= 1;
         self.0.cvar.notify_all();
     }
 }

src/table.rs

@@ -252,7 +252,10 @@ impl Table {
                 }
 
                 let allocated_idx = self.push_page::<T>(ingredient, memo_types.clone());
-                assert_eq!(allocated_idx, page_idx);
+                assert_eq!(
+                    allocated_idx, page_idx,
+                    "allocated index does not match requested index"
+                );
             }
         };
     }

src/views.rs

@@ -108,7 +108,11 @@ impl Views {
         &self,
         func: fn(NonNull<Concrete>) -> NonNull<DbView>,
     ) -> &DatabaseDownCaster<DbView> {
-        assert_eq!(self.source_type_id, TypeId::of::<Concrete>());
+        assert_eq!(
+            self.source_type_id,
+            TypeId::of::<Concrete>(),
+            "mismatched source type"
+        );
         let target_type_id = TypeId::of::<DbView>();
         if let Some((_, caster)) = self
             .view_casters

src/zalsa_local.rs

@@ -1173,7 +1173,7 @@ impl ActiveQueryGuard<'_> {
         unsafe {
             self.local_state.with_query_stack_unchecked_mut(|stack| {
                 #[cfg(debug_assertions)]
-                assert_eq!(stack.len(), self.push_len);
+                assert_eq!(stack.len(), self.push_len, "mismatched push and pop");
                 let frame = stack.last_mut().unwrap();
                 frame.tracked_struct_ids_mut().seed(tracked_struct_ids);
             })
@@ -1195,7 +1195,7 @@ impl ActiveQueryGuard<'_> {
         unsafe {
             self.local_state.with_query_stack_unchecked_mut(|stack| {
                 #[cfg(debug_assertions)]
-                assert_eq!(stack.len(), self.push_len);
+                assert_eq!(stack.len(), self.push_len, "mismatched push and pop");
                 let frame = stack.last_mut().unwrap();
                 frame.seed_iteration(durability, changed_at, edges, untracked_read, tracked_ids);
             })