Create the Pin for selecting on !Unpin futures lazily inside the poll function

Author: Matthias247

futures-select-macro/src/lib.rs

@@ -177,21 +177,28 @@ pub fn select(input: TokenStream) -> TokenStream {
                     // This prevents creating redundant stack space
                     // for them.
                     // Passing Futures by path requires those Futures to implement Unpin.
+                    // We check for this condition here in order to be able to
+                    // safely use Pin::new_unchecked(&mut #path) later on.
+                    future_let_bindings.push(quote! {
+                        #futures_crate::async_await::assert_fused_future(&mut #path);
+                        #futures_crate::async_await::assert_unpin(&mut #path);
+                    });
                     path
                 },
                 _ => {
                     // Bind and pin the resulting Future on the stack. This is
                     // necessary to support direct select! calls on !Unpin
-                    // Futures.
+                    // Futures. The Future is not explicitly pinned here with
+                    // a Pin call, but assumed as pinned. The actual Pin is
+                    // created inside the poll() function below to defer the
+                    // creation of the temporary pointer, which would otherwise
+                    // increase the size of the generated Future.
                     // Safety: This is safe since the lifetime of the Future
                     // is totally constraint to the lifetime of the select!
                     // expression, and the Future can't get moved inside it
                     // (it is shadowed).
                     future_let_bindings.push(quote! {
                         let mut #variant_name = #expr;
-                        let mut #variant_name = unsafe {
-                            ::core::pin::Pin::new_unchecked(&mut #variant_name)
-                        };
                     });
                     parse_quote! { #variant_name }
                 }
@@ -203,8 +210,19 @@ pub fn select(input: TokenStream) -> TokenStream {
     // to use for polling that individual future. These will then be put in an array.
     let poll_functions = bound_future_names.iter().zip(variant_names.iter())
         .map(|(bound_future_name, variant_name)| {
+            // Below we lazily create the Pin on the Future below.
+            // This is done in order to avoid allocating memory in the generator
+            // for the Pin variable.
+            // Safety: This is safe because one of the following condition applies:
+            // 1. The Future is passed by the caller by name, and we assert that
+            //    it implements Unpin.
+            // 2. The Future is created in scope of the select! function and will
+            //    not be moved for the duration of it. It is thereby stack-pinned
             quote! {
                 let mut #variant_name = |__cx: &mut #futures_crate::task::Context<'_>| {
+                    let mut #bound_future_name = unsafe {
+                        ::core::pin::Pin::new_unchecked(&mut #bound_future_name)
+                    };
                     if #futures_crate::future::FusedFuture::is_terminated(&#bound_future_name) {
                         None
                     } else {

futures/tests/async_await_macros.rs

@@ -169,6 +169,24 @@ fn select_on_non_unpin_expressions() {
         5
     }};
 
+    let res = block_on(async {
+        let select_res;
+        select! {
+            value_1 = make_non_unpin_fut().fuse() => { select_res = value_1 },
+            value_2 = make_non_unpin_fut().fuse() => { select_res = value_2 },
+        };
+        select_res
+    });
+    assert_eq!(res, 5);
+}
+
+#[test]
+fn select_on_non_unpin_expressions_with_default() {
+    // The returned Future is !Unpin
+    let make_non_unpin_fut = || { async {
+        5
+    }};
+
     let res = block_on(async {
         let select_res;
         select! {
@@ -181,6 +199,25 @@ fn select_on_non_unpin_expressions() {
     assert_eq!(res, 5);
 }
 
+#[test]
+fn select_on_non_unpin_size() {
+    // The returned Future is !Unpin
+    let make_non_unpin_fut = || { async {
+        5
+    }};
+
+    let fut = async {
+        let select_res;
+        select! {
+            value_1 = make_non_unpin_fut().fuse() => { select_res = value_1 },
+            value_2 = make_non_unpin_fut().fuse() => { select_res = value_2 },
+        };
+        select_res
+    };
+
+    assert_eq!(48, std::mem::size_of_val(&fut));
+}
+
 #[test]
 fn select_can_be_used_as_expression() {
     block_on(async {