internal_iterator/
from_fn_impl.rs

1use core::marker::PhantomData;
2use super::*;
3
4/// An iterator returned by [`from_fn`].
5pub struct FromFn<F, R> {
6    f: F,
7    _marker: PhantomData<fn() -> R>,
8}
9
10/// A helper type used in [`from_fn`].
11///
12/// It represents a value that the iterator is stopped with, which `from_fn`
13/// closure needs to pass back to the caller. It has a lifetime parameter to
14/// prevent values from different scopes from getting mixed up (i.e. if
15/// `from_fn` is used within another `from_fn`).
16pub struct BreakValue<'a> {
17    _invariant: PhantomData<fn(&'a ()) -> &'a ()>,
18}
19
20impl<F, R> InternalIterator for FromFn<F, R>
21where
22    F: for<'a> FnOnce(&mut dyn FnMut(R) -> ControlFlow<BreakValue<'a>>) -> ControlFlow<BreakValue<'a>>,
23{
24    type Item = R;
25
26    fn try_for_each<RR, FF>(self, mut f: FF) -> ControlFlow<RR>
27    where
28        FF: FnMut(Self::Item) -> ControlFlow<RR>,
29    {
30        let mut result = ControlFlow::Continue(());
31        (self.f)(&mut |item| {
32            match f(item) {
33                ControlFlow::Continue(()) => ControlFlow::Continue(()),
34                ControlFlow::Break(res) => {
35                    result = ControlFlow::Break(res);
36                    ControlFlow::Break(BreakValue { _invariant: PhantomData })
37                }
38            }
39        });
40        result
41    }
42}
43
44/// Creates an internal iterator from provided closure.
45///
46/// Provided closure should be equivalent to how
47/// [`InternalIterator::try_for_each`] would be implemented - it should call its
48/// parameter with every value that needs to be yielded from the iterator,
49/// respect returned `ControlFlow` value, and return
50/// `ControlFlow::Continue(())` at the end. Type signature ensures most of
51/// that, you only need to take care to use `?` operator on every call when
52/// yielding values.
53///
54/// If you want to construct an [`InternalIterator`] from a closure that yields
55/// items one by one, like you would with [`std::iter::from_fn`], then you can
56/// use that function with a conversion to internal iterator:
57/// `std::iter::from_fn(f).into_internal()`.
58///
59/// ```
60/// # use internal_iterator::InternalIterator;
61/// # use std::ops::ControlFlow;
62/// let x = 2;
63/// let y = Some(3);
64///
65/// let iter = internal_iterator::from_fn(|f| {
66///     f(1)?;
67///     f(x)?;
68///     if let Some(value) = y {
69///         f(value)?;
70///     }
71///     ControlFlow::Continue(())
72/// });
73///
74/// let values = iter.collect::<Vec<_>>();
75/// assert_eq!(values, [1, 2, 3])
76/// ```
77///
78/// Note that [`InternalIterator::try_for_each`] function is generic, but
79/// generic closures are not possible. Therefore this function utilizes dynamic
80/// dispatch (to be able to provide any function to the closure), and a marker
81/// type for break value with the actual value passed via a side channel
82/// (to be able to use any type as break value). Because of this, iterators
83/// constructed by [`from_fn`] might be optimized more poorly. If the need
84/// arises such iterators can always be rewritten as explicit structs with
85/// generic implementations of [`InternalIterator::try_for_each`], although that
86/// will require manually handling captured variables (whereas compiler does
87/// that for you when using closures).
88pub fn from_fn<F, R>(f: F) -> FromFn<F, R>
89where
90    F: for<'a> FnOnce(&mut dyn FnMut(R) -> ControlFlow<BreakValue<'a>>) -> ControlFlow<BreakValue<'a>>,
91{
92    FromFn { f, _marker: PhantomData }
93}
94
95#[test]
96fn pipeline_tests() {
97    fn check(iter: impl InternalIterator<Item = i32>, expect: &[i32]) {
98        assert_eq!(iter.collect::<Vec<_>>(), expect);
99    }
100    fn make_iter() -> impl InternalIterator<Item = i32> {
101        from_fn(|f| {
102            f(1)?;
103            f(2)?;
104            f(3)?;
105            f(4)?;
106            ControlFlow::Continue(())
107        })
108    }
109    check(make_iter(), &[1, 2, 3, 4]);
110    check(make_iter().map(|x| x * 2), &[2, 4, 6, 8]);
111    check(make_iter().filter(|&x| x > 2), &[3, 4]);
112    check(make_iter().filter(|&x| x <= 2), &[1, 2]);
113    check(make_iter().chain([7, 8, 9]), &[1, 2, 3, 4, 7, 8, 9]);
114}
115
116#[test]
117fn can_be_static() {
118    fn use_static_iterator(iter: impl InternalIterator<Item = i32> + 'static) {
119        assert_eq!(iter.collect::<Vec<_>>(), &[1, 2, 3]);
120    }
121    let mut a = 1;
122    let iter = from_fn(move |f| {
123        f(a)?;
124        a += 1;
125        f(a)?;
126        a += 1;
127        f(a)?;
128        ControlFlow::Continue(())
129    });
130    use_static_iterator(iter);
131}
132
133/// ```compile_fail
134/// use std::ops::ControlFlow;
135/// use internal_iterator::InternalIterator;
136///
137/// internal_iterator::from_fn(|f| {
138///     f(1)?;
139///     let mut slot = ControlFlow::Continue(());
140///     internal_iterator::from_fn(|f2| {
141///         slot = f2(42); // this assignment is supposed to fail to compile,
142///                        // BreakValue can be moved between from_fn invocations
143///         ControlFlow::Continue(())
144///     }).collect::<Vec<_>>();
145///     slot
146/// });
147/// ```
148fn __compile_fail_check() {}
internal_iterator/from_fn_impl.rs

internal_iterator/
from_fn_impl.rs
