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//! | Dispatch | Direct for `Stopper`, else vtable | Same (wraps a `StopToken`) |
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` (or a
42/// [`BoxedStop`](crate::BoxedStop), which wraps one) and reuses it instead of
43/// double-wrapping. No-op stops (`Unstoppable`) are stored as `None` —
44/// `check()` short-circuits without any vtable dispatch. A `Stopper` is
45/// checked as a direct atomic load, without a vtable. A `SyncStopper` keeps
46/// its own `Arc` (no allocation) and is checked through the vtable.
47///
48/// # Example
49///
50/// ```rust
51/// use almost_enough::{StopToken, Stopper, Stop, StopReason};
52///
53/// let stopper = Stopper::new();
54/// let stop = StopToken::new(stopper.clone());
55/// let stop2 = stop.clone(); // cheap Arc clone
56///
57/// stopper.cancel();
58/// assert!(stop.should_stop());
59/// assert!(stop2.should_stop()); // both see cancellation
60/// ```
61pub struct StopToken {
62    inner: StopTokenInner,
63}
64
65/// Dispatch enum — avoids a vtable for the common `Stopper` case.
66///
67/// Keep it at three arms: with four, LLVM compiles the `match` in `check` to
68/// a jump table, an indirect jump on every check on x86-64.
69enum StopTokenInner {
70    /// No-op (Unstoppable). check() → Ok(()), no dispatch.
71    None,
72    /// Direct atomic load with Relaxed ordering (Stopper).
73    Relaxed(Arc<crate::stopper::StopperInner>),
74    /// Everything else — vtable dispatch.
75    Dyn(Arc<dyn Stop + Send + Sync>),
76}
77
78impl StopToken {
79    /// Create a new `StopToken` from any [`Stop`] implementation.
80    ///
81    /// If `stop` is already a `StopToken`, it is unwrapped instead of
82    /// double-wrapping (no extra indirection).
83    #[inline]
84    pub fn new<T: Stop + 'static>(stop: T) -> Self {
85        // Fast path: no-op stops skip all wrapping
86        if !stop.may_stop() {
87            return Self {
88                inner: StopTokenInner::None,
89            };
90        }
91        // Collapse StopToken nesting
92        if TypeId::of::<T>() == TypeId::of::<StopToken>() {
93            let any_ref: &dyn Any = &stop;
94            let inner = any_ref.downcast_ref::<StopToken>().unwrap();
95            let result = inner.clone();
96            drop(stop);
97            return result;
98        }
99        // A BoxedStop is a StopToken that can't be cloned: reuse it
100        if let Some(token) = boxed_token(&stop) {
101            return token;
102        }
103        // Stopper: direct atomic, no vtable dispatch
104        if TypeId::of::<T>() == TypeId::of::<crate::Stopper>() {
105            let any_ref: &dyn Any = &stop;
106            let stopper = any_ref.downcast_ref::<crate::Stopper>().unwrap();
107            let result = Self {
108                inner: StopTokenInner::Relaxed(stopper.inner.clone()),
109            };
110            drop(stop);
111            return result;
112        }
113        // SyncStopper: reuse its Arc behind the vtable, not Arc<SyncStopper>
114        if TypeId::of::<T>() == TypeId::of::<crate::SyncStopper>() {
115            let any_ref: &dyn Any = &stop;
116            let stopper = any_ref.downcast_ref::<crate::SyncStopper>().unwrap();
117            let result = Self {
118                inner: StopTokenInner::Dyn(stopper.inner.clone()),
119            };
120            drop(stop);
121            return result;
122        }
123        Self {
124            inner: StopTokenInner::Dyn(Arc::new(stop)),
125        }
126    }
127
128    /// Create a `StopToken` from an existing `Arc<T>` without re-wrapping.
129    ///
130    /// ```rust
131    /// use almost_enough::{StopToken, Stopper, Stop};
132    /// # #[cfg(feature = "std")]
133    /// # fn main() {
134    /// use std::sync::Arc;
135    ///
136    /// let stopper = Arc::new(Stopper::new());
137    /// let stop = StopToken::from_arc(stopper);
138    /// assert!(!stop.should_stop());
139    /// # }
140    /// # #[cfg(not(feature = "std"))]
141    /// # fn main() {}
142    /// ```
143    #[inline]
144    pub fn from_arc<T: Stop + 'static>(arc: Arc<T>) -> Self {
145        if !arc.may_stop() {
146            return Self {
147                inner: StopTokenInner::None,
148            };
149        }
150        if TypeId::of::<T>() == TypeId::of::<StopToken>() {
151            let any_ref: &dyn Any = &*arc;
152            let inner = any_ref.downcast_ref::<StopToken>().unwrap();
153            return inner.clone();
154        }
155        // A Stopper's flag is checked directly, as in `new`. A SyncStopper
156        // stays behind the vtable, where its checks are Acquire loads.
157        if TypeId::of::<T>() == TypeId::of::<crate::Stopper>() {
158            let any_ref: &dyn Any = &*arc;
159            let stopper = any_ref.downcast_ref::<crate::Stopper>().unwrap();
160            return Self {
161                inner: StopTokenInner::Relaxed(stopper.inner.clone()),
162            };
163        }
164        Self {
165            inner: StopTokenInner::Dyn(arc as Arc<dyn Stop + Send + Sync>),
166        }
167    }
168}
169
170/// The token inside `stop` if it is a (deprecated) `BoxedStop`.
171#[allow(deprecated)]
172#[inline]
173fn boxed_token<T: Stop + 'static>(stop: &T) -> Option<StopToken> {
174    let any_ref: &dyn Any = stop;
175    any_ref
176        .downcast_ref::<crate::BoxedStop>()
177        .map(|boxed| boxed.0.clone())
178}
179
180impl Clone for StopTokenInner {
181    #[inline]
182    fn clone(&self) -> Self {
183        match self {
184            Self::None => Self::None,
185            Self::Relaxed(arc) => Self::Relaxed(Arc::clone(arc)),
186            Self::Dyn(arc) => Self::Dyn(Arc::clone(arc)),
187        }
188    }
189}
190
191impl Clone for StopToken {
192    #[inline]
193    fn clone(&self) -> Self {
194        Self {
195            inner: self.inner.clone(),
196        }
197    }
198}
199
200impl Stop for StopToken {
201    #[inline(always)]
202    #[track_caller]
203    fn check(&self) -> Result<(), StopReason> {
204        match &self.inner {
205            StopTokenInner::None => Ok(()),
206            StopTokenInner::Relaxed(inner) => inner.check(),
207            StopTokenInner::Dyn(inner) => inner.check(),
208        }
209    }
210
211    #[inline(always)]
212    #[track_caller]
213    fn should_stop(&self) -> bool {
214        match &self.inner {
215            StopTokenInner::None => false,
216            StopTokenInner::Relaxed(inner) => inner.should_stop(),
217            StopTokenInner::Dyn(inner) => inner.should_stop(),
218        }
219    }
220
221    #[inline(always)]
222    fn may_stop(&self) -> bool {
223        !matches!(self.inner, StopTokenInner::None)
224    }
225}
226
227/// Zero-cost conversion: reuses the Stopper's Arc. Direct atomic dispatch, no vtable.
228impl From<crate::Stopper> for StopToken {
229    #[inline]
230    fn from(stopper: crate::Stopper) -> Self {
231        Self {
232            inner: StopTokenInner::Relaxed(stopper.inner),
233        }
234    }
235}
236
237/// Reuses the SyncStopper's Arc (no allocation); checks go through the vtable.
238impl From<crate::SyncStopper> for StopToken {
239    #[inline]
240    fn from(stopper: crate::SyncStopper) -> Self {
241        Self {
242            inner: StopTokenInner::Dyn(stopper.inner),
243        }
244    }
245}
246
247impl core::fmt::Debug for StopToken {
248    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
249        f.debug_tuple("StopToken").finish()
250    }
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256    use crate::{FnStop, StopSource, Stopper, Unstoppable};
257
258    #[test]
259    fn boxed_and_arced_dyn_stops_are_stops_with_almost_enough_alone() {
260        // `alloc` turns on enough's `alloc`, which implements `Stop` for these.
261        fn takes(stop: impl Stop + 'static) -> StopToken {
262            StopToken::new(stop)
263        }
264        let arc: Arc<dyn Stop> = Arc::new(Stopper::new());
265        assert!(takes(arc).check().is_ok());
266        let boxed: alloc::boxed::Box<dyn Stop> = alloc::boxed::Box::new(Stopper::cancelled());
267        assert!(takes(boxed).should_stop());
268    }
269
270    #[test]
271    fn from_arc_checks_a_stopper_directly_and_a_sync_stopper_through_the_vtable() {
272        let stopper = Stopper::new();
273        let token = StopToken::from_arc(Arc::new(stopper.clone()));
274        assert_eq!(direct(&token), "relaxed");
275        assert!(!token.should_stop());
276        stopper.cancel();
277        assert!(token.should_stop());
278
279        let sync = crate::SyncStopper::new();
280        let token = StopToken::from_arc(Arc::new(sync.clone()));
281        assert_eq!(direct(&token), "dyn");
282        sync.cancel();
283        assert!(token.should_stop());
284    }
285
286    fn direct(stop: &StopToken) -> &'static str {
287        match &stop.inner {
288            StopTokenInner::None => "none",
289            StopTokenInner::Relaxed(_) => "relaxed",
290            StopTokenInner::Dyn(_) => "dyn",
291        }
292    }
293
294    #[test]
295    #[allow(deprecated)] // Exercises the deprecated BoxedStop.
296    fn boxed_stop_takes_the_same_fast_paths() {
297        use crate::{BoxedStop, SyncStopper};
298        assert_eq!(direct(&BoxedStop::new(Unstoppable).0), "none");
299        let stopper = Stopper::new();
300        let boxed = BoxedStop::new(stopper.clone());
301        assert_eq!(direct(&boxed.0), "relaxed");
302        let StopTokenInner::Relaxed(inner) = &boxed.0.inner else {
303            unreachable!()
304        };
305        assert!(
306            Arc::ptr_eq(inner, &stopper.inner),
307            "reuses the Stopper's Arc"
308        );
309        let sync = SyncStopper::new();
310        let boxed = BoxedStop::new(sync.clone());
311        assert_eq!(direct(&boxed.0), "dyn");
312        let StopTokenInner::Dyn(inner) = &boxed.0.inner else {
313            unreachable!()
314        };
315        assert_eq!(
316            Arc::as_ptr(inner).cast::<()>(),
317            Arc::as_ptr(&sync.inner).cast::<()>(),
318            "reuses the SyncStopper's Arc"
319        );
320        assert_eq!(direct(&BoxedStop::new(StopSource::new()).0), "dyn");
321    }
322
323    #[test]
324    #[allow(deprecated)] // Exercises the deprecated BoxedStop.
325    fn boxed_stop_and_stop_token_nest_without_wrapping() {
326        use crate::BoxedStop;
327        let token = StopToken::new(FnStop::new(|| false));
328        let StopTokenInner::Dyn(original) = &token.inner else {
329            unreachable!()
330        };
331        let boxed = BoxedStop::new(token.clone());
332        let again = BoxedStop::new(boxed);
333        let back = StopToken::new(again);
334        let StopTokenInner::Dyn(inner) = &back.inner else {
335            unreachable!()
336        };
337        assert!(Arc::ptr_eq(inner, original));
338    }
339
340    #[test]
341    #[allow(deprecated)] // Exercises the deprecated BoxedStop.
342    fn boxed_stop_sees_cancellation_through_every_path() {
343        use crate::{BoxedStop, SyncStopper};
344        let stopper = Stopper::new();
345        let sync = SyncStopper::new();
346        let boxed = [
347            BoxedStop::new(stopper.clone()),
348            BoxedStop::new(sync.clone()),
349            BoxedStop::new(BoxedStop::new(StopToken::new(stopper.clone()))),
350        ];
351        assert!(boxed.iter().all(|stop| stop.check().is_ok()));
352        stopper.cancel();
353        sync.cancel();
354        assert!(boxed.iter().all(|stop| stop.should_stop()));
355        assert_eq!(boxed[0].check(), Err(StopReason::Cancelled));
356    }
357
358    #[test]
359    fn from_unstoppable() {
360        let stop = StopToken::new(Unstoppable);
361        assert!(!stop.should_stop());
362        assert!(stop.check().is_ok());
363    }
364
365    #[test]
366    fn from_stopper() {
367        let stopper = Stopper::new();
368        let stop = StopToken::new(stopper.clone());
369
370        assert!(!stop.should_stop());
371
372        stopper.cancel();
373
374        assert!(stop.should_stop());
375        assert_eq!(stop.check(), Err(StopReason::Cancelled));
376    }
377
378    #[test]
379    fn clone_is_cheap() {
380        let stopper = Stopper::new();
381        let stop = StopToken::new(stopper.clone());
382        let stop2 = stop.clone();
383
384        stopper.cancel();
385
386        // Both clones see the cancellation (shared state)
387        assert!(stop.should_stop());
388        assert!(stop2.should_stop());
389    }
390
391    #[cfg(feature = "std")]
392    #[test]
393    fn clone_send_to_thread() {
394        let stopper = Stopper::new();
395        let stop = StopToken::new(stopper.clone());
396
397        let handle = std::thread::spawn({
398            let stop = stop.clone();
399            move || stop.should_stop()
400        });
401
402        stopper.cancel();
403        // Thread may or may not see cancellation depending on timing
404        let _ = handle.join().unwrap();
405    }
406
407    #[test]
408    fn is_send_sync() {
409        fn assert_send_sync<T: Send + Sync>() {}
410        assert_send_sync::<StopToken>();
411    }
412
413    #[test]
414    fn debug_format() {
415        let stop = StopToken::new(Unstoppable);
416        let debug = alloc::format!("{:?}", stop);
417        assert!(debug.contains("StopToken"));
418    }
419
420    #[test]
421    fn may_stop_delegates() {
422        assert!(!StopToken::new(Unstoppable).may_stop());
423        assert!(StopToken::new(Stopper::new()).may_stop());
424    }
425
426    #[test]
427    fn unstoppable_is_none_internally() {
428        let stop = StopToken::new(Unstoppable);
429        assert!(!stop.may_stop());
430        assert!(stop.check().is_ok());
431    }
432
433    #[test]
434    fn collapses_nested_dyn_stop() {
435        let inner = StopToken::new(Unstoppable);
436        let outer = StopToken::new(inner);
437        assert!(!outer.may_stop());
438    }
439
440    #[test]
441    fn collapses_nested_stopper() {
442        let stopper = Stopper::new();
443        let inner = StopToken::new(stopper.clone());
444        let outer = StopToken::new(inner.clone());
445
446        // Both share the same Arc chain
447        stopper.cancel();
448        assert!(outer.should_stop());
449        assert!(inner.should_stop());
450    }
451
452    #[test]
453    fn from_arc() {
454        let stopper = Arc::new(Stopper::new());
455        let cancel_handle = stopper.clone();
456        let stop = StopToken::from_arc(stopper);
457
458        assert!(!stop.should_stop());
459        cancel_handle.cancel();
460        assert!(stop.should_stop());
461    }
462
463    #[test]
464    fn from_non_clone_fn_stop() {
465        // FnStop with non-Clone closure — StopToken doesn't need Clone on T
466        let flag = Arc::new(core::sync::atomic::AtomicBool::new(false));
467        let flag2 = flag.clone();
468        let stop = StopToken::new(FnStop::new(move || {
469            flag2.load(core::sync::atomic::Ordering::Relaxed)
470        }));
471
472        assert!(!stop.should_stop());
473
474        // Clone the StopToken (shares the Arc, not the closure)
475        let stop2 = stop.clone();
476        flag.store(true, core::sync::atomic::Ordering::Relaxed);
477
478        assert!(stop.should_stop());
479        assert!(stop2.should_stop());
480    }
481
482    #[test]
483    fn avoids_monomorphization() {
484        fn process(stop: StopToken) -> bool {
485            stop.should_stop()
486        }
487
488        assert!(!process(StopToken::new(Unstoppable)));
489        assert!(!process(StopToken::new(StopSource::new())));
490        assert!(!process(StopToken::new(Stopper::new())));
491    }
492
493    #[test]
494    fn hot_loop_pattern() {
495        let stop = StopToken::new(Unstoppable);
496        for _ in 0..1000 {
497            assert!(stop.check().is_ok()); // None path, no dispatch
498        }
499    }
500
501    #[test]
502    fn from_stopper_zero_cost() {
503        let stopper = Stopper::new();
504        let cancel = stopper.clone();
505        let stop: StopToken = stopper.into(); // zero-cost: reuses Arc
506
507        assert!(!stop.should_stop());
508        cancel.cancel();
509        assert!(stop.should_stop()); // same Arc, same AtomicBool
510    }
511
512    #[test]
513    fn from_sync_stopper_reuses_its_arc_behind_the_vtable() {
514        let stopper = crate::SyncStopper::new();
515        let cancel = stopper.clone();
516        let stop: StopToken = stopper.into();
517        assert_eq!(direct(&stop), "dyn");
518        let StopTokenInner::Dyn(inner) = &stop.inner else {
519            unreachable!()
520        };
521        assert_eq!(
522            Arc::as_ptr(inner).cast::<()>(),
523            Arc::as_ptr(&cancel.inner).cast::<()>()
524        );
525
526        assert!(!stop.should_stop());
527        cancel.cancel();
528        assert!(stop.should_stop());
529        assert_eq!(stop.check(), Err(StopReason::Cancelled));
530    }
531
532    #[test]
533    fn new_stopper_flattens() {
534        // StopToken::new(Stopper) should reuse the Stopper's Arc,
535        // not double-wrap in Arc<Stopper{Arc<AtomicBool>}>
536        let stopper = Stopper::new();
537        let cancel = stopper.clone();
538        let stop = StopToken::new(stopper);
539
540        cancel.cancel();
541        assert!(stop.should_stop());
542    }
543
544    #[test]
545    fn new_sync_stopper_flattens() {
546        let stopper = crate::SyncStopper::new();
547        let cancel = stopper.clone();
548        let stop = StopToken::new(stopper);
549
550        cancel.cancel();
551        assert!(stop.should_stop());
552    }
553
554    #[test]
555    fn from_arc_collapses_dynstop() {
556        // from_arc(Arc<StopToken>) should reuse inner, not double-wrap
557        let inner = StopToken::new(Stopper::new());
558        let arc = alloc::sync::Arc::new(inner);
559        let stop = StopToken::from_arc(arc);
560        assert!(!stop.should_stop());
561    }
562
563    #[test]
564    fn stopper_inner_debug() {
565        let stop = Stopper::new();
566        let debug = alloc::format!("{:?}", stop);
567        assert!(debug.contains("cancelled"));
568    }
569
570    #[test]
571    fn sync_stopper_inner_debug() {
572        let stop = crate::SyncStopper::new();
573        let debug = alloc::format!("{:?}", stop);
574        assert!(debug.contains("cancelled"));
575    }
576
577    #[test]
578    fn from_stopper_clone_shares_state() {
579        let stopper = Stopper::new();
580        let stop: StopToken = stopper.clone().into();
581        let stop2 = stop.clone(); // Arc clone of the flattened inner
582
583        stopper.cancel();
584        // All three (stopper, stop, stop2) share the same AtomicBool
585        assert!(stop.should_stop());
586        assert!(stop2.should_stop());
587    }
588}
almost_enough/stop_token.rs

almost_enough/
stop_token.rs
