diff --git a/futures-macro/src/lib.rs b/futures-macro/src/lib.rs index 19643021b8..6357266314 100644 --- a/futures-macro/src/lib.rs +++ b/futures-macro/src/lib.rs @@ -60,5 +60,5 @@ pub fn test_internal(input: TokenStream, item: TokenStream) -> TokenStream { #[cfg_attr(fn_like_proc_macro, proc_macro)] #[cfg_attr(not(fn_like_proc_macro), proc_macro_hack::proc_macro_hack)] pub fn stream_select_internal(input: TokenStream) -> TokenStream { - crate::stream_select::stream_select(input) + crate::stream_select::stream_select(input.into()).unwrap_or_else(syn::Error::into_compile_error).into() } diff --git a/futures-macro/src/stream_select.rs b/futures-macro/src/stream_select.rs index 28dae740d6..936b15eb97 100644 --- a/futures-macro/src/stream_select.rs +++ b/futures-macro/src/stream_select.rs @@ -1,19 +1,18 @@ -use proc_macro::TokenStream; +use proc_macro2::TokenStream; use quote::{format_ident, quote, ToTokens}; use syn::{parse::Parser, punctuated::Punctuated, Expr, Index, Token}; /// The `stream_select!` macro. -pub(crate) fn stream_select(input: TokenStream) -> TokenStream { +pub(crate) fn stream_select(input: TokenStream) -> Result { let args = Punctuated::::parse_terminated - .parse(input) - .expect("macro expects a comma separated list of expressions"); + .parse2(input)?; let generic_idents = (0..args.len()).map(|i| format_ident!("_{}", i)).collect::>(); let field_idents = (0..args.len()).map(|i| format_ident!("__{}", i)).collect::>(); let field_idents_2 = (0..args.len()).map(|i| format_ident!("___{}", i)).collect::>(); let field_indices = (0..args.len()).map(Index::from).collect::>(); let args = args.iter().map(|e| e.to_token_stream()); - TokenStream::from(quote! { + Ok(TokenStream::from(quote! { { #[derive(Debug)] struct StreamSelect<#(#generic_idents),*> (#(Option<#generic_idents>),*); @@ -106,5 +105,5 @@ pub(crate) fn stream_select(input: TokenStream) -> TokenStream { StreamSelect(#(Some(#args)),*) } - }) + })) } diff --git a/futures/tests/async_await_macros.rs b/futures/tests/async_await_macros.rs index 19833d01ca..397de24ade 100644 --- a/futures/tests/async_await_macros.rs +++ b/futures/tests/async_await_macros.rs @@ -4,7 +4,7 @@ use futures::future::{self, poll_fn, FutureExt}; use futures::sink::SinkExt; use futures::stream::StreamExt; use futures::task::{Context, Poll}; -use futures::{join, pending, pin_mut, poll, select, select_biased, try_join}; +use futures::{join, pending, pin_mut, poll, select, select_biased, stream, stream_select, try_join}; use std::mem; #[test] @@ -308,6 +308,43 @@ fn select_on_mutable_borrowing_future_with_same_borrow_in_block_and_default() { }); } +#[test] +#[allow(unused_assignments)] +fn stream_select() { + // stream_select! macro + block_on(async { + let endless_ints = |i| stream::iter(vec![i].into_iter().cycle()); + + let mut endless_ones = stream_select!(endless_ints(1i32), stream::pending()); + assert_eq!(endless_ones.next().await, Some(1)); + assert_eq!(endless_ones.next().await, Some(1)); + + let mut finite_list = stream_select!(stream::iter(vec![1, 2, 3].into_iter())); + assert_eq!(finite_list.next().await, Some(1)); + assert_eq!(finite_list.next().await, Some(2)); + assert_eq!(finite_list.next().await, Some(3)); + assert_eq!(finite_list.next().await, None); + + let endless_mixed = + stream_select!(endless_ints(1i32), endless_ints(2), endless_ints(3)); + // Take 100, and assert a somewhat even distribution of values. + // The fairness is randomized, but over 100 samples we should be pretty close to even. + // This test may be a bit flaky. Feel free to adjust the margins as you see fit. + let mut count = 0; + let results = endless_mixed + .take_while(move |_| { + count += 1; + let ret = count < 100; + async move { ret } + }) + .collect::>() + .await; + assert!(results.iter().filter(|x| **x == 1).count() >= 29); + assert!(results.iter().filter(|x| **x == 2).count() >= 29); + assert!(results.iter().filter(|x| **x == 3).count() >= 29); + }); +} + #[test] fn join_size() { let fut = async { diff --git a/futures/tests/macro-tests/src/main.rs b/futures/tests/macro-tests/src/main.rs index dcfca9bfd5..fef24a99e3 100644 --- a/futures/tests/macro-tests/src/main.rs +++ b/futures/tests/macro-tests/src/main.rs @@ -1,7 +1,7 @@ // Check that it works even if proc-macros are reexported. fn main() { - use futures04::{executor::block_on, future, stream, StreamExt}; + use futures04::{executor::block_on, future}; // join! macro let _ = block_on(async { @@ -65,37 +65,4 @@ fn main() { _ = b => unreachable!(), }; }); - - // stream_select! macro - let _ = block_on(async { - let endless_ints = |i| stream::iter(vec![i].into_iter().cycle()); - - let mut endless_ones = futures04::stream_select!(endless_ints(1i32), stream::pending()); - assert_eq!(endless_ones.next().await, Some(1)); - assert_eq!(endless_ones.next().await, Some(1)); - - let mut finite_list = futures04::stream_select!(stream::iter(vec![1, 2, 3].into_iter())); - assert_eq!(finite_list.next().await, Some(1)); - assert_eq!(finite_list.next().await, Some(2)); - assert_eq!(finite_list.next().await, Some(3)); - assert_eq!(finite_list.next().await, None); - - let endless_mixed = - futures04::stream_select!(endless_ints(1i32), endless_ints(2), endless_ints(3)); - // Take 100, and assert a somewhat even distribution of values. - // The fairness is randomized, but over 100 samples we should be pretty close to even. - // This test may be a bit flaky. Feel free to adjust the margins as you see fit. - let mut count = 0; - let results = endless_mixed - .take_while(move |_| { - count += 1; - let ret = count < 100; - async move { ret } - }) - .collect::>() - .await; - assert!(results.iter().filter(|x| **x == 1).count() >= 29); - assert!(results.iter().filter(|x| **x == 2).count() >= 29); - assert!(results.iter().filter(|x| **x == 3).count() >= 29); - }); }