1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
#[cfg(feature = "parallel")]
mod stepped {
    use crate::parallel::num_threads;

    /// An iterator adaptor to allow running computations using [`in_parallel()`][crate::parallel::in_parallel()] in a step-wise manner, see the [module docs][crate::parallel]
    /// for details.
    pub struct Stepwise<Reduce: super::Reduce> {
        /// This field is first to assure it's dropped first and cause threads that are dropped next to stop their loops
        /// as sending results fails when the receiver is dropped.
        receive_result: std::sync::mpsc::Receiver<Reduce::Input>,
        /// `join()` will be called on these guards to assure every thread tries to send through a closed channel. When
        /// that happens, they break out of their loops.
        threads: Vec<std::thread::JoinHandle<()>>,
        /// The reducer is called only in the thread using the iterator, dropping it has no side effects.
        reducer: Option<Reduce>,
    }

    impl<Reduce: super::Reduce> Drop for Stepwise<Reduce> {
        fn drop(&mut self) {
            let (_, sink) = std::sync::mpsc::channel();
            drop(std::mem::replace(&mut self.receive_result, sink));

            let mut last_err = None;
            for handle in std::mem::take(&mut self.threads) {
                if let Err(err) = handle.join() {
                    last_err = Some(err);
                };
            }
            if let Some(thread_err) = last_err {
                std::panic::resume_unwind(thread_err);
            }
        }
    }

    impl<Reduce: super::Reduce> Stepwise<Reduce> {
        /// Instantiate a new iterator and start working in threads.
        /// For a description of parameters, see [`in_parallel()`][crate::parallel::in_parallel()].
        pub fn new<InputIter, ThreadStateFn, ConsumeFn, I, O, S>(
            input: InputIter,
            thread_limit: Option<usize>,
            new_thread_state: ThreadStateFn,
            consume: ConsumeFn,
            reducer: Reduce,
        ) -> Self
        where
            InputIter: Iterator<Item = I> + Send + 'static,
            ThreadStateFn: Fn(usize) -> S + Send + Clone + 'static,
            ConsumeFn: Fn(I, &mut S) -> O + Send + Clone + 'static,
            Reduce: super::Reduce<Input = O> + 'static,
            I: Send + 'static,
            O: Send + 'static,
        {
            let num_threads = num_threads(thread_limit);
            let mut threads = Vec::with_capacity(num_threads + 1);
            let receive_result = {
                let (send_input, receive_input) = crossbeam_channel::bounded::<I>(num_threads);
                let (send_result, receive_result) = std::sync::mpsc::sync_channel::<O>(num_threads);
                for thread_id in 0..num_threads {
                    let handle = std::thread::spawn({
                        let send_result = send_result.clone();
                        let receive_input = receive_input.clone();
                        let new_thread_state = new_thread_state.clone();
                        let consume = consume.clone();
                        move || {
                            let mut state = new_thread_state(thread_id);
                            for item in receive_input {
                                if send_result.send(consume(item, &mut state)).is_err() {
                                    break;
                                }
                            }
                        }
                    });
                    threads.push(handle);
                }
                threads.push(std::thread::spawn(move || {
                    for item in input {
                        if send_input.send(item).is_err() {
                            break;
                        }
                    }
                }));
                receive_result
            };
            Stepwise {
                threads,
                receive_result,
                reducer: Some(reducer),
            }
        }

        /// Consume the iterator by finishing its iteration and calling [`Reduce::finalize()`][crate::parallel::Reduce::finalize()].
        pub fn finalize(mut self) -> Result<Reduce::Output, Reduce::Error> {
            for value in self.by_ref() {
                drop(value?);
            }
            self.reducer
                .take()
                .expect("this is the last call before consumption")
                .finalize()
        }
    }

    impl<Reduce: super::Reduce> Iterator for Stepwise<Reduce> {
        type Item = Result<Reduce::FeedProduce, Reduce::Error>;

        fn next(&mut self) -> Option<<Self as Iterator>::Item> {
            self.receive_result
                .recv()
                .ok()
                .and_then(|input| self.reducer.as_mut().map(|r| r.feed(input)))
        }
    }

    impl<R: super::Reduce> super::Finalize for Stepwise<R> {
        type Reduce = R;

        fn finalize(
            self,
        ) -> Result<
            <<Self as super::Finalize>::Reduce as super::Reduce>::Output,
            <<Self as super::Finalize>::Reduce as super::Reduce>::Error,
        > {
            Stepwise::finalize(self)
        }
    }
}

#[cfg(not(feature = "parallel"))]
mod stepped {
    /// An iterator adaptor to allow running computations using [`in_parallel()`][crate::parallel::in_parallel()] in a step-wise manner, see the [module docs][crate::parallel]
    /// for details.
    pub struct Stepwise<InputIter, ConsumeFn, ThreadState, Reduce> {
        input: InputIter,
        consume: ConsumeFn,
        thread_state: ThreadState,
        reducer: Reduce,
    }

