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

//! Iterators over `PrefixSum`.
//!
//! # Example
//!
//! ```
//! use prefix_sum::PrefixSum;
//!
//! let mut sum = PrefixSum::new(4);
//! sum[1..=2] += 2;
//!
//! let mut iter = sum.iter();
//! assert_eq!(iter.next(), Some(0));
//! assert_eq!(iter.next(), Some(2));
//! assert_eq!(iter.next(), Some(2));
//! assert_eq!(iter.next(), Some(0));
//! assert_eq!(iter.next(), None);
//! ```

use crate::{PrefixSum, Summable};

/// An iterator through a `PrefixSum` that clones each return value.
///
/// # Example
///
/// ```
/// use prefix_sum::PrefixSum;
///
/// let mut sum = PrefixSum::new(4);
/// sum[1..=2] += 3;
/// sum[2..] += 1;
///
/// let mut iter = sum.iter();
/// assert_eq!(iter.next(), Some(0));
/// assert_eq!(iter.next(), Some(3));
/// assert_eq!(iter.next(), Some(4));
/// assert_eq!(iter.next(), Some(1));
/// assert_eq!(iter.next(), None);
/// ```
pub struct Iter<'a, T: Summable + Clone> {
    sum: T,
    inner: std::slice::Iter<'a, T>,
}
impl<'a, T: Summable + Clone> Iterator for Iter<'a, T> {
    type Item = T;
    #[inline]
    fn next(&mut self) -> Option<T> {
        match self.inner.next() {
            None => None,
            Some(value) => {
                self.sum.add_assign_ref(value);
                Some(self.sum.clone())
            }
        }
    }
}
impl<'a, T: Summable + Clone> IntoIterator for &'a PrefixSum<T> {
    type IntoIter = Iter<'a, T>;
    type Item = T;
    #[inline]
    fn into_iter(self) -> Iter<'a, T> {
        Iter {
            sum: T::zero(),
            inner: self.values[..self.values.len() - 1].iter(),
        }
    }
}

/// An iterator through a `PrefixSum`.
///
/// # Example
///
/// ```
/// use prefix_sum::PrefixSum;
///
/// let mut sum = PrefixSum::new(4);
/// sum[1..=2] += 3;
/// sum[2..] += 1;
///
/// let mut iter = sum.into_iter();
/// assert_eq!(iter.next(), Some(0));
/// assert_eq!(iter.next(), Some(3));
/// assert_eq!(iter.next(), Some(4));
/// assert_eq!(iter.next(), Some(1));
/// assert_eq!(iter.next(), None);
/// ```
pub struct IntoIter<T: Summable> {
    sum: T,
    inner: std::vec::IntoIter<T>,
}
impl<T: Summable> Iterator for IntoIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        match self.inner.next() {
            None => None,
            Some(mut next) => {
                next.add_assign_ref(&self.sum);
                Some(std::mem::replace(&mut self.sum, next))
            }
        }
    }
}
impl<T: Summable> IntoIterator for PrefixSum<T> {
    type IntoIter = IntoIter<T>;
    type Item = T;
    #[inline]
    fn into_iter(self) -> IntoIter<T> {
        let mut iter = self.values.into_iter();
        match iter.next() {
            None => IntoIter {
                sum: T::zero(),
                inner: iter,
            },
            Some(sum) => IntoIter { inner: iter, sum },
        }
    }
}