Skip to content

Commit a35cdd4

Browse files
committed
Implement iter::Sum and iter::Product for Option
This is similar to the existing implementation for `Result`. It will take each item into the accumulator unless a `None` is returned.
1 parent f22dca0 commit a35cdd4

File tree

2 files changed

+129
-0
lines changed

2 files changed

+129
-0
lines changed

src/libcore/iter/traits/accum.rs

+113
Original file line numberDiff line numberDiff line change
@@ -223,3 +223,116 @@ impl<T, U, E> Product<Result<U, E>> for Result<T, E>
223223
ResultShunt::process(iter, |i| i.product())
224224
}
225225
}
226+
227+
/// An iterator adapter that produces output as long as the underlying
228+
/// iterator produces `Option::Some` values.
229+
struct OptionShunt<I> {
230+
iter: I,
231+
exited_early: bool,
232+
}
233+
234+
impl<I, T> OptionShunt<I>
235+
where
236+
I: Iterator<Item = Option<T>>,
237+
{
238+
/// Process the given iterator as if it yielded a `T` instead of a
239+
/// `Option<T>`. Any `None` value will stop the inner iterator and
240+
/// the overall result will be a `None`.
241+
pub fn process<F, U>(iter: I, mut f: F) -> Option<U>
242+
where
243+
F: FnMut(&mut Self) -> U,
244+
{
245+
let mut shunt = OptionShunt::new(iter);
246+
let value = f(shunt.by_ref());
247+
shunt.reconstruct(value)
248+
}
249+
250+
fn new(iter: I) -> Self {
251+
OptionShunt {
252+
iter,
253+
exited_early: false,
254+
}
255+
}
256+
257+
/// Consume the adapter and rebuild a `Option` value.
258+
fn reconstruct<U>(self, val: U) -> Option<U> {
259+
if self.exited_early {
260+
None
261+
} else {
262+
Some(val)
263+
}
264+
}
265+
}
266+
267+
impl<I, T> Iterator for OptionShunt<I>
268+
where
269+
I: Iterator<Item = Option<T>>,
270+
{
271+
type Item = T;
272+
273+
fn next(&mut self) -> Option<Self::Item> {
274+
match self.iter.next() {
275+
Some(Some(v)) => Some(v),
276+
Some(None) => {
277+
self.exited_early = true;
278+
None
279+
}
280+
None => None,
281+
}
282+
}
283+
284+
fn size_hint(&self) -> (usize, Option<usize>) {
285+
if self.exited_early {
286+
(0, Some(0))
287+
} else {
288+
let (_, upper) = self.iter.size_hint();
289+
(0, upper)
290+
}
291+
}
292+
}
293+
294+
#[stable(feature = "iter_arith_traits_option", since = "1.34.0")]
295+
impl<T, U> Sum<Option<U>> for Option<T>
296+
where
297+
T: Sum<U>,
298+
{
299+
/// Takes each element in the `Iterator`: if it is a `None`, no further
300+
/// elements are taken, and the `None` is returned. Should no `None` occur,
301+
/// the sum of all elements is returned.
302+
///
303+
/// # Examples
304+
///
305+
/// This sums up every integer in a vector, rejecting the sum if a negative
306+
/// element is encountered:
307+
///
308+
/// ```
309+
/// let v = vec![1, 2];
310+
/// let res: Option<i32> = v.iter().map(|&x: &i32|
311+
/// if x < 0 { None }
312+
/// else { Some(x) }
313+
/// ).sum();
314+
/// assert_eq!(res, Some(3));
315+
/// ```
316+
fn sum<I>(iter: I) -> Option<T>
317+
where
318+
I: Iterator<Item = Option<U>>,
319+
{
320+
OptionShunt::process(iter, |i| i.sum())
321+
}
322+
}
323+
324+
#[stable(feature = "iter_arith_traits_option", since = "1.34.0")]
325+
impl<T, U> Product<Option<U>> for Option<T>
326+
where
327+
T: Product<U>,
328+
{
329+
/// Takes each element in the `Iterator`: if it is a `None`, no further
330+
/// elements are taken, and the `None` is returned. Should no `None` occur,
331+
/// the product of all elements is returned.
332+
fn product<I>(iter: I) -> Option<T>
333+
where
334+
I: Iterator<Item = Option<U>>,
335+
{
336+
OptionShunt::process(iter, |i| i.product())
337+
}
338+
}

src/libcore/tests/iter.rs

+16
Original file line numberDiff line numberDiff line change
@@ -1066,6 +1066,14 @@ fn test_iterator_sum_result() {
10661066
assert_eq!(v.iter().cloned().sum::<Result<i32, _>>(), Err(()));
10671067
}
10681068

1069+
#[test]
1070+
fn test_iterator_sum_option() {
1071+
let v: &[Option<i32>] = &[Some(1), Some(2), Some(3), Some(4)];
1072+
assert_eq!(v.iter().cloned().sum::<Option<i32>>(), Some(10));
1073+
let v: &[Option<i32>] = &[Some(1), None, Some(3), Some(4)];
1074+
assert_eq!(v.iter().cloned().sum::<Option<i32>>(), None);
1075+
}
1076+
10691077
#[test]
10701078
fn test_iterator_product() {
10711079
let v: &[i32] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
@@ -1082,6 +1090,14 @@ fn test_iterator_product_result() {
10821090
assert_eq!(v.iter().cloned().product::<Result<i32, _>>(), Err(()));
10831091
}
10841092

1093+
#[test]
1094+
fn test_iterator_product_option() {
1095+
let v: &[Option<i32>] = &[Some(1), Some(2), Some(3), Some(4)];
1096+
assert_eq!(v.iter().cloned().product::<Option<i32>>(), Some(24));
1097+
let v: &[Option<i32>] = &[Some(1), None, Some(3), Some(4)];
1098+
assert_eq!(v.iter().cloned().product::<Option<i32>>(), None);
1099+
}
1100+
10851101
#[test]
10861102
fn test_iterator_max() {
10871103
let v: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];

0 commit comments

Comments
 (0)