itertools_num/
ext.rs

1
2use std::ops::Add;
3use num_traits::{Zero};
4
5/// Extension trait for iterators: extra adaptors and methods
6/// for numerical iterators
7pub trait ItertoolsNum : Iterator {
8    /// Return an iterator that produces the sequence of cumulative sums
9    /// of the base iterator. The type of the sum is `S`.
10    fn cumsum<S>(self) -> Cumsum<Self, S>
11        where Self: Sized,
12              S: Add<Self::Item, Output=S>,
13              S: Zero,
14    {
15        cumsum(self)
16    }
17}
18
19impl<I: ?Sized> ItertoolsNum for I where I: Iterator { }
20
21pub struct Cumsum<I, S> {
22    sum: S,
23    iter: I,
24}
25
26fn cumsum<I, S>(iter: I) -> Cumsum<I, S>
27    where I: Iterator,
28          S: Add<I::Item, Output=S>,
29          S: Zero,
30{
31    Cumsum {
32        sum: Zero::zero(),
33        iter: iter,
34    }
35}
36
37impl<I, S> Iterator for Cumsum<I, S>
38    where I: Iterator,
39          S: Add<I::Item, Output=S>,
40          S: Zero + Clone,
41{
42    type Item = S;
43    fn next(&mut self) -> Option<Self::Item> {
44        let z = &mut self.sum;
45        self.iter.next().map(|x| {
46            *z = z.clone() + x;
47            z.clone()
48        })
49    }
50
51    fn size_hint(&self) -> (usize, Option<usize>) {
52        self.iter.size_hint()
53    }
54}
55
56impl<I, S> ExactSizeIterator for Cumsum<I, S>
57    where I: Iterator,
58          S: Add<I::Item, Output=S>,
59          S: Zero + Clone,
60{ }
61
62#[test]
63fn test_cumsum() {
64    let data = [1., 2., 3.];
65    let mut iter = data.iter().cumsum::<f64>();
66    assert_eq!(iter.next(), Some(1.));
67    assert_eq!(iter.next(), Some(3.));
68    assert_eq!(iter.next(), Some(6.));
69
70}