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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use core::marker::PhantomData;
use super::*;

/// An iterator returned by [`from_fn`].
pub struct FromFn<F, R> {
    f: F,
    _marker: PhantomData<fn() -> R>,
}

/// A helper type used in [`from_fn`].
///
/// It represents a value that the iterator is stopped with, which `from_fn`
/// closure needs to pass back to the caller. It has a lifetime parameter to
/// prevent values from different scopes from getting mixed up (i.e. if
/// `from_fn` is used within another `from_fn`).
pub struct BreakValue<'a> {
    _invariant: PhantomData<fn(&'a ()) -> &'a ()>,
}

impl<F, R> InternalIterator for FromFn<F, R>
where
    F: for<'a> FnOnce(&mut dyn FnMut(R) -> ControlFlow<BreakValue<'a>>) -> ControlFlow<BreakValue<'a>>,
{
    type Item = R;

    fn try_for_each<RR, FF>(self, mut f: FF) -> ControlFlow<RR>
    where
        FF: FnMut(Self::Item) -> ControlFlow<RR>,
    {
        let mut result = ControlFlow::Continue(());
        (self.f)(&mut |item| {
            match f(item) {
                ControlFlow::Continue(()) => ControlFlow::Continue(()),
                ControlFlow::Break(res) => {
                    result = ControlFlow::Break(res);
                    ControlFlow::Break(BreakValue { _invariant: PhantomData })
                }
            }
        });
        result
    }
}

/// Creates an internal iterator from provided closure.
///
/// Provided closure should be equivalent to how
/// [`InternalIterator::try_for_each`] would be implemented - it should call its
/// parameter with every value that needs to be yielded from the iterator,
/// respect returned `ControlFlow` value, and return
/// `ControlFlow::Continue(())` at the end. Type signature ensures most of
/// that, you only need to take care to use `?` operator on every call when
/// yielding values.
///
/// If you want to construct an [`InternalIterator`] from a closure that yields
/// items one by one, like you would with [`std::iter::from_fn`], then you can
/// use that function with a conversion to internal iterator:
/// `std::iter::from_fn(f).into_internal()`.
///
/// ```
/// # use internal_iterator::InternalIterator;
/// # use std::ops::ControlFlow;
/// let x = 2;
/// let y = Some(3);
///
/// let iter = internal_iterator::from_fn(|f| {
///     f(1)?;
///     f(x)?;
///     if let Some(value) = y {
///         f(value)?;
///     }
///     ControlFlow::Continue(())
/// });
///
/// let values = iter.collect::<Vec<_>>();
/// assert_eq!(values, [1, 2, 3])
/// ```
///
/// Note that [`InternalIterator::try_for_each`] function is generic, but
/// generic closures are not possible. Therefore this function utilizes dynamic
/// dispatch (to be able to provide any function to the closure), and a marker
/// type for break value with the actual value passed via a side channel
/// (to be able to use any type as break value). Because of this, iterators
/// constructed by [`from_fn`] might be optimized more poorly. If the need
/// arises such iterators can always be rewritten as explicit structs with
/// generic implementations of [`InternalIterator::try_for_each`], although that
/// will require manually handling captured variables (whereas compiler does
/// that for you when using closures).
pub fn from_fn<F, R>(f: F) -> FromFn<F, R>
where
    F: for<'a> FnOnce(&mut dyn FnMut(R) -> ControlFlow<BreakValue<'a>>) -> ControlFlow<BreakValue<'a>>,
{
    FromFn { f, _marker: PhantomData }
}

#[test]
fn pipeline_tests() {
    fn check(iter: impl InternalIterator<Item = i32>, expect: &[i32]) {
        assert_eq!(iter.collect::<Vec<_>>(), expect);
    }
    fn make_iter() -> impl InternalIterator<Item = i32> {
        from_fn(|f| {
            f(1)?;
            f(2)?;
            f(3)?;
            f(4)?;
            ControlFlow::Continue(())
        })
    }
    check(make_iter(), &[1, 2, 3, 4]);
    check(make_iter().map(|x| x * 2), &[2, 4, 6, 8]);
    check(make_iter().filter(|&x| x > 2), &[3, 4]);
    check(make_iter().filter(|&x| x <= 2), &[1, 2]);
    check(make_iter().chain([7, 8, 9]), &[1, 2, 3, 4, 7, 8, 9]);
}

#[test]
fn can_be_static() {
    fn use_static_iterator(iter: impl InternalIterator<Item = i32> + 'static) {
        assert_eq!(iter.collect::<Vec<_>>(), &[1, 2, 3]);
    }
    let mut a = 1;
    let iter = from_fn(move |f| {
        f(a)?;
        a += 1;
        f(a)?;
        a += 1;
        f(a)?;
        ControlFlow::Continue(())
    });
    use_static_iterator(iter);
}

/// ```compile_fail
/// use std::ops::ControlFlow;
/// use internal_iterator::InternalIterator;
///
/// internal_iterator::from_fn(|f| {
///     f(1)?;
///     let mut slot = ControlFlow::Continue(());
///     internal_iterator::from_fn(|f2| {
///         slot = f2(42); // this assignment is supposed to fail to compile,
///                        // BreakValue can be moved between from_fn invocations
///         ControlFlow::Continue(())
///     }).collect::<Vec<_>>();
///     slot
/// });
/// ```
fn __compile_fail_check() {}