    impl<InputIter, ConsumeFn, Reduce, I, O, S> Stepwise<InputIter, ConsumeFn, S, Reduce>
    where
        InputIter: Iterator<Item = I>,
        ConsumeFn: Fn(I, &mut S) -> O,
        Reduce: super::Reduce<Input = O>,
    {
        /// Instantiate a new iterator.
        /// For a description of parameters, see [`in_parallel()`][crate::parallel::in_parallel()].
        pub fn new<ThreadStateFn>(
            input: InputIter,
            _thread_limit: Option<usize>,
            new_thread_state: ThreadStateFn,
            consume: ConsumeFn,
            reducer: Reduce,
        ) -> Self
        where
            ThreadStateFn: Fn(usize) -> S,
        {
            Stepwise {
                input,
                consume,
                thread_state: new_thread_state(0),
                reducer,
            }
        }

        /// Consume the iterator by finishing its iteration and calling [`Reduce::finalize()`][crate::parallel::Reduce::finalize()].
        pub fn finalize(mut self) -> Result<Reduce::Output, Reduce::Error> {
            for value in self.by_ref() {
                drop(value?);
            }
            self.reducer.finalize()
        }
    }

    impl<InputIter, ConsumeFn, ThreadState, Reduce, I, O> Iterator for Stepwise<InputIter, ConsumeFn, ThreadState, Reduce>
    where
        InputIter: Iterator<Item = I>,
        ConsumeFn: Fn(I, &mut ThreadState) -> O,
        Reduce: super::Reduce<Input = O>,
    {
        type Item = Result<Reduce::FeedProduce, Reduce::Error>;

        fn next(&mut self) -> Option<<Self as Iterator>::Item> {
            self.input
                .next()
                .map(|input| self.reducer.feed((self.consume)(input, &mut self.thread_state)))
        }
    }

    impl<InputIter, ConsumeFn, R, I, O, S> super::Finalize for Stepwise<InputIter, ConsumeFn, S, R>
    where
        InputIter: Iterator<Item = I>,
        ConsumeFn: Fn(I, &mut S) -> O,
        R: super::Reduce<Input = O>,
    {
        type Reduce = R;

        fn finalize(
            self,
        ) -> Result<
            <<Self as super::Finalize>::Reduce as super::Reduce>::Output,
            <<Self as super::Finalize>::Reduce as super::Reduce>::Error,
        > {
            Stepwise::finalize(self)
        }
    }
}

use std::marker::PhantomData;

pub use stepped::Stepwise;

/// An trait for aggregating items commonly produced in threads into a single result, without itself
/// needing to be thread safe.
pub trait Reduce {
    /// The type fed to the reducer in the [`feed()`][Reduce::feed()] method.
    ///
    /// It's produced by a function that may run on multiple threads.
    type Input;
    /// The type produced in Ok(…) by [`feed()`][Reduce::feed()].
    /// Most reducers by nature use `()` here as the value is in the aggregation.
    /// However, some may use it to collect statistics only and return their Input
    /// in some form as a result here for [`Stepwise`] to be useful.
    type FeedProduce;
    /// The type produced once by the [`finalize()`][Reduce::finalize()] method.
    ///
    /// For traditional reducers, this is the value produced by the entire operation.
    /// For those made for step-wise iteration this may be aggregated statistics.
    type Output;
    /// The error type to use for all methods of this trait.
    type Error;
    /// Called each time a new `item` was produced in order to aggregate it into the final result.
    ///
    /// If an `Error` is returned, the entire operation will be stopped.
    fn feed(&mut self, item: Self::Input) -> Result<Self::FeedProduce, Self::Error>;
    /// Called once once all items where passed to `feed()`, producing the final `Output` of the operation or an `Error`.
    fn finalize(self) -> Result<Self::Output, Self::Error>;
}

/// An identity reducer for those who want to use [`Stepwise`] or [`in_parallel()`][crate::parallel::in_parallel()]
/// without the use of non-threaded reduction of products created in threads.
pub struct IdentityWithResult<Input, Error> {
    _input: PhantomData<Input>,
    _error: PhantomData<Error>,
}

impl<Input, Error> Default for IdentityWithResult<Input, Error> {
    fn default() -> Self {
        IdentityWithResult {
            _input: Default::default(),
            _error: Default::default(),
        }
    }
}

impl<Input, Error> Reduce for IdentityWithResult<Input, Error> {
    type Input = Result<Input, Self::Error>;
    type FeedProduce = Input;
    type Output = ();
    type Error = Error;

    fn feed(&mut self, item: Self::Input) -> Result<Self::FeedProduce, Self::Error> {
        item
    }

    fn finalize(self) -> Result<Self::Output, Self::Error> {
        Ok(())
    }
}

/// A trait reflecting the `finalize()` method of [`Reduce`] implementations
pub trait Finalize {
    /// An implementation of [`Reduce`]
    type Reduce: self::Reduce;

    /// Similar to the [`Reduce::finalize()`] method
    fn finalize(
        self,
    ) -> Result<<<Self as Finalize>::Reduce as self::Reduce>::Output, <<Self as Finalize>::Reduce as self::Reduce>::Error>;
}