cumsum.rs

use std::ops::Add;

struct CumSum<T, I> {
    acc: Option<T>,
    iter: I,
}

impl<T, I> Iterator for CumSum<T, I>
where
    T: Add<Output = T> + Copy,
    I: Iterator<Item = T>
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.iter.next() {
            Some(next) => {
                self.acc = match self.acc {
                    Some(acc) => Some(acc + next),
                    None => Some(next)
                };
                self.acc
            }
            None => None
        }
    }
}

pub fn cumsum_iter<T, I>(iter: I) -> impl Iterator<Item = T>
where
    T: Add<Output = T> + Copy,
    I: Iterator<Item = T>,
{
    CumSum { iter, acc: None }
}

pub fn cumsum_scan<T, I>(iter: I) -> impl Iterator<Item = T>
where
    T: Add<Output = T> + Copy,
    I: Iterator<Item = T>,
{
    iter.scan(None, |st: &mut Option<T>, item: T| {
        *st = Some(match st {
            Some(st_inner) => *st_inner + item,
            None => item 
        });
        *st
    })
}

pub fn cumsum_vec<T: Add<Output = T> + Copy>(v: &[T]) -> Vec<T>
{
    let mut ret = Vec::from(v);
    for i in 1..ret.len() {
        ret[i] = ret[i] + ret[i-1];
    }
    ret
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn iter() {
        let mut iter = cumsum_iter(1..4);
        assert_eq!(iter.next().unwrap(), 1);
        assert_eq!(iter.next().unwrap(), 3);
        assert_eq!(iter.next().unwrap(), 6);
        assert!(iter.next().is_none());
    }

    #[test]
    fn scan() {
        let mut iter = cumsum_scan(1..4);
        assert_eq!(iter.next().unwrap(), 1);
        assert_eq!(iter.next().unwrap(), 3);
        assert_eq!(iter.next().unwrap(), 6);
        assert!(iter.next().is_none());
    }

    #[test]
    fn vec() {
        let iter = cumsum_vec(&vec![1, 2, 3]);
        assert_eq!(iter, vec![1, 3, 6]);
    }
}