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


use std::ops::Add;
use num_traits::{Zero};

/// Extension trait for iterators: extra adaptors and methods
/// for numerical iterators
pub trait ItertoolsNum : Iterator {
    /// Return an iterator that produces the sequence of cumulative sums
    /// of the base iterator. The type of the sum is `S`.
    fn cumsum<S>(self) -> Cumsum<Self, S>
        where Self: Sized,
              S: Add<Self::Item, Output=S>,
              S: Zero,
    {
        cumsum(self)
    }
}

impl<I: ?Sized> ItertoolsNum for I where I: Iterator { }

pub struct Cumsum<I, S> {
    sum: S,
    iter: I,
}

fn cumsum<I, S>(iter: I) -> Cumsum<I, S>
    where I: Iterator,
          S: Add<I::Item, Output=S>,
          S: Zero,
{
    Cumsum {
        sum: Zero::zero(),
        iter: iter,
    }
}

impl<I, S> Iterator for Cumsum<I, S>
    where I: Iterator,
          S: Add<I::Item, Output=S>,
          S: Zero + Clone,
{
    type Item = S;
    fn next(&mut self) -> Option<Self::Item> {
        let z = &mut self.sum;
        self.iter.next().map(|x| {
            *z = z.clone() + x;
            z.clone()
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<I, S> ExactSizeIterator for Cumsum<I, S>
    where I: Iterator,
          S: Add<I::Item, Output=S>,
          S: Zero + Clone,
{ }

#[test]
fn test_cumsum() {
    let data = [1., 2., 3.];
    let mut iter = data.iter().cumsum::<f64>();
    assert_eq!(iter.next(), Some(1.));
    assert_eq!(iter.next(), Some(3.));
    assert_eq!(iter.next(), Some(6.));

}