Provide `select!` macro (#2152)

Provides a `select!` macro for concurrently waiting on multiple async
expressions. The macro has similar goals and syntax as the one provided
by the `futures` crate, but differs significantly in implementation.

First, this implementation does not require special traits to be
implemented on futures or streams (i.e., no `FuseFuture`). A design goal
is to be able to pass a "plain" async fn result into the select! macro.

Even without `FuseFuture`, this `select!` implementation is able to
handle all cases the `futures::select!` macro can handle. It does this
by supporting pre-poll conditions on branches and result pattern
matching. For pre-conditions, each branch is able to include a condition
that disables the branch if it evaluates to false. This allows the user
to guard futures that have already been polled, preventing double
polling. Pattern matching can be used to disable streams that complete.

A second big difference is the macro is implemented almost entirely as a
declarative macro. The biggest advantage to using this strategy is that
the user will not need to alter the rustc recursion limit except in the
most extreme cases.

The resulting future also tends to be smaller in many cases.

3 files changed, 55 insertions, 2 deletions

diff --git a/tokio-macros/Cargo.toml b/tokio-macros/Cargo.toml
index d69d48d6..3cc3a3e4 100644
--- a/tokio-macros/Cargo.toml
+++ b/tokio-macros/Cargo.toml
@@ -25,6 +25,7 @@ proc-macro = true
 [features]
 
 [dependencies]
+proc-macro2 = "1.0.7"
 quote = "1"
 syn = { version = "1.0.3", features = ["full"] }
 
diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
index 09b8b093..17554470 100644
--- a/tokio-macros/src/lib.rs
+++ b/tokio-macros/src/lib.rs
@@ -14,10 +14,11 @@
 
 //! Macros for use with Tokio
 
-mod entry;
-
 extern crate proc_macro;
 
+mod entry;
+mod select;
+
 use proc_macro::TokenStream;
 
 /// Marks async function to be executed by selected runtime.
@@ -198,3 +199,11 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
 pub fn test_basic(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::test(args, item, false)
 }
+
+/// Implementation detail of the `select!` macro. This macro is **not** intended
+/// to be used as part of the public API and is permitted to change.
+#[proc_macro]
+#[doc(hidden)]
+pub fn select_priv_declare_output_enum(input: TokenStream) -> TokenStream {
+    select::declare_output_enum(input)
+}
diff --git a/tokio-macros/src/select.rs b/tokio-macros/src/select.rs
new file mode 100644
index 00000000..ddb2e6a0
--- /dev/null
+++ b/tokio-macros/src/select.rs
@@ -0,0 +1,43 @@
+use proc_macro::{TokenStream, TokenTree};
+use proc_macro2::Span;
+use quote::quote;
+use syn::Ident;
+
+pub(crate) fn declare_output_enum(input: TokenStream) -> TokenStream {
+    // passed in is: `(_ _ _)` with one `_` per branch
+    let branches = match input.into_iter().next() {
+        Some(TokenTree::Group(group)) => group.stream().into_iter().count(),
+        _ => panic!("unexpected macro input"),
+    };
+
+    let variants = (0..branches)
+        .map(|num| Ident::new(&format!("_{}", num), Span::call_site()))
+        .collect::<Vec<_>>();
+
+    // Use a bitfield to track which futures completed
+    let mask = Ident::new(
+        if branches <= 8 {
+            "u8"
+        } else if branches <= 16 {
+            "u16"
+        } else if branches <= 32 {
+            "u32"
+        } else if branches <= 64 {
+            "u64"
+        } else {
+            panic!("up to 64 branches supported");
+        },
+        Span::call_site(),
+    );
+
+    TokenStream::from(quote! {
+        pub(super) enum Out<#( #variants ),*> {
+            #( #variants(#variants), )*
+            // Include a `Disabled` variant signifying that all select branches
+            // failed to resolve.
+            Disabled,
+        }
+
+        pub(super) type Mask = #mask;
+    })
+}