Implemented 'panic_fuse' operation

asparit/src/core/iterator.rs

@@ -25,6 +25,7 @@ use crate::{
         map_with::MapWith,
         max::{Max, MaxBy, MaxByKey},
         min::{Min, MinBy, MinByKey},
+        panic_fuse::PanicFuse,
         product::Product,
         reduce::{Reduce, ReduceWith},
         sum::Sum,
@@ -1529,6 +1530,40 @@ pub trait ParallelIterator<'a>: Sized + Send {
         WhileSome::new(self)
     }
 
+    /// Wraps an iterator with a fuse in case of panics, to halt all threads
+    /// as soon as possible.
+    ///
+    /// Panics within parallel iterators are always propagated to the caller,
+    /// but they don't always halt the rest of the iterator right away, due to
+    /// the internal semantics of [`join`]. This adaptor makes a greater effort
+    /// to stop processing other items sooner, with the cost of additional
+    /// synchronization overhead, which may also inhibit some optimizations.
+    ///
+    /// [`join`]: ../fn.join.html#panics
+    ///
+    /// # Examples
+    ///
+    /// If this code didn't use `panic_fuse()`, it would continue processing
+    /// many more items in other threads (with long sleep delays) before the
+    /// panic is finally propagated.
+    ///
+    /// ```should_panic
+    /// use rayon::prelude::*;
+    /// use std::{thread, time};
+    ///
+    /// (0..1_000_000)
+    ///     .into_par_iter()
+    ///     .panic_fuse()
+    ///     .for_each(|i| {
+    ///         // simulate some work
+    ///         thread::sleep(time::Duration::from_secs(1));
+    ///         assert!(i > 0); // oops!
+    ///     });
+    /// ```
+    fn panic_fuse(self) -> PanicFuse<Self> {
+        PanicFuse::new(self)
+    }
+
     /// Creates a fresh collection containing all the elements produced
     /// by this parallel iterator.
     ///

asparit/src/inner/mod.rs

@@ -16,6 +16,7 @@ pub mod map_with;
 pub mod max;
 pub mod min;
 pub mod noop;
+pub mod panic_fuse;
 pub mod product;
 pub mod reduce;
 pub mod sum;
