almost_enough/
stop_token.rs

1//! Arc-based cloneable dynamic dispatch for Stop.
2//!
3//! This module provides [`StopToken`], a shared-ownership wrapper that enables
4//! dynamic dispatch without monomorphization bloat, with cheap `Clone`.
5//!
6//! # StopToken vs BoxedStop
7//!
8//! | | `StopToken` | `BoxedStop` |
9//! |---|-----------|-------------|
10//! | Clone | Yes (Arc increment) | No |
11//! | Storage | `Arc<dyn Stop>` | `Box<dyn Stop>` |
12//! | Send to threads | Clone and move | Must wrap in Arc yourself |
13//! | Use case | Default choice | When Clone is unwanted |
14//!
15//! # Example
16//!
17//! ```rust
18//! use almost_enough::{StopToken, Stopper, Unstoppable, Stop};
19//!
20//! let stopper = Stopper::new();
21//! let stop = StopToken::new(stopper.clone());
22//! let stop2 = stop.clone(); // Arc increment, no allocation
23//!
24//! stopper.cancel();
25//! assert!(stop.should_stop());
26//! assert!(stop2.should_stop());
27//! ```
28
29use alloc::sync::Arc;
30use core::any::{Any, TypeId};
31
32use crate::{Stop, StopReason};
33
34/// A shared-ownership [`Stop`] implementation with cheap `Clone`.
35///
36/// Wraps any `Stop` in an `Arc` for shared ownership across threads.
37/// Cloning is an atomic increment — no heap allocation.
38///
39/// # Indirection Collapsing
40///
41/// `StopToken::new()` detects when you pass another `StopToken` and unwraps
42/// it instead of double-wrapping. No-op stops (`Unstoppable`) are stored
43/// as `None` — `check()` short-circuits without any vtable dispatch.
44///
45/// # Example
46///
47/// ```rust
48/// use almost_enough::{StopToken, Stopper, Stop, StopReason};
49///
50/// let stopper = Stopper::new();
51/// let stop = StopToken::new(stopper.clone());
52/// let stop2 = stop.clone(); // cheap Arc clone
53///
54/// stopper.cancel();
55/// assert!(stop.should_stop());
56/// assert!(stop2.should_stop()); // both see cancellation
57/// ```
58pub struct StopToken {
59    inner: StopTokenInner,
60}
61
62/// Dispatch enum — avoids vtable for the common Stopper/SyncStopper cases.
63enum StopTokenInner {
64    /// No-op (Unstoppable). check() → Ok(()), no dispatch.
65    None,
66    /// Direct atomic load with Relaxed ordering (Stopper).
67    Relaxed(Arc<crate::stopper::StopperInner>),
68    /// Direct atomic load with Acquire ordering (SyncStopper).
69    Acquire(Arc<crate::sync_stopper::SyncStopperInner>),
70    /// Everything else — vtable dispatch.
71    Dyn(Arc<dyn Stop + Send + Sync>),
72}
73
74impl StopToken {
75    /// Create a new `StopToken` from any [`Stop`] implementation.
76    ///
77    /// If `stop` is already a `StopToken`, it is unwrapped instead of
78    /// double-wrapping (no extra indirection).
79    #[inline]
80    pub fn new<T: Stop + 'static>(stop: T) -> Self {
81        // Fast path: no-op stops skip all wrapping
82        if !stop.may_stop() {
83            return Self {
84                inner: StopTokenInner::None,
85            };
86        }
87        // Collapse StopToken nesting
88        if TypeId::of::<T>() == TypeId::of::<StopToken>() {
89            let any_ref: &dyn Any = &stop;
90            let inner = any_ref.downcast_ref::<StopToken>().unwrap();
91            let result = inner.clone();
92            drop(stop);
93            return result;
94        }
95        // Stopper: direct atomic, no vtable dispatch
96        if TypeId::of::<T>() == TypeId::of::<crate::Stopper>() {
97            let any_ref: &dyn Any = &stop;
98            let stopper = any_ref.downcast_ref::<crate::Stopper>().unwrap();
99            let result = Self {
100                inner: StopTokenInner::Relaxed(stopper.inner.clone()),
101            };
102            drop(stop);
103            return result;
104        }
105        // SyncStopper: direct atomic with Acquire ordering
106        if TypeId::of::<T>() == TypeId::of::<crate::SyncStopper>() {
107            let any_ref: &dyn Any = &stop;
108            let stopper = any_ref.downcast_ref::<crate::SyncStopper>().unwrap();
109            let result = Self {
110                inner: StopTokenInner::Acquire(stopper.inner.clone()),
111            };
112            drop(stop);
113            return result;
114        }
115        Self {
116            inner: StopTokenInner::Dyn(Arc::new(stop)),
117        }
118    }
119
120    /// Create a `StopToken` from an existing `Arc<T>` without re-wrapping.
121    ///
122    /// ```rust
123    /// use almost_enough::{StopToken, Stopper, Stop};
124    /// # #[cfg(feature = "std")]
125    /// # fn main() {
126    /// use std::sync::Arc;
127    ///
128    /// let stopper = Arc::new(Stopper::new());
129    /// let stop = StopToken::from_arc(stopper);
130    /// assert!(!stop.should_stop());
131    /// # }
132    /// # #[cfg(not(feature = "std"))]
133    /// # fn main() {}
134    /// ```
135    #[inline]
136    pub fn from_arc<T: Stop + 'static>(arc: Arc<T>) -> Self {
137        if !arc.may_stop() {
138            return Self {
139                inner: StopTokenInner::None,
140            };
141        }
142        if TypeId::of::<T>() == TypeId::of::<StopToken>() {
143            let any_ref: &dyn Any = &*arc;
144            let inner = any_ref.downcast_ref::<StopToken>().unwrap();
145            return inner.clone();
146        }
147        Self {
148            inner: StopTokenInner::Dyn(arc as Arc<dyn Stop + Send + Sync>),
149        }
150    }
151}
152
153impl Clone for StopTokenInner {
154    #[inline]
155    fn clone(&self) -> Self {
156        match self {
157            Self::None => Self::None,
158            Self::Relaxed(arc) => Self::Relaxed(Arc::clone(arc)),
159            Self::Acquire(arc) => Self::Acquire(Arc::clone(arc)),
160            Self::Dyn(arc) => Self::Dyn(Arc::clone(arc)),
161        }
162    }
163}
164
165impl Clone for StopToken {
166    #[inline]
167    fn clone(&self) -> Self {
168        Self {
169            inner: self.inner.clone(),
170        }
171    }
172}
173
174impl Stop for StopToken {
175    #[inline(always)]
176    fn check(&self) -> Result<(), StopReason> {
177        match &self.inner {
178            StopTokenInner::None => Ok(()),
179            StopTokenInner::Relaxed(inner) => inner.check(),
180            StopTokenInner::Acquire(inner) => inner.check(),
181            StopTokenInner::Dyn(inner) => inner.check(),
182        }
183    }
184
185    #[inline(always)]
186    fn should_stop(&self) -> bool {
187        match &self.inner {
188            StopTokenInner::None => false,
189            StopTokenInner::Relaxed(inner) => inner.should_stop(),
190            StopTokenInner::Acquire(inner) => inner.should_stop(),
191            StopTokenInner::Dyn(inner) => inner.should_stop(),
192        }
193    }
194
195    #[inline(always)]
196    fn may_stop(&self) -> bool {
197        !matches!(self.inner, StopTokenInner::None)
198    }
199}
200
201/// Zero-cost conversion: reuses the Stopper's Arc. Direct atomic dispatch, no vtable.
202impl From<crate::Stopper> for StopToken {
203    #[inline]
204    fn from(stopper: crate::Stopper) -> Self {
205        Self {
206            inner: StopTokenInner::Relaxed(stopper.inner),
207        }
208    }
209}
210
211/// Zero-cost conversion: reuses the SyncStopper's Arc. Direct atomic dispatch.
212impl From<crate::SyncStopper> for StopToken {
213    #[inline]
214    fn from(stopper: crate::SyncStopper) -> Self {
215        Self {
216            inner: StopTokenInner::Acquire(stopper.inner),
217        }
218    }
219}
220
221impl core::fmt::Debug for StopToken {
222    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
223        f.debug_tuple("StopToken").finish()
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use crate::{FnStop, StopSource, Stopper, Unstoppable};
231
232    #[test]
233    fn from_unstoppable() {
234        let stop = StopToken::new(Unstoppable);
235        assert!(!stop.should_stop());
236        assert!(stop.check().is_ok());
237    }
238
239    #[test]
240    fn from_stopper() {
241        let stopper = Stopper::new();
242        let stop = StopToken::new(stopper.clone());
243
244        assert!(!stop.should_stop());
245
246        stopper.cancel();
247
248        assert!(stop.should_stop());
249        assert_eq!(stop.check(), Err(StopReason::Cancelled));
250    }
251
252    #[test]
253    fn clone_is_cheap() {
254        let stopper = Stopper::new();
255        let stop = StopToken::new(stopper.clone());
256        let stop2 = stop.clone();
257
258        stopper.cancel();
259
260        // Both clones see the cancellation (shared state)
261        assert!(stop.should_stop());
262        assert!(stop2.should_stop());
263    }
264
265    #[cfg(feature = "std")]
266    #[test]
267    fn clone_send_to_thread() {
268        let stopper = Stopper::new();
269        let stop = StopToken::new(stopper.clone());
270
271        let handle = std::thread::spawn({
272            let stop = stop.clone();
273            move || stop.should_stop()
274        });
275
276        stopper.cancel();
277        // Thread may or may not see cancellation depending on timing
278        let _ = handle.join().unwrap();
279    }
280
281    #[test]
282    fn is_send_sync() {
283        fn assert_send_sync<T: Send + Sync>() {}
284        assert_send_sync::<StopToken>();
285    }
286
287    #[test]
288    fn debug_format() {
289        let stop = StopToken::new(Unstoppable);
290        let debug = alloc::format!("{:?}", stop);
291        assert!(debug.contains("StopToken"));
292    }
293
294    #[test]
295    fn may_stop_delegates() {
296        assert!(!StopToken::new(Unstoppable).may_stop());
297        assert!(StopToken::new(Stopper::new()).may_stop());
298    }
299
300    #[test]
301    fn unstoppable_is_none_internally() {
302        let stop = StopToken::new(Unstoppable);
303        assert!(!stop.may_stop());
304        assert!(stop.check().is_ok());
305    }
306
307    #[test]
308    fn collapses_nested_dyn_stop() {
309        let inner = StopToken::new(Unstoppable);
310        let outer = StopToken::new(inner);
311        assert!(!outer.may_stop());
312    }
313
314    #[test]
315    fn collapses_nested_stopper() {
316        let stopper = Stopper::new();
317        let inner = StopToken::new(stopper.clone());
318        let outer = StopToken::new(inner.clone());
319
320        // Both share the same Arc chain
321        stopper.cancel();
322        assert!(outer.should_stop());
323        assert!(inner.should_stop());
324    }
325
326    #[test]
327    fn from_arc() {
328        let stopper = Arc::new(Stopper::new());
329        let cancel_handle = stopper.clone();
330        let stop = StopToken::from_arc(stopper);
331
332        assert!(!stop.should_stop());
333        cancel_handle.cancel();
334        assert!(stop.should_stop());
335    }
336
337    #[test]
338    fn from_non_clone_fn_stop() {
339        // FnStop with non-Clone closure — StopToken doesn't need Clone on T
340        let flag = Arc::new(core::sync::atomic::AtomicBool::new(false));
341        let flag2 = flag.clone();
342        let stop = StopToken::new(FnStop::new(move || {
343            flag2.load(core::sync::atomic::Ordering::Relaxed)
344        }));
345
346        assert!(!stop.should_stop());
347
348        // Clone the StopToken (shares the Arc, not the closure)
349        let stop2 = stop.clone();
350        flag.store(true, core::sync::atomic::Ordering::Relaxed);
351
352        assert!(stop.should_stop());
353        assert!(stop2.should_stop());
354    }
355
356    #[test]
357    fn avoids_monomorphization() {
358        fn process(stop: StopToken) -> bool {
359            stop.should_stop()
360        }
361
362        assert!(!process(StopToken::new(Unstoppable)));
363        assert!(!process(StopToken::new(StopSource::new())));
364        assert!(!process(StopToken::new(Stopper::new())));
365    }
366
367    #[test]
368    fn hot_loop_pattern() {
369        let stop = StopToken::new(Unstoppable);
370        for _ in 0..1000 {
371            assert!(stop.check().is_ok()); // None path, no dispatch
372        }
373    }
374
375    #[test]
376    fn from_stopper_zero_cost() {
377        let stopper = Stopper::new();
378        let cancel = stopper.clone();
379        let stop: StopToken = stopper.into(); // zero-cost: reuses Arc
380
381        assert!(!stop.should_stop());
382        cancel.cancel();
383        assert!(stop.should_stop()); // same Arc, same AtomicBool
384    }
385
386    #[test]
387    fn from_sync_stopper_zero_cost() {
388        let stopper = crate::SyncStopper::new();
389        let cancel = stopper.clone();
390        let stop: StopToken = stopper.into();
391
392        assert!(!stop.should_stop());
393        cancel.cancel();
394        assert!(stop.should_stop());
395    }
396
397    #[test]
398    fn new_stopper_flattens() {
399        // StopToken::new(Stopper) should reuse the Stopper's Arc,
400        // not double-wrap in Arc<Stopper{Arc<AtomicBool>}>
401        let stopper = Stopper::new();
402        let cancel = stopper.clone();
403        let stop = StopToken::new(stopper);
404
405        cancel.cancel();
406        assert!(stop.should_stop());
407    }
408
409    #[test]
410    fn new_sync_stopper_flattens() {
411        let stopper = crate::SyncStopper::new();
412        let cancel = stopper.clone();
413        let stop = StopToken::new(stopper);
414
415        cancel.cancel();
416        assert!(stop.should_stop());
417    }
418
419    #[test]
420    fn from_arc_collapses_dynstop() {
421        // from_arc(Arc<StopToken>) should reuse inner, not double-wrap
422        let inner = StopToken::new(Stopper::new());
423        let arc = alloc::sync::Arc::new(inner);
424        let stop = StopToken::from_arc(arc);
425        assert!(!stop.should_stop());
426    }
427
428    #[test]
429    fn stopper_inner_debug() {
430        let stop = Stopper::new();
431        let debug = alloc::format!("{:?}", stop);
432        assert!(debug.contains("cancelled"));
433    }
434
435    #[test]
436    fn sync_stopper_inner_debug() {
437        let stop = crate::SyncStopper::new();
438        let debug = alloc::format!("{:?}", stop);
439        assert!(debug.contains("cancelled"));
440    }
441
442    #[test]
443    fn from_stopper_clone_shares_state() {
444        let stopper = Stopper::new();
445        let stop: StopToken = stopper.clone().into();
446        let stop2 = stop.clone(); // Arc clone of the flattened inner
447
448        stopper.cancel();
449        // All three (stopper, stop, stop2) share the same AtomicBool
450        assert!(stop.should_stop());
451        assert!(stop2.should_stop());
452    }
453}
almost_enough/stop_token.rs

almost_enough/
stop_token.rs
