yield_return/
local_iter.rs

1use std::{
2    cell::RefCell,
3    future::Future,
4    iter::FusedIterator,
5    ops::{Deref, DerefMut},
6    pin::Pin,
7    rc::Rc,
8    task::{Context, Poll},
9};
10
11use futures_core::{FusedStream, Stream};
12
13use crate::utils::noop_waker;
14
15struct Sender<T>(Rc<RefCell<Option<T>>>);
16
17impl<T> Sender<T> {
18    #[track_caller]
19    fn set(&self, value: T) {
20        let mut data = self.0.borrow_mut();
21        assert!(data.is_none(), "The result of `ret` is not await.");
22        *data = Some(value);
23    }
24}
25
26impl<T> Future for Sender<T> {
27    type Output = ();
28    fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
29        if self.0.borrow().is_some() {
30            Poll::Pending
31        } else {
32            Poll::Ready(())
33        }
34    }
35}
36
37/// Non-`Send` iterator context.
38///
39/// This type does not implement `Send`.
40pub struct LocalIterContext<T>(Sender<T>);
41
42impl<T> LocalIterContext<T> {
43    /// Yields a single value. Similar to C#'s `yield return` or Python's `yield`.
44    #[track_caller]
45    pub fn ret(&mut self, value: T) -> impl Future<Output = ()> + '_ {
46        self.0.set(value);
47        &mut self.0
48    }
49
50    /// Yields all values from an iterator. Similar to Python's `yield from` or JavaScript's `yield*`.
51    pub async fn ret_iter(&mut self, iter: impl IntoIterator<Item = T>) {
52        for value in iter {
53            self.ret(value).await;
54        }
55    }
56}
57
58struct Data<'a, T> {
59    value: Rc<RefCell<Option<T>>>,
60    fut: Option<Pin<Box<dyn Future<Output = ()> + 'a>>>,
61}
62impl<T> Data<'_, T> {
63    fn poll_next(&mut self, cx: &mut Context) -> Poll<Option<T>> {
64        let Some(fut) = &mut self.fut else {
65            return Poll::Ready(None);
66        };
67        let poll = fut.as_mut().poll(cx);
68        match poll {
69            Poll::Ready(_) => {
70                assert!(
71                    self.value.borrow().is_none(),
72                    "The result of `ret` is not await."
73                );
74                self.fut = None;
75                Poll::Ready(None)
76            }
77            Poll::Pending => {
78                if let Some(value) = self.value.borrow_mut().take() {
79                    Poll::Ready(Some(value))
80                } else {
81                    Poll::Pending
82                }
83            }
84        }
85    }
86}
87
88/// Non-`Send` iterator implemented using async functions.
89///
90/// This type does not implement `Send`.
91pub struct LocalIter<'a, T>(Data<'a, T>);
92
93impl<'a, T: 'a> LocalIter<'a, T> {
94    /// Create an iterator from an asynchronous function.
95    ///
96    /// # Example
97    /// ```
98    /// use yield_return::Yield;
99    /// let iter = Yield::new(|mut y| async move {
100    ///     y.ret(1).await;
101    ///     y.ret(2).await;
102    /// });
103    /// let list: Vec<_> = iter.collect();
104    /// assert_eq!(list, vec![1, 2]);
105    /// ```
106    pub fn new<Fut: Future<Output = ()> + 'a>(f: impl FnOnce(LocalIterContext<T>) -> Fut) -> Self {
107        let value = Rc::new(RefCell::new(None));
108        let cx = LocalIterContext(Sender(value.clone()));
109        let fut: Pin<Box<dyn Future<Output = ()> + 'a>> = Box::pin(f(cx));
110        let fut = Some(fut);
111        Self(Data { value, fut })
112    }
113}
114
115impl<T> Iterator for LocalIter<'_, T> {
116    type Item = T;
117    #[track_caller]
118    fn next(&mut self) -> Option<Self::Item> {
119        match self.0.poll_next(&mut Context::from_waker(&noop_waker())) {
120            Poll::Ready(value) => value,
121            Poll::Pending => panic!("`YieldContext::ret` is not called."),
122        }
123    }
124}
125impl<T> FusedIterator for LocalIter<'_, T> {}
126
127/// Non-`Send` stream context.
128///
129/// This type does not implement `Send`.
130pub struct LocalAsyncIterContext<T>(LocalIterContext<T>);
131impl<T> Deref for LocalAsyncIterContext<T> {
132    type Target = LocalIterContext<T>;
133    fn deref(&self) -> &Self::Target {
134        &self.0
135    }
136}
137impl<T> DerefMut for LocalAsyncIterContext<T> {
138    fn deref_mut(&mut self) -> &mut Self::Target {
139        &mut self.0
140    }
141}
142
143/// Non-`Send` stream implemented using async functions.
144///
145/// This type does not implement `Send`.
146pub struct LocalAsyncIter<'a, T>(LocalIter<'a, T>);
147
148impl<'a, T: 'a> LocalAsyncIter<'a, T> {
149    /// Create a stream from an asynchronous function.
150    ///
151    /// # Example
152    /// ```
153    /// use yield_return::LocalAsyncIter;
154    /// # futures::executor::block_on(async {
155    /// let iter = LocalAsyncIter::new(|mut y| async move {
156    ///     y.ret(1).await;
157    ///     y.ret(2).await;
158    /// });
159    /// let list: Vec<_> = futures::StreamExt::collect(iter).await;
160    /// assert_eq!(list, vec![1, 2]);
161    /// # });
162    /// ```
163    pub fn new<Fut: Future<Output = ()> + 'a>(
164        f: impl FnOnce(LocalAsyncIterContext<T>) -> Fut,
165    ) -> Self {
166        Self(LocalIter::new(|cx| f(LocalAsyncIterContext(cx))))
167    }
168}
169
170impl<T> Stream for LocalAsyncIter<'_, T> {
171    type Item = T;
172
173    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
174        self.0 .0.poll_next(cx)
175    }
176}
177impl<T> FusedStream for LocalAsyncIter<'_, T> {
178    fn is_terminated(&self) -> bool {
179        self.0 .0.fut.is_none()
180    }
181}
yield_return/local_iter.rs

yield_return/
local_iter.rs
