unfold_iter/
lib.rs

1#![no_std]
2
3pub fn unfold<T, F>(init: T, next: F) -> Unfold<T, F>
4where
5    F: FnMut(&T) -> T,
6{
7    Unfold { state: init, next }
8}
9
10pub struct Unfold<T, F> {
11    state: T,
12    next: F,
13}
14
15impl<T, F> Iterator for Unfold<T, F>
16where
17    F: FnMut(&T) -> T,
18{
19    type Item = T;
20
21    fn next(&mut self) -> Option<Self::Item> {
22        let next = (self.next)(&self.state);
23        let prev = core::mem::replace(&mut self.state, next);
24        Some(prev)
25    }
26}
27
28pub fn try_unfold<T, F>(init: T, next: F) -> TryUnfold<T, F>
29where
30    F: FnMut(&T) -> Option<T>,
31{
32    let state = Some(init);
33    TryUnfold { state, next }
34}
35
36pub struct TryUnfold<T, F> {
37    state: Option<T>,
38    next: F,
39}
40
41impl<T, F> Iterator for TryUnfold<T, F>
42where
43    F: FnMut(&T) -> Option<T>,
44{
45    type Item = T;
46
47    fn next(&mut self) -> Option<Self::Item> {
48        let prev = self.state.take()?;
49        self.state = (self.next)(&prev);
50        Some(prev)
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[test]
59    fn test_unfold() {
60        let mut iter = unfold(0, |&x| x + 1);
61        assert_eq!(iter.next(), Some(0));
62        assert_eq!(iter.next(), Some(1));
63        assert_eq!(iter.next(), Some(2));
64    }
65
66    #[test]
67    fn test_try_unfold() {
68        let mut iter = try_unfold(0, |&x| (x < 2).then(|| x + 1));
69        assert_eq!(iter.next(), Some(0));
70        assert_eq!(iter.next(), Some(1));
71        assert_eq!(iter.next(), Some(2));
72        assert_eq!(iter.next(), None);
73    }
74}
unfold_iter/lib.rs

unfold_iter/
lib.rs
