├── .github └── workflows │ └── test.yml ├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── async-gen-macros ├── Cargo.toml └── src │ └── lib.rs ├── benches ├── simple.rs └── utils.rs ├── examples └── tcp_accept.rs ├── src └── lib.rs └── tests ├── spans_preserved.rs ├── stream.rs └── try_stream.rs /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: 4 | push: 5 | branches: [ "master" ] 6 | pull_request: 7 | branches: [ "master" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | test: 14 | runs-on: ubuntu-latest 15 | 16 | steps: 17 | - uses: actions/checkout@v3 18 | - name: Install Miri 19 | run: | 20 | rustup toolchain install nightly --component miri 21 | cargo +nightly miri setup 22 | 23 | - name: Run clippy 24 | run: cargo clippy 25 | 26 | - name: Run tests 27 | run: cargo test 28 | 29 | - name: Run miri test 30 | run: cargo +nightly miri test --test stream --test try_stream 31 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | /Cargo.lock 3 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = ["async-gen-macros"] 3 | 4 | [package] 5 | name = "async-gen" 6 | version = "0.2.3" 7 | edition = "2021" 8 | 9 | license = "MIT" 10 | keywords = ["async", "generator"] 11 | authors = ["Nur "] 12 | repository = "https://github.com/nurmohammed840/async-gen" 13 | description = "Async generator in stable rust using async/await" 14 | 15 | exclude = ["/async-gen-macros", "/tests", ".*"] 16 | 17 | [[bench]] 18 | name = "simple" 19 | harness = false 20 | 21 | [dependencies] 22 | futures-core = "0.3" 23 | async-gen-macros = { version = "0.3", path = "./async-gen-macros" } 24 | pin-project-lite = "0.2" 25 | 26 | [dev-dependencies] 27 | futures-util = "0.3" 28 | tokio = { version = "1", features = ["full"] } 29 | 30 | async-stream = "*" 31 | criterion = "*" 32 | futures = { version = "*", features = ["default"] } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Nur 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | This library provides a way to create asynchronous generator using the `async/await` feature in stable Rust. 2 | 3 | It is similar to [async-stream](https://docs.rs/async-stream/latest/async_stream/), 4 | But closely mimics the [Coroutine API](https://doc.rust-lang.org/std/ops/trait.Coroutine.html), 5 | Allowing the generator also return a value upon completion, in addition to yielding intermediate values. 6 | 7 | # Installation 8 | 9 | Add it as a dependency to your Rust project by adding the following line to your `Cargo.toml` file: 10 | 11 | 12 | ```toml 13 | [dependencies] 14 | async-gen = "0.2" 15 | ``` 16 | 17 | # Examples 18 | 19 | ```rust 20 | use std::pin::pin; 21 | use async_gen::{gen, GeneratorState}; 22 | 23 | #[tokio::main] 24 | async fn main() { 25 | let g = gen! { 26 | yield 42; 27 | return "42" 28 | }; 29 | let mut g = pin!(g); 30 | assert_eq!(g.resume().await, GeneratorState::Yielded(42)); 31 | assert_eq!(g.resume().await, GeneratorState::Complete("42")); 32 | } 33 | ``` -------------------------------------------------------------------------------- /async-gen-macros/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "async-gen-macros" 3 | version = "0.3.0" 4 | edition = "2021" 5 | 6 | license = "MIT" 7 | keywords = ["async", "generator"] 8 | authors = ["Nur "] 9 | repository = "https://github.com/nurmohammed840/async-gen" 10 | description = "Async generator in stable rust using async/await" 11 | 12 | [lib] 13 | proc-macro = true 14 | 15 | [dependencies] 16 | -------------------------------------------------------------------------------- /async-gen-macros/src/lib.rs: -------------------------------------------------------------------------------- 1 | use proc_macro::*; 2 | 3 | #[proc_macro] 4 | pub fn gen_inner(input: TokenStream) -> TokenStream { 5 | let mut tokens = input.into_iter(); 6 | 7 | let Some(TokenTree::Group(crate_path)) = tokens.next() else { 8 | unimplemented!() 9 | }; 10 | let crate_path = crate_path.stream(); 11 | 12 | let mut has_yielded = false; 13 | let output = out(tokens, &mut has_yielded); 14 | 15 | let mut o = TokenStream::new(); 16 | o.extend(crate_path.clone()); 17 | o.push_colon2(); 18 | o.push_ident("gen"); 19 | 20 | o.push_group(Delimiter::Parenthesis, |o| { 21 | o.push_punct('|'); 22 | o.push_ident("mut"); 23 | o.push_ident("yield_"); 24 | 25 | if !has_yielded { 26 | o.push_punct(':'); 27 | o.extend(crate_path); 28 | o.push_colon2(); 29 | o.push_ident("Yield"); 30 | } 31 | 32 | o.push_punct('|'); 33 | o.push_ident("async"); 34 | o.push_ident("move"); 35 | o.push_group(Delimiter::Brace, |o| { 36 | o.push_ident("let"); 37 | o.push_ident("v"); 38 | o.push_punct('='); 39 | o.push_ident("async"); 40 | 41 | o.push(output); 42 | 43 | o.push_punct('.'); 44 | o.push_ident("await"); 45 | o.push_punct(';'); 46 | o.push_ident("yield_"); 47 | o.push_punct('.'); 48 | o.push_ident("return_"); 49 | o.push_group(Delimiter::Parenthesis, |o| { 50 | o.push_ident("v"); 51 | }); 52 | }); 53 | }); 54 | o 55 | } 56 | 57 | fn out(mut tokens: token_stream::IntoIter, has_yielded: &mut bool) -> Group { 58 | let mut o = TokenStream::new(); 59 | 60 | while let Some(tt) = tokens.next() { 61 | match tt { 62 | TokenTree::Ident(name) if name.to_string() == "yield" => { 63 | *has_yielded = true; 64 | let mut expr = TokenStream::new(); 65 | for tt in &mut tokens { 66 | match tt { 67 | TokenTree::Punct(p) if p.as_char() == ';' => break, 68 | _ => expr.push(tt), 69 | } 70 | } 71 | if expr.is_empty() { 72 | expr.push(Group::new(Delimiter::Parenthesis, TokenStream::new())); 73 | }; 74 | o.push_ident("yield_"); 75 | o.push_punct('.'); 76 | o.push_ident("yield_"); 77 | o.push(Group::new(Delimiter::Parenthesis, expr)); 78 | o.push_punct('.'); 79 | o.push_ident("await"); 80 | o.push_punct(';'); 81 | } 82 | TokenTree::Group(g) if g.delimiter() == Delimiter::Brace => { 83 | o.push(out(g.stream().into_iter(), has_yielded)); 84 | } 85 | _ => o.push(tt), 86 | } 87 | } 88 | Group::new(Delimiter::Brace, o) 89 | } 90 | 91 | trait TokenStreamExt { 92 | fn push(&mut self, token: U) 93 | where 94 | U: Into; 95 | 96 | #[inline] 97 | fn push_punct(&mut self, ch: char) { 98 | self.push(Punct::new(ch, Spacing::Alone)) 99 | } 100 | 101 | #[inline] 102 | fn push_ident(&mut self, name: &str) { 103 | self.push(Ident::new(name, Span::call_site())) 104 | } 105 | 106 | #[inline] 107 | fn push_group(&mut self, delimiter: Delimiter, f: impl FnOnce(&mut TokenStream)) { 108 | let mut stream = TokenStream::new(); 109 | f(&mut stream); 110 | self.push(Group::new(delimiter, stream)) 111 | } 112 | 113 | #[inline] 114 | fn push_colon2(&mut self) { 115 | self.push(Punct::new(':', Spacing::Joint)); 116 | self.push(Punct::new(':', Spacing::Alone)); 117 | } 118 | } 119 | 120 | impl TokenStreamExt for TokenStream { 121 | #[inline] 122 | fn push(&mut self, token: U) 123 | where 124 | U: Into, 125 | { 126 | self.extend(std::iter::once(token.into())); 127 | } 128 | } 129 | -------------------------------------------------------------------------------- /benches/simple.rs: -------------------------------------------------------------------------------- 1 | mod utils; 2 | 3 | use criterion::{criterion_group, criterion_main, Criterion}; 4 | use futures_util::StreamExt; 5 | use std::pin::pin; 6 | 7 | const ITER: usize = 1000000; 8 | 9 | async fn async_gen_sum(iter: usize) -> usize { 10 | let mut gen = pin!(async_gen::gen! { 11 | for i in 1..=iter { 12 | yield i; 13 | } 14 | }); 15 | let mut i = 0; 16 | while let Some(v) = gen.next().await { 17 | i += v; 18 | } 19 | i 20 | } 21 | 22 | async fn async_stream_sum(iter: usize) -> usize { 23 | let mut gen = pin!(async_stream::stream! { 24 | for i in 1..=iter { 25 | yield i; 26 | } 27 | }); 28 | let mut i = 0; 29 | while let Some(v) = gen.next().await { 30 | i += v; 31 | } 32 | i 33 | } 34 | 35 | pub fn simple_benchmark(c: &mut Criterion) { 36 | c.bench_function("async_gen", |b| { 37 | b.iter(|| utils::poll_once(async_gen_sum(ITER))) 38 | }); 39 | c.bench_function("async_stream", |b| { 40 | b.iter(|| utils::poll_once(async_stream_sum(ITER))) 41 | }); 42 | } 43 | criterion_group!(benches, simple_benchmark); 44 | criterion_main!(benches); 45 | -------------------------------------------------------------------------------- /benches/utils.rs: -------------------------------------------------------------------------------- 1 | use std::pin::pin; 2 | use std::{future::Future, task::*}; 3 | 4 | static DATA: () = (); 5 | static VTABLE: RawWakerVTable = RawWakerVTable::new(|_| raw_waker(), no_op, no_op, no_op); 6 | fn no_op(_: *const ()) {} 7 | fn raw_waker() -> RawWaker { 8 | RawWaker::new(&DATA, &VTABLE) 9 | } 10 | 11 | #[inline] 12 | pub fn poll_once(fut: impl Future) -> T { 13 | let waker: Waker = unsafe { Waker::from_raw(raw_waker()) }; 14 | let mut cx = Context::from_waker(&waker); 15 | let fut = pin!(fut); 16 | match fut.poll(&mut cx) { 17 | Poll::Ready(val) => val, 18 | Poll::Pending => panic!("pending"), 19 | } 20 | } 21 | -------------------------------------------------------------------------------- /examples/tcp_accept.rs: -------------------------------------------------------------------------------- 1 | use async_gen::gen; 2 | use futures_util::StreamExt; 3 | use std::pin::pin; 4 | use tokio::net::TcpListener; 5 | 6 | #[tokio::main] 7 | async fn main() { 8 | let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); 9 | 10 | let mut incoming = pin!(gen! { 11 | loop { 12 | let (socket, _) = listener.accept().await.unwrap(); 13 | yield socket; 14 | } 15 | }); 16 | 17 | while let Some(v) = incoming.next().await { 18 | println!("handle = {:?}", v); 19 | } 20 | } 21 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #![doc = include_str!("../README.md")] 2 | #![warn(missing_docs)] 3 | 4 | pub use futures_core; 5 | use pin_project_lite::pin_project; 6 | use std::{ 7 | cell::UnsafeCell, 8 | future::{poll_fn, Future}, 9 | pin::Pin, 10 | sync::Arc, 11 | task::{Context, Poll}, 12 | }; 13 | 14 | /// The result of a generator resumption. 15 | /// 16 | /// This enum is returned from the `Generator::resume` method and indicates the 17 | /// possible return values of a generator. Currently this corresponds to either 18 | /// a suspension point (`Yielded`) or a termination point (`Complete`). 19 | #[derive(Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Debug, Hash)] 20 | pub enum GeneratorState { 21 | /// The generator suspended with a value. 22 | /// 23 | /// This state indicates that a generator has been suspended, and typically 24 | /// corresponds to a `yield` statement. The value provided in this variant 25 | /// corresponds to the expression passed to `yield` and allows generators to 26 | /// provide a value each time they yield. 27 | Yielded(Y), 28 | 29 | /// The generator completed with a return value. 30 | /// 31 | /// This state indicates that a generator has finished execution with the 32 | /// provided value. Once a generator has returned `Complete` it is 33 | /// considered a programmer error to call `resume` again. 34 | Complete(R), 35 | } 36 | 37 | /// Generators, also commonly referred to as coroutines. 38 | pub trait AsyncGenerator { 39 | /// The type of value this generator yields. 40 | /// 41 | /// This associated type corresponds to the `yield` expression and the 42 | /// values which are allowed to be returned each time a generator yields. 43 | /// For example an iterator-as-a-generator would likely have this type as 44 | /// `T`, the type being iterated over. 45 | type Yield; 46 | 47 | /// The type of value this generator returns. 48 | /// 49 | /// This corresponds to the type returned from a generator either with a 50 | /// `return` statement or implicitly as the last expression of a generator 51 | /// literal. For example futures would use this as `Result` as it 52 | /// represents a completed future. 53 | type Return; 54 | 55 | /// Resumes the execution of this generator. 56 | /// 57 | /// This function will resume execution of the generator or start execution 58 | /// if it hasn't already. This call will return back into the generator's 59 | /// last suspension point, resuming execution from the latest `yield`. The 60 | /// generator will continue executing until it either yields or returns, at 61 | /// which point this function will return. 62 | /// 63 | /// # Return value 64 | /// 65 | /// The `GeneratorState` enum returned from this function indicates what 66 | /// state the generator is in upon returning. If the `Yielded` variant is 67 | /// returned then the generator has reached a suspension point and a value 68 | /// has been yielded out. Generators in this state are available for 69 | /// resumption at a later point. 70 | /// 71 | /// If `Complete` is returned then the generator has completely finished 72 | /// with the value provided. It is invalid for the generator to be resumed 73 | /// again. 74 | /// 75 | /// # Panics 76 | /// 77 | /// This function may panic if it is called after the `Complete` variant has 78 | /// been returned previously. While generator literals in the language are 79 | /// guaranteed to panic on resuming after `Complete`, this is not guaranteed 80 | /// for all implementations of the `Generator` trait. 81 | fn poll_resume( 82 | self: Pin<&mut Self>, 83 | cx: &mut Context<'_>, 84 | ) -> Poll>; 85 | } 86 | 87 | struct Inner { 88 | data: UnsafeCell>, 89 | } 90 | 91 | unsafe impl Send for Inner {} 92 | unsafe impl Sync for Inner {} 93 | 94 | #[doc(hidden)] 95 | pub struct Yield { 96 | inner: Arc>, 97 | } 98 | 99 | #[doc(hidden)] 100 | pub struct Return(T); 101 | 102 | impl Yield { 103 | /// Same as `yield` keyword. 104 | /// 105 | /// It pauses execution and the value is returned to the generator's caller. 106 | pub async fn yield_(&mut self, val: Y) { 107 | // SEAFTY: this function is marked with `&mut self` 108 | // 109 | // And `Yield<()>` can't escape from this closure: 110 | // 111 | // gen(|y: Yield<()>| async { 112 | // // `y` can't escape from this closure. and owned by `async` body 113 | // y.return_(()) 114 | // }); 115 | unsafe { 116 | *self.inner.data.get() = Some(val); 117 | } 118 | 119 | poll_fn(|_| { 120 | if unsafe { (*self.inner.data.get()).is_some() } { 121 | return Poll::Pending; 122 | } 123 | Poll::Ready(()) 124 | }) 125 | .await 126 | } 127 | 128 | #[inline] 129 | pub fn return_(self, _v: R) -> Return { 130 | Return(_v) 131 | } 132 | } 133 | 134 | pin_project! { 135 | /// Represent an asyncronus generator. It implementations [`AsyncGenerator`] trait. 136 | /// 137 | /// This `struct` is created by [`gen()`]. See its documentation for more details. 138 | pub struct AsyncGen { 139 | inner: Arc>, 140 | #[pin] 141 | fut: Fut, 142 | } 143 | } 144 | 145 | impl AsyncGen 146 | where 147 | Fut: Future>, 148 | { 149 | /// See [`AsyncGenerator::poll_resume`] for more details. 150 | #[doc(hidden)] 151 | pub fn poll_resume(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 152 | let me = self.project(); 153 | match me.fut.poll(cx) { 154 | Poll::Ready(Return(val)) => Poll::Ready(GeneratorState::Complete(val)), 155 | Poll::Pending => { 156 | // SEAFTY: We just return from `me.fut`, 157 | // So this is safe and unique access to `me.inner.data` 158 | unsafe { 159 | if (*me.inner.data.get()).is_some() { 160 | return Poll::Ready(GeneratorState::Yielded( 161 | (*me.inner.data.get()).take().unwrap_unchecked(), 162 | )); 163 | } 164 | } 165 | Poll::Pending 166 | } 167 | } 168 | } 169 | 170 | #[inline] 171 | /// See [`AsyncGenerator::poll_resume`] for more details. 172 | pub async fn resume(self: &mut Pin<&mut Self>) -> GeneratorState { 173 | poll_fn(|cx| self.as_mut().poll_resume(cx)).await 174 | } 175 | } 176 | 177 | impl AsyncGen 178 | where 179 | Fut: Future>, 180 | { 181 | #[inline] 182 | /// Creates an async iterator from this generator. 183 | /// 184 | /// See [`AsyncIter`] for more details. 185 | pub fn into_async_iter(self) -> AsyncIter { 186 | AsyncIter::from(self) 187 | } 188 | 189 | #[doc(hidden)] 190 | pub fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 191 | let me = self.project(); 192 | match me.fut.poll(cx) { 193 | Poll::Ready(Return(())) => Poll::Ready(None), 194 | Poll::Pending => { 195 | // SEAFTY: We just return from `me.fut`, 196 | // So this is safe and unique access to `me.inner.data` 197 | unsafe { 198 | if (*me.inner.data.get()).is_some() { 199 | return Poll::Ready((*me.inner.data.get()).take()); 200 | } 201 | } 202 | Poll::Pending 203 | } 204 | } 205 | } 206 | } 207 | 208 | impl futures_core::Stream for AsyncGen 209 | where 210 | Fut: Future>, 211 | { 212 | type Item = Y; 213 | 214 | #[inline] 215 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 216 | AsyncGen::poll_next(self, cx) 217 | } 218 | } 219 | 220 | impl AsyncGenerator for AsyncGen 221 | where 222 | Fut: Future>, 223 | { 224 | type Yield = Y; 225 | type Return = R; 226 | 227 | #[inline] 228 | fn poll_resume( 229 | self: Pin<&mut Self>, 230 | cx: &mut Context<'_>, 231 | ) -> Poll> { 232 | AsyncGen::poll_resume(self, cx) 233 | } 234 | } 235 | 236 | pin_project! { 237 | /// An async iterator over the values yielded by an underlying generator. 238 | /// 239 | /// ## Example 240 | /// 241 | /// ``` 242 | /// use async_gen::{gen, AsyncIter}; 243 | /// use futures_core::Stream; 244 | /// use futures_util::StreamExt; 245 | /// 246 | /// fn get_async_iter() -> impl Stream { 247 | /// AsyncIter::from(gen! { 248 | /// yield 1; 249 | /// yield 2; 250 | /// yield 3; 251 | /// }) 252 | /// } 253 | /// 254 | /// #[tokio::main] 255 | /// async fn main() { 256 | /// let it = get_async_iter(); 257 | /// let v: Vec<_> = it.collect().await; 258 | /// assert_eq!(v, [1, 2, 3]); 259 | /// } 260 | /// ``` 261 | #[derive(Clone)] 262 | pub struct AsyncIter { 263 | #[pin] 264 | gen: G, 265 | } 266 | } 267 | 268 | impl From for AsyncIter { 269 | #[inline] 270 | fn from(gen: G) -> Self { 271 | AsyncIter { gen } 272 | } 273 | } 274 | 275 | impl> AsyncIter { 276 | /// Attempt to pull out the next value of this async iterator, registering the 277 | /// current task for wakeup if the value is not yet available, and returning 278 | /// `None` if the async iterator is exhausted. 279 | /// 280 | /// # Return value 281 | /// 282 | /// There are several possible return values, each indicating a distinct 283 | /// async iterator state: 284 | /// 285 | /// - `Poll::Pending` means that this async iterator's next value is not ready 286 | /// yet. Implementations will ensure that the current task will be notified 287 | /// when the next value may be ready. 288 | /// 289 | /// - `Poll::Ready(Some(val))` means that the async iterator has successfully 290 | /// produced a value, `val`, and may produce further values on subsequent 291 | /// `poll_next` calls. 292 | /// 293 | /// - `Poll::Ready(None)` means that the async iterator has terminated, and 294 | /// `poll_next` should not be invoked again. 295 | /// 296 | /// # Panics 297 | /// 298 | /// Once an async iterator has finished (returned `Ready(None)` from `poll_next`), calling its 299 | /// `poll_next` method again may panic, block forever, or cause other kinds of 300 | /// problems; the `AsyncIterator` trait places no requirements on the effects of 301 | /// such a call. However, as the `poll_next` method is not marked `unsafe`, 302 | /// Rust's usual rules apply: calls must never cause undefined behavior 303 | /// (memory corruption, incorrect use of `unsafe` functions, or the like), 304 | /// regardless of the async iterator's state. 305 | #[inline] 306 | pub fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 307 | self.project().gen.poll_resume(cx).map(|s| match s { 308 | GeneratorState::Yielded(val) => Some(val), 309 | GeneratorState::Complete(()) => None, 310 | }) 311 | } 312 | } 313 | 314 | impl> futures_core::Stream for AsyncIter { 315 | type Item = G::Yield; 316 | #[inline] 317 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 318 | AsyncIter::poll_next(self, cx) 319 | } 320 | } 321 | 322 | /// Creates a new generator, which implements the [`AsyncGenerator`] trait. 323 | /// 324 | /// Also see [`gen!`] macro for more details. 325 | /// 326 | /// ## Examples 327 | /// 328 | /// ``` 329 | /// use async_gen::{gen, AsyncGen, AsyncGenerator, Return}; 330 | /// use std::future::Future; 331 | /// 332 | /// fn example() { 333 | /// let g = gen(|mut c| async { 334 | /// c.yield_(42).await; 335 | /// c.return_("42") 336 | /// }); 337 | /// 338 | /// check_type_1(&g); 339 | /// check_type_2(&g); 340 | /// } 341 | /// fn check_type_1(_: &AsyncGen>, i32>) {} 342 | /// fn check_type_2(_: &impl AsyncGenerator) {} 343 | /// ``` 344 | pub fn gen(fut: impl FnOnce(Yield) -> Fut) -> AsyncGen 345 | where 346 | Fut: Future>, 347 | { 348 | let inner = Arc::new(Inner { 349 | data: UnsafeCell::new(None), 350 | }); 351 | let fut = fut(Yield { 352 | inner: inner.clone(), 353 | }); 354 | AsyncGen { inner, fut } 355 | } 356 | 357 | /// A macro for creating generator. 358 | /// 359 | /// Also see [`gen()`] function for more details. 360 | /// 361 | /// ## Examples 362 | /// 363 | /// ``` 364 | /// use std::pin::pin; 365 | /// use async_gen::{gen, GeneratorState}; 366 | /// 367 | /// # #[tokio::main] 368 | /// # async fn main() { 369 | /// let gen = gen! { 370 | /// yield 42; 371 | /// return "42" 372 | /// }; 373 | /// let mut g = pin!(gen); 374 | /// assert_eq!(g.resume().await, GeneratorState::Yielded(42)); 375 | /// assert_eq!(g.resume().await, GeneratorState::Complete("42")); 376 | /// # } 377 | /// ``` 378 | #[macro_export] 379 | macro_rules! gen { 380 | ($($tt:tt)*) => { 381 | $crate::__private::gen_inner!(($crate) $($tt)*) 382 | } 383 | } 384 | 385 | #[doc(hidden)] 386 | pub mod __private { 387 | pub use async_gen_macros::*; 388 | } 389 | -------------------------------------------------------------------------------- /tests/spans_preserved.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::never_loop)] 2 | 3 | use async_gen::gen; 4 | use futures_util::stream::StreamExt; 5 | use std::pin::pin; 6 | 7 | #[tokio::test] 8 | async fn spans_preserved() { 9 | let mut s = pin!(gen! { 10 | assert_eq!(line!(), 10); 11 | }); 12 | 13 | while s.next().await.is_some() { 14 | unreachable!(); 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /tests/stream.rs: -------------------------------------------------------------------------------- 1 | use async_gen::{gen, GeneratorState}; 2 | use futures::executor::block_on; 3 | use futures_core::Stream; 4 | use futures_util::stream::StreamExt; 5 | use std::pin::pin; 6 | 7 | #[test] 8 | fn noop_stream() { 9 | block_on(async { 10 | let mut gen = pin!(gen! {}); 11 | assert_eq!(gen.resume().await, GeneratorState::Complete(())); 12 | }) 13 | } 14 | 15 | #[test] 16 | fn empty_stream() { 17 | block_on(async { 18 | let mut ran = false; 19 | { 20 | let r = &mut ran; 21 | let mut gen = pin!(gen! { 22 | *r = true; 23 | println!("hello world!"); 24 | }); 25 | assert_eq!(gen.resume().await, GeneratorState::Complete(())); 26 | } 27 | assert!(ran); 28 | }); 29 | } 30 | 31 | #[test] 32 | fn yield_single_value() { 33 | block_on(async { 34 | let mut s = pin!(gen! { 35 | yield "hello"; 36 | }); 37 | assert_eq!(s.resume().await, GeneratorState::Yielded("hello")); 38 | assert_eq!(s.resume().await, GeneratorState::Complete(())); 39 | }) 40 | } 41 | 42 | #[test] 43 | fn fused() { 44 | block_on(async { 45 | let s = pin!(gen! { 46 | yield "hello"; 47 | }); 48 | let mut s = s.fuse(); 49 | assert_eq!(s.next().await, Some("hello")); 50 | assert_eq!(s.next().await, None); 51 | assert_eq!(s.next().await, None); 52 | }); 53 | } 54 | 55 | #[test] 56 | fn yield_multi_value() { 57 | block_on(async { 58 | let mut s = pin!(gen! { 59 | yield "hello"; 60 | yield "world"; 61 | yield "dizzy"; 62 | }); 63 | assert_eq!(s.resume().await, GeneratorState::Yielded("hello")); 64 | assert_eq!(s.resume().await, GeneratorState::Yielded("world")); 65 | assert_eq!(s.resume().await, GeneratorState::Yielded("dizzy")); 66 | assert_eq!(s.resume().await, GeneratorState::Complete(())); 67 | }) 68 | } 69 | 70 | #[test] 71 | fn return_stream() { 72 | block_on(async { 73 | fn build_stream() -> impl Stream { 74 | gen! { 75 | yield 1; 76 | yield 2; 77 | yield 3; 78 | } 79 | } 80 | let s = build_stream(); 81 | 82 | let values: Vec<_> = s.collect().await; 83 | assert_eq!(3, values.len()); 84 | assert_eq!(1, values[0]); 85 | assert_eq!(2, values[1]); 86 | assert_eq!(3, values[2]); 87 | }) 88 | } 89 | 90 | #[test] 91 | fn consume_channel() { 92 | block_on(async { 93 | let (tx, mut rx) = tokio::sync::mpsc::channel(10); 94 | let mut s = pin!(gen! { 95 | while let Some(v) = rx.recv().await { 96 | yield v; 97 | } 98 | }); 99 | for i in 0..3 { 100 | assert!(tx.send(i).await.is_ok()); 101 | assert_eq!(Some(i), s.next().await); 102 | } 103 | drop(tx); 104 | assert_eq!(None, s.next().await); 105 | }); 106 | } 107 | 108 | #[test] 109 | fn borrow_self() { 110 | block_on(async { 111 | struct Data(String); 112 | 113 | impl Data { 114 | fn stream(&self) -> impl Stream + '_ { 115 | gen! { 116 | yield &self.0[..]; 117 | } 118 | } 119 | } 120 | 121 | let data = Data("hello".to_string()); 122 | let mut s = pin!(data.stream()); 123 | assert_eq!(Some("hello"), s.next().await); 124 | }) 125 | } 126 | 127 | #[test] 128 | fn stream_in_stream() { 129 | block_on(async { 130 | let s = gen! { 131 | let mut s = pin!(gen! { 132 | for i in 0..3 { 133 | yield i; 134 | } 135 | }); 136 | while let Some(v) = s.next().await { 137 | yield v; 138 | } 139 | }; 140 | let values: Vec<_> = s.collect().await; 141 | assert_eq!(3, values.len()); 142 | }) 143 | } 144 | 145 | #[test] 146 | fn yield_non_unpin_value() { 147 | block_on(async { 148 | let s: Vec<_> = gen! { 149 | for i in 0..3 { 150 | yield async move { i }; 151 | } 152 | } 153 | .buffered(1) 154 | .collect() 155 | .await; 156 | 157 | assert_eq!(s, vec![0, 1, 2]); 158 | }) 159 | } 160 | 161 | #[test] 162 | fn unit_yield_in_select() { 163 | block_on(async { 164 | async fn do_stuff_async() {} 165 | 166 | let s = gen! { 167 | tokio::select! { 168 | _ = do_stuff_async() => { yield }, 169 | else => { yield }, 170 | }; 171 | }; 172 | let values: Vec<_> = s.collect().await; 173 | assert_eq!(values.len(), 1); 174 | }) 175 | } 176 | 177 | #[test] 178 | fn yield_with_select() { 179 | block_on(async { 180 | async fn do_stuff_async() {} 181 | async fn more_async_work() {} 182 | 183 | let s = gen! { 184 | tokio::select! { 185 | _ = do_stuff_async() => { yield "hey" }, 186 | _ = more_async_work() => { yield "hey" }, 187 | else => { yield "hey" }, 188 | }; 189 | }; 190 | let values: Vec<_> = s.collect().await; 191 | assert_eq!(values, vec!["hey"]); 192 | }) 193 | } 194 | -------------------------------------------------------------------------------- /tests/try_stream.rs: -------------------------------------------------------------------------------- 1 | use async_gen::{gen, GeneratorState}; 2 | use futures::executor::block_on; 3 | use std::pin::pin; 4 | 5 | #[test] 6 | fn single_err() { 7 | block_on(async { 8 | let mut s = pin!(gen! { 9 | if true { 10 | Err("hello")?; 11 | } else { 12 | yield "world"; 13 | } 14 | Result::<_, &str>::Ok(()) 15 | }); 16 | assert_eq!(s.resume().await, GeneratorState::Complete(Err("hello"))); 17 | }) 18 | } 19 | 20 | #[test] 21 | fn yield_then_err() { 22 | block_on(async { 23 | let mut s = pin!(gen! { 24 | yield "hello"; 25 | Err("world")?; 26 | Ok(()) 27 | }); 28 | assert_eq!(s.resume().await, GeneratorState::Yielded("hello")); 29 | assert_eq!(s.resume().await, GeneratorState::Complete(Err("world"))); 30 | }) 31 | } 32 | --------------------------------------------------------------------------------