@@ -54,6 +55,7 @@ mod tests {
             .chain(b)
             .update(|x| x.push(0))
             .flatten_iter()
+            .panic_fuse()
             .map(Some)
             .while_some()
             .map_init(

asparit/src/inner/panic_fuse.rs

@@ -0,0 +1,401 @@
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
+use std::thread::panicking;
+
+use crate::{
+    Consumer, Executor, Folder, IndexedParallelIterator, IndexedProducer, IndexedProducerCallback,
+    ParallelIterator, Producer, ProducerCallback, Reducer,
+};
+
+/* PanicFuse */
+
+pub struct PanicFuse<X> {
+    base: X,
+}
+
+impl<X> PanicFuse<X> {
+    pub fn new(base: X) -> Self {
+        Self { base }
+    }
+}
+
+impl<'a, X> ParallelIterator<'a> for PanicFuse<X>
+where
+    X: ParallelIterator<'a>,
+{
+    type Item = X::Item;
+
+    fn drive<E, C, D, R>(self, executor: E, base: C) -> E::Result
+    where
+        E: Executor<'a, D>,
+        C: Consumer<Self::Item, Result = D, Reducer = R> + 'a,
+        D: Send + 'a,
+        R: Reducer<D> + Send + 'a,
+    {
+        let consumer = PanicFuseConsumer {
+            base,
+            fuse: Fuse::default(),
+        };
+
+        self.base.drive(executor, consumer)
+    }
+
+    fn with_producer<CB>(self, base: CB) -> CB::Output
+    where
+        CB: ProducerCallback<'a, Self::Item>,
+    {
+        self.base.with_producer(PanicFuseCallback { base })
+    }
+
+    fn len_hint_opt(&self) -> Option<usize> {
+        self.base.len_hint_opt()
+    }
+}
+
+impl<'a, X> IndexedParallelIterator<'a> for PanicFuse<X>
+where
+    X: IndexedParallelIterator<'a>,
+{
+    fn drive_indexed<E, C, D, R>(self, executor: E, base: C) -> E::Result
+    where
+        E: Executor<'a, D>,
+        C: Consumer<Self::Item, Result = D, Reducer = R> + 'a,
+        D: Send + 'a,
+        R: Reducer<D> + Send + 'a,
+    {
+        let consumer = PanicFuseConsumer {
+            base,
+            fuse: Fuse::default(),
+        };
+
+        self.base.drive_indexed(executor, consumer)
+    }
+
+    fn with_producer_indexed<CB>(self, base: CB) -> CB::Output
+    where
+        CB: IndexedProducerCallback<'a, Self::Item>,
+    {
+        self.base.with_producer_indexed(PanicFuseCallback { base })
+    }
+
+    fn len_hint(&self) -> usize {
+        self.base.len_hint()
+    }
+}
+
+/* Fuse */
+
+struct Fuse(Arc<AtomicBool>);
+
+impl Fuse {
+    #[inline]
+    fn panicked(&self) -> bool {
+        self.0.load(Ordering::Relaxed)
+    }
+}
+
+impl Default for Fuse {
+    fn default() -> Self {
+        Self(Arc::new(AtomicBool::new(false)))
+    }
+}
+
+impl Clone for Fuse {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl Drop for Fuse {
+    fn drop(&mut self) {
+        if panicking() {
+            self.0.store(true, Ordering::Relaxed);
+        }
+    }
+}
+
+/* PanicFuseCallback */
+
+struct PanicFuseCallback<CB> {
+    base: CB,
+}
+
+impl<'a, CB, T> ProducerCallback<'a, T> for PanicFuseCallback<CB>
+where
+    CB: ProducerCallback<'a, T>,
+{
+    type Output = CB::Output;
+
+    fn callback<P>(self, base: P) -> Self::Output
+    where
+        P: Producer<Item = T> + 'a,
+    {
+        self.base.callback(PanicFuseProducer {
+            base,
+            fuse: Fuse::default(),
+        })
+    }
+}
+
+impl<'a, CB, T> IndexedProducerCallback<'a, T> for PanicFuseCallback<CB>
+where
+    CB: IndexedProducerCallback<'a, T>,
+{
+    type Output = CB::Output;
+
+    fn callback<P>(self, base: P) -> Self::Output
+    where
+        P: IndexedProducer<Item = T> + 'a,
+    {
+        self.base.callback(PanicFuseProducer {
+            base,
+            fuse: Fuse::default(),
+        })
+    }
+}
+
+/* PanicFuseProducer */
+
+struct PanicFuseProducer<P> {
+    base: P,
+    fuse: Fuse,
+}
+
+impl<P> Producer for PanicFuseProducer<P>
+where
+    P: Producer,
+{
+    type Item = P::Item;
+    type IntoIter = PanicFuseIter<P::IntoIter>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        PanicFuseIter {
+            base: self.base.into_iter(),
+            fuse: self.fuse,
+        }
+    }
+
+    fn split(self) -> (Self, Option<Self>) {
+        let fuse = self.fuse;
+        let (left, right) = self.base.split();
+
+        let left = Self {
+            base: left,
+            fuse: fuse.clone(),
+        };
+        let right = right.map(|base| Self { base, fuse });
+
+        (left, right)
+    }
+
+    fn splits(&self) -> Option<usize> {
+        self.base.splits()
+    }
+
+    fn fold_with<F>(self, base: F) -> F
+    where
+        F: Folder<Self::Item>,
+    {
+        let folder = PanicFuseFolder {
+            base,
+            fuse: self.fuse,
+        };
+
+        self.base.fold_with(folder).base
+    }
+}
+
+impl<P> IndexedProducer for PanicFuseProducer<P>
+where
+    P: IndexedProducer,
+{
+    type Item = P::Item;
+    type IntoIter = PanicFuseIter<P::IntoIter>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        PanicFuseIter {
+            base: self.base.into_iter(),
+            fuse: self.fuse,
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.base.len()
+    }
+
+    fn split_at(self, index: usize) -> (Self, Self) {
+        let (left, right) = self.base.split_at(index);
+
+        let left = Self {
+            base: left,
+            fuse: self.fuse.clone(),
+        };
+        let right = Self {
+            base: right,
+            fuse: self.fuse,
+        };
+
+        (left, right)
+    }
+
+    fn splits(&self) -> Option<usize> {
+        self.base.splits()
+    }
+
+    fn min_len(&self) -> Option<usize> {
+        self.base.min_len()
+    }
+
+    fn max_len(&self) -> Option<usize> {
+        self.base.max_len()
+    }
+
+    fn fold_with<F>(self, base: F) -> F
+    where
+        F: Folder<Self::Item>,
+    {
+        let folder = PanicFuseFolder {
+            base,
+            fuse: self.fuse.clone(),
+        };
+
+        self.base.fold_with(folder).base
+    }
+}
+
+/* PanicFuseIter */
+
+struct PanicFuseIter<I> {
+    base: I,
+    fuse: Fuse,
+}
+
+impl<I> Iterator for PanicFuseIter<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.fuse.panicked() {
+            None
+        } else {
+            self.base.next()
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.base.size_hint()
+    }
+}
+
+impl<I> DoubleEndedIterator for PanicFuseIter<I>
+where
+    I: DoubleEndedIterator,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.fuse.panicked() {
+            None
+        } else {
+            self.base.next_back()
+        }
+    }
+}
+
+impl<I> ExactSizeIterator for PanicFuseIter<I> where I: ExactSizeIterator {}
+
+/* PanicFuseConsumer */
+
+struct PanicFuseConsumer<C> {
+    base: C,
+    fuse: Fuse,
+}
+
+impl<C, T> Consumer<T> for PanicFuseConsumer<C>
+where
+    C: Consumer<T>,
+{
+    type Folder = PanicFuseFolder<C::Folder>;
+    type Reducer = C::Reducer;
+    type Result = C::Result;
+
+    fn split(self) -> (Self, Self, Self::Reducer) {
+        let (left, right, reducer) = self.base.split();
+
+        let left = Self {
+            base: left,
+            fuse: self.fuse.clone(),
+        };
+        let right = Self {
+            base: right,
+            fuse: self.fuse,
+        };
+
+        (left, right, reducer)
+    }
+
+    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
+        let (left, right, reducer) = self.base.split_at(index);
+
+        let left = Self {
+            base: left,
+            fuse: self.fuse.clone(),
+        };
+        let right = Self {
+            base: right,
+            fuse: self.fuse,
+        };
+
+        (left, right, reducer)
+    }
+
+    fn into_folder(self) -> Self::Folder {
+        PanicFuseFolder {
+            base: self.base.into_folder(),
+            fuse: self.fuse,
+        }
+    }
+
+    fn is_full(&self) -> bool {
+        self.base.is_full()
+    }
+}
+
+/* PanicFuseFolder */
+
+struct PanicFuseFolder<F> {
+    base: F,
+    fuse: Fuse,
+}
+
+impl<F, T> Folder<T> for PanicFuseFolder<F>
+where
+    F: Folder<T>,
+{
+    type Result = F::Result;
+
+    fn consume(mut self, item: T) -> Self {
+        self.base = self.base.consume(item);
+
+        self
+    }
+
+    fn consume_iter<X>(mut self, iter: X) -> Self
+    where
+        X: IntoIterator<Item = T>,
+    {
+        self.base = self.base.consume_iter(iter);
+
+        self
+    }
+
+    fn complete(self) -> Self::Result {
+        self.base.complete()
+    }
+
+    fn is_full(&self) -> bool {
+        self.base.is_full() || self.fuse.panicked()
+    }
+}