iterx/
lib.rs

1use std::iter::{Map, Scan, Take, Zip};
2
3#[cfg(test)]
4use std::ops::{Add, Mul};
5
6pub fn option_lift<T: Copy>(f: &dyn Fn(&mut T, T) -> T) -> impl Fn(&mut T, T) -> Option<T> + '_ {
7    |acc, x| {
8        *acc = f(&mut *acc, x);
9        Some(*acc)
10    }
11}
12
13#[cfg(test)]
14fn w<T: Copy, O>(binop: &dyn Fn(T, T) -> O) -> impl Fn(T) -> O + '_ {
15    |x| binop(x, x)
16}
17
18// TODO
19// 1. Figure out &Add::add
20// 2. Use in rust-tx
21
22pub trait Iterx: Iterator + Clone {
23    fn drop_last(self) -> Take<Self>
24    where
25        Self: Sized,
26    {
27        let m = self.clone().count() - 1;
28        self.take(m)
29    }
30
31    fn mark_last(self) -> MarkLast<Self> {
32        MarkLast(false, self.peekable())
33    }
34
35    /// * Name scan_ so as to not collide with [`scan`][std::iter::Iterator::scan]
36    /// * [`scan`][std::iter::Iterator::scan] has been renamed in this library to
37    fn scan_<F>(self, f: F) -> Scan_<Self, Self::Item, F>
38    where
39        Self: Sized,
40        F: FnMut(Self::Item, Self::Item) -> Self::Item,
41    {
42        Scan_::new(self, f)
43    }
44
45    /// * A renaming of [`scan`][std::iter::Iterator::scan]
46    fn scan_while<St, B, F>(self, initial_state: St, f: F) -> Scan<Self, St, F>
47    where
48        Self: Sized,
49        F: FnMut(&mut St, Self::Item) -> Option<B>,
50    {
51        self.scan(initial_state, f)
52    }
53
54    /// * A convenience function for `std::iter::once(value).chain(self)`
55    fn prepend(
56        self,
57        value: Self::Item,
58    ) -> std::iter::Chain<std::iter::Once<<Self as Iterator>::Item>, Self>
59    where
60        Self: Sized,
61    {
62        std::iter::once(value).chain(self)
63    }
64
65    /// * The analog to `scan_` that takes an initial value.
66    fn prescan<St, F>(self, initial_state: St, f: F) -> Prescan<Self, St, F>
67    where
68        Self: Sized,
69        F: FnMut(St, Self::Item) -> St,
70    {
71        Prescan::new(self, initial_state, f)
72    }
73
74    /// * A convenience function that fuses [`map`][std::iter::Iterator::map] and [`zip`][std::iter::Iterator::zip]
75    fn zip_map<U, T, F>(
76        self,
77        other: U,
78        f: F,
79    ) -> Map<
80        Zip<Self, U::IntoIter>,
81        Box<dyn Fn((<Self as Iterator>::Item, <U as IntoIterator>::Item)) -> T>,
82    >
83    where
84        Self: Sized,
85        U: IntoIterator,
86        F: Fn(<Self as Iterator>::Item, <U as IntoIterator>::Item) -> T + 'static,
87    {
88        self.zip(other.into_iter())
89            .map(Box::new(move |(x, y)| f(x, y)))
90    }
91}
92
93impl<T: ?Sized + std::clone::Clone> Iterx for T where T: Iterator {}
94
95#[derive(Clone)]
96pub struct Prescan<I, St, F> {
97    iter: I,
98    f: F,
99    state: Option<St>,
100}
101
102impl<I, St, F> Prescan<I, St, F> {
103    fn new(iter: I, initial_state: St, f: F) -> Self {
104        Self {
105            iter,
106            f,
107            state: Some(initial_state),
108        }
109    }
110}
111
112impl<I, St: Clone, F> Iterator for Prescan<I, St, F>
113where
114    I: Iterator,
115    F: FnMut(St, I::Item) -> St,
116{
117    type Item = St;
118
119    fn next(&mut self) -> Option<Self::Item> {
120        let state = self.state.take()?;
121
122        if let Some(x) = self.iter.next() {
123            self.state = Some((self.f)(state.clone(), x));
124        }
125
126        Some(state)
127    }
128
129    fn size_hint(&self) -> (usize, Option<usize>) {
130        let (lb, ub) = self.iter.size_hint();
131        (lb + 1, ub.map(|ub| ub + 1))
132    }
133}
134
135#[derive(Clone)]
136pub struct Scan_<I, T, F> {
137    iter: I,
138    f: F,
139    state: Option<T>,
140    is_first: bool,
141}
142
143impl<I, T, F> Scan_<I, T, F>
144where
145    I: Iterator<Item = T>,
146{
147    fn new(iter: I, f: F) -> Self {
148        Self {
149            iter,
150            f,
151            state: None,
152            is_first: true,
153        }
154    }
155}
156
157impl<I, T: Clone, F> Iterator for Scan_<I, T, F>
158where
159    I: Iterator<Item = T>,
160    F: FnMut(T, T) -> T,
161{
162    type Item = T;
163
164    fn next(&mut self) -> Option<Self::Item> {
165        if self.is_first {
166            self.state = Some(self.iter.next())?;
167            self.is_first = false;
168        }
169
170        let state = self.state.take()?;
171
172        if let Some(x) = self.iter.next() {
173            self.state = Some((self.f)(state.clone(), x));
174        }
175
176        Some(state)
177    }
178
179    fn size_hint(&self) -> (usize, Option<usize>) {
180        self.iter.size_hint()
181    }
182}
183
184pub struct MarkLast<I>(bool, std::iter::Peekable<I>)
185where
186    I: Iterator;
187
188impl<I> Iterator for MarkLast<I>
189where
190    I: Iterator,
191{
192    type Item = (bool, I::Item);
193
194    fn next(&mut self) -> Option<Self::Item> {
195        self.1.next().map(|e| (self.1.peek().is_none(), e))
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    use itertools::assert_equal;
204
205    #[test]
206    fn test_scan_while() {
207        assert_equal(
208            vec![1, 1, 1]
209                .into_iter()
210                .scan_while(0, &option_lift(&|a, b| *a + b)),
211            1..=3,
212        );
213    }
214
215    #[test]
216    fn test_w_combinator() {
217        assert_eq!(w(&Add::add)(1), 2);
218    }
219
220    #[test]
221    fn test_scan_() {
222        assert_equal(vec![1].into_iter().scan_(|x, y| x + y), 1..2);
223        // assert_equal(vec![1].into_iter().scan_(&Add::add), 1..2); TODO
224        assert_equal(vec![1, 1, 1].into_iter().scan_(|x, y| x + y), 1..=3);
225        assert_equal((1..=5).scan_(|x, y| x + y), vec![1, 3, 6, 10, 15]);
226    }
227
228    #[test]
229    fn test_prescan() {
230        assert_equal(vec![1].into_iter().prescan(0, |x, y| x + y), 0..=1);
231        assert_equal(vec![1, 1, 1].into_iter().prescan(0, |x, y| x + y), 0..=3);
232        assert_equal((1..=5).prescan(0, |x, y| x + y), vec![0, 1, 3, 6, 10, 15]);
233        assert_equal((1..=5).prescan(0, |x, y| x + y).skip(2), vec![3, 6, 10, 15]);
234    }
235
236    #[test]
237    fn test_zip_map() {
238        assert_equal((1..5).zip_map(1..5, |a, b| a + b), vec![2, 4, 6, 8]);
239        assert_equal((1..5).zip_map(1..5, &Add::add), vec![2, 4, 6, 8]);
240        assert_equal((1..5).zip_map(1..5, |a, b| a * b), vec![1, 4, 9, 16]);
241        assert_equal((1..5).zip_map(1..5, &Mul::mul), vec![1, 4, 9, 16]);
242        assert_equal((1..5).zip_map(2..6, &std::cmp::max), 2..6);
243        assert_equal((1..5).zip_map(2..6, &std::cmp::min), 1..5);
244    }
245
246    #[test]
247    fn test_drop_last() {
248        assert_equal((1..4).drop_last(), vec![1, 2]);
249        assert_equal((1..5).drop_last(), vec![1, 2, 3]);
250    }
251
252    #[test]
253    fn test_prepend() {
254        assert_equal((1..4).prepend(0), 0..4);
255        assert_equal((2..5).prepend(1), 1..5);
256    }
257
258    #[test]
259    fn test_mark_last() {
260        assert_equal((1..3).mark_last(), vec![(false, 1), (true, 2)]);
261        assert_equal((1..4).mark_last(), vec![(false, 1), (false, 2), (true, 3)]);
262    }
263}
iterx/lib.rs

iterx/
lib.rs
