1use core::pin::Pin;
2
3use crate::stream::Stream;
4use crate::task::{Context, Poll};
5
6/// A stream that yields elements by calling a closure.
7///
8/// This stream is created by the [`from_fn`] function. See its
9/// documentation for more.
10///
11/// [`from_fn`]: fn.from_fn.html
12#[derive(Clone, Debug)]
13pub struct FromFn<F> {
14 f: F,
15}
16
17impl<F> Unpin for FromFn<F> {}
18
19/// Creates a new stream where to produce each new element a provided closure is called.
20///
21/// This allows creating a custom stream with any behaviour without using the more verbose
22/// syntax of creating a dedicated type and implementing a `Stream` trait for it.
23///
24/// # Examples
25///
26/// ```
27/// # async_std::task::block_on(async {
28/// #
29/// use async_std::prelude::*;
30/// use async_std::stream;
31///
32/// let mut count = 0u8;
33/// let s = stream::from_fn(|| {
34/// count += 1;
35/// if count > 3 {
36/// None
37/// } else {
38/// Some(count)
39/// }
40/// });
41///
42/// pin_utils::pin_mut!(s);
43///
44/// assert_eq!(s.next().await, Some(1));
45/// assert_eq!(s.next().await, Some(2));
46/// assert_eq!(s.next().await, Some(3));
47/// assert_eq!(s.next().await, None);
48/// #
49/// # })
50/// ```
51pub fn from_fn<T, F>(f: F) -> FromFn<F>
52where
53 F: FnMut() -> Option<T>,
54{
55 FromFn { f }
56}
57
58impl<T, F> Stream for FromFn<F>
59where
60 F: FnMut() -> Option<T>,
61{
62 type Item = T;
63
64 fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
65 let item: Option = (&mut self.f)();
66 Poll::Ready(item)
67 }
68}
69