refactor: better tests

Author: ChieloNewctle

README.md

@@ -24,15 +24,15 @@ use futures_executor::block_on;
 
 fn main() {
     block_on(async {
-        let (wg, handle) = MonoWaitGroup::new();
+        let (wg, token) = MonoWaitGroup::new();
 
         thread::spawn(move || {
             println!("Worker started");
             // Long-running task...
             thread::sleep(Duration::from_secs(1));
             println!("Worker finished");
-            // Handle is dropped here, signaling completion
-            handle.done();
+            // Token is dropped here, signaling completion
+            token.release();
         });
 
         // Wait for the task to complete
@@ -54,16 +54,16 @@ use futures_executor::block_on;
 
 fn main() {
     block_on(async {
-        let (wg, base_handle) = WaitGroup::new();
+        let (wg, factory) = WaitGroup::new();
 
-        for (i, handle) in repeat_n(base_handle, 8).enumerate() {
+        for (i, token) in repeat_n(factory.into_token(), 8).enumerate() {
             thread::spawn(move || {
                 println!("Task {i} started");
                 // Long-running task...
                 thread::sleep(Duration::from_secs(1));
                 println!("Task {i} finished");
-                // Handle is dropped here, signaling completion
-                handle.done();
+                // Token is dropped here, signaling completion
+                token.release();
             });
         }
 
@@ -86,18 +86,20 @@ use tokio::time::sleep;
 
 #[tokio::main]
 async fn main() {
-    let (wg, base_handle) = WaitGroup::new();
+    let (wg, factory) = WaitGroup::new();
 
-    for (i, handle) in repeat_n(base_handle, 8).enumerate() {
-        let task = async move {
-            println!("Task {i} started");
-            // Long-running task...
-            sleep(Duration::from_secs(1)).await;
-            println!("Task {i} finished");
+    factory.scope(|token| {
+        for (i, token) in repeat_n(token, 8).enumerate() {
+            let task = async move {
+                println!("Task {i} started");
+                // Long-running task...
+                sleep(Duration::from_secs(1)).await;
+                println!("Task {i} finished");
+            }
+            .release_on_ready(token);
+            tokio::spawn(task);
         }
-        .with_worker_handle(handle);
-        tokio::spawn(task);
-    }
+    });
 
     // Wait for the task to complete
     wg.await;

src/ext.rs

@@ -6,13 +6,15 @@ use core::{
 use derive_more::Into;
 use pin_project_lite::pin_project;
 
-use crate::{GroupToken, MonoGroupToken};
+use crate::{GroupToken, MonoGroupToken, group::GroupTokenFactory};
 
 pub trait GroupTokenExt<T>: Sized {
+    #[inline]
     fn release_on_drop(self, token: T) -> GroupTokenReleaseOnDrop<Self, T> {
         GroupTokenReleaseOnDrop { inner: self, token }
     }
 
+    #[inline]
     fn release_on_ready(self, token: T) -> GroupTokenReleaseOnReady<Self, T> {
         GroupTokenReleaseOnReady {
             inner: self,
@@ -23,6 +25,7 @@ pub trait GroupTokenExt<T>: Sized {
 
 trait GroupTokenType {}
 
+impl GroupTokenType for GroupTokenFactory {}
 impl GroupTokenType for GroupToken {}
 impl GroupTokenType for MonoGroupToken {}
 
@@ -47,20 +50,24 @@ pin_project! {
 }
 
 impl<F, T> GroupTokenReleaseOnDrop<F, T> {
+    #[inline]
     pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
         self.project().inner
     }
 
+    #[inline]
     pub fn group_token(&self) -> &T {
         &self.token
     }
 }
 
 impl<F, T> GroupTokenReleaseOnReady<F, T> {
+    #[inline]
     pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
         self.project().inner
     }
 
+    #[inline]
     pub fn group_token(&self) -> Option<&T> {
         self.token.as_ref()
     }
@@ -69,6 +76,7 @@ impl<F, T> GroupTokenReleaseOnReady<F, T> {
 impl<F: Future, T> Future for GroupTokenReleaseOnDrop<F, T> {
     type Output = F::Output;
 
+    #[inline]
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         self.inner_pin().poll(cx)
     }
@@ -77,6 +85,7 @@ impl<F: Future, T> Future for GroupTokenReleaseOnDrop<F, T> {
 impl<F: Future, T> Future for GroupTokenReleaseOnReady<F, T> {
     type Output = F::Output;
 
+    #[inline]
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.project();
         let res = this.inner.poll(cx);

src/group.rs

@@ -3,7 +3,7 @@ use core::{
     task::{Context, Poll},
 };
 
-use derive_more::Debug;
+use derive_more::{Debug, Into};
 
 use crate::{
     layout::SharedLayout,
@@ -77,34 +77,42 @@ pub struct WaitGroup(#[debug("done: {}", _0.is_done())] WaitGroupWrapper<TwinRef
 pub struct MonoWaitGroup(#[debug("done: {}", _0.is_done())] WaitGroupWrapper<TwinRef<MonoLayout>>);
 
 /// Clonable group token.
-#[allow(unused)]
 #[must_use]
 #[derive(Clone, Debug)]
-pub struct GroupToken(#[debug("done: {}", _0.is_done())] ClonableTwinRef<SharedLayout>);
+pub struct GroupToken(
+    #[allow(unused)]
+    #[debug("done: {}", _0.is_done())]
+    ClonableTwinRef<SharedLayout>,
+);
 
 /// Non-clonable group token.
 #[must_use]
 #[derive(Debug)]
 pub struct MonoGroupToken(#[debug("done: {}", _0.is_done())] TwinRef<MonoLayout>);
 
+/// Factory of `GroupToken`.
+#[must_use]
+#[derive(Debug, Into)]
+pub struct GroupTokenFactory(GroupToken);
+
 impl WaitGroup {
-    /// Creates a new `WaitGroup` and a clonable `GroupToken`.
-    ///
-    /// The `WaitGroup` is used to await the completion of tasks. The
-    /// `GroupToken` is used to signal task completion.
+    /// Creates a new `WaitGroup` and a `GroupTokenFactory`.
     ///
     /// # Examples
     ///
     /// ```
     /// use compact_waitgroup::WaitGroup;
     ///
-    /// let (wg, token) = WaitGroup::new();
-    /// // ... distribute token ...
+    /// let (wg, factory) = WaitGroup::new();
+    /// // ... distribute token with factory ...
     /// ```
-    pub fn new() -> (Self, GroupToken) {
+    pub fn new() -> (Self, GroupTokenFactory) {
         let inner = SharedLayout::new();
         let (wg, token) = TwinRef::new_clonable(inner);
-        (Self(WaitGroupWrapper::new(wg)), GroupToken(token))
+        (
+            Self(WaitGroupWrapper::new(wg)),
+            GroupTokenFactory(GroupToken(token)),
+        )
     }
 
     /// Checks if the `WaitGroup` has completed.
@@ -186,6 +194,26 @@ impl Future for MonoWaitGroup {
     }
 }
 
+impl GroupTokenFactory {
+    /// Consumes the inner token.
+    ///
+    /// This is equivalent to dropping the factory.
+    #[inline]
+    pub fn release(self) {
+        drop(self);
+    }
+
+    #[inline]
+    pub fn into_token(self) -> GroupToken {
+        self.0
+    }
+
+    #[inline]
+    pub fn scope<T, F: FnOnce(GroupToken) -> T>(self, func: F) -> T {
+        func(self.into_token())
+    }
+}
+
 impl GroupToken {
     /// Consumes the token.
     ///
@@ -204,6 +232,16 @@ impl MonoGroupToken {
     pub fn release(self) {
         drop(self);
     }
+
+    #[inline]
+    pub fn into_token(self) -> Self {
+        self
+    }
+
+    #[inline]
+    pub fn scope<T, F: FnOnce(MonoGroupToken) -> T>(self, func: F) -> T {
+        func(self.into_token())
+    }
 }
 
 impl Drop for MonoGroupToken {

src/lib.rs

@@ -27,16 +27,18 @@
 //! ```rust
 //! # use compact_waitgroup::WaitGroup;
 //! # futures_executor::block_on(async {
-//! let (wg, token) = WaitGroup::new();
-//! let token_cloned = token.clone();
-//! assert!(!wg.is_done());
-//! std::thread::spawn(move || {
-//!     // Long-running task
-//!     token_cloned.release();
-//! });
-//! std::thread::spawn(move || {
-//!     // Another long-running task
-//!     token.release();
+//! let (wg, factory) = WaitGroup::new();
+//! factory.scope(|token| {
+//!     let token_cloned = token.clone();
+//!     assert!(!wg.is_done());
+//!     std::thread::spawn(move || {
+//!         // Long-running task
+//!         token_cloned.release();
+//!     });
+//!     std::thread::spawn(move || {
+//!         // Another long-running task
+//!         token.release();
+//!     });
 //! });
 //! // Wait for all tasks to complete
 //! wg.await;