Improve support for mutable borrows in select! blocks

As reported in #1811 mutable accesses were not allowed inside select!
blocks if a Future inside the same expression borrowed it. This is not
necessarily required, since the actual polling of the Futures occurs
before the result blocks are evaluated.

This change allows borrowing a variable inside a Future and referring to
the same variable inside other select branches. The change introduces
additional scoping into the select macro which separates the future poll
phase from executing the result blocks phase.

Fixes #1811

Author: Matthias247

futures-select-macro/src/lib.rs

@@ -250,58 +250,74 @@ pub fn select(input: TokenStream) -> TokenStream {
         #complete_branch
     };
 
-    let await_and_select = if let Some(default_expr) = parsed.default {
+    let await_select_fut = if parsed.default.is_some() {
+        // For select! with default this returns the Poll result
         quote! {
-            if let #futures_crate::task::Poll::Ready(x) =
-                __poll_fn(&mut #futures_crate::task::Context::from_waker(
-                    #futures_crate::task::noop_waker_ref()
-                ))
-            {
-                match x { #branches }
-            } else {
-                #default_expr
-            };
+            __poll_fn(&mut #futures_crate::task::Context::from_waker(
+                #futures_crate::task::noop_waker_ref()
+            ))
+        }
+    } else {
+        quote! {
+            #futures_crate::future::poll_fn(__poll_fn).await
+        }
+    };
+
+    let execute_result_expr = if let Some(default_expr) = &parsed.default {
+        // For select! with default __select_result is a Poll, otherwise not
+        quote! {
+            match __select_result {
+                #futures_crate::task::Poll::Ready(result) => match result {
+                    #branches
+                },
+                _ => #default_expr
+            }
         }
     } else {
         quote! {
-            match #futures_crate::future::poll_fn(__poll_fn).await {
+            match __select_result {
                 #branches
             }
         }
     };
 
     TokenStream::from(quote! { {
         #enum_item
-        #( #future_let_bindings )*
-
-        let mut __poll_fn = |__cx: &mut #futures_crate::task::Context<'_>| {
-            let mut __any_polled = false;
-
-            #( #poll_functions )*
 
-            let mut __select_arr = [#( #variant_names ),*];
-            #futures_crate::async_await::shuffle(&mut __select_arr);
-            for poller in &mut __select_arr {
-                let poller: &mut &mut dyn FnMut(
-                    &mut #futures_crate::task::Context<'_>
-                ) -> Option<#futures_crate::task::Poll<_>> = poller;
-                match poller(__cx) {
-                    Some(x @ #futures_crate::task::Poll::Ready(_)) =>
-                        return x,
-                    Some(#futures_crate::task::Poll::Pending) => {
-                        __any_polled = true;
+        let __select_result = {
+            #( #future_let_bindings )*
+
+            let mut __poll_fn = |__cx: &mut #futures_crate::task::Context<'_>| {
+                let mut __any_polled = false;
+
+                #( #poll_functions )*
+
+                let mut __select_arr = [#( #variant_names ),*];
+                #futures_crate::async_await::shuffle(&mut __select_arr);
+                for poller in &mut __select_arr {
+                    let poller: &mut &mut dyn FnMut(
+                        &mut #futures_crate::task::Context<'_>
+                    ) -> Option<#futures_crate::task::Poll<_>> = poller;
+                    match poller(__cx) {
+                        Some(x @ #futures_crate::task::Poll::Ready(_)) =>
+                            return x,
+                        Some(#futures_crate::task::Poll::Pending) => {
+                            __any_polled = true;
+                        }
+                        None => {}
                     }
-                    None => {}
                 }
-            }
 
-            if !__any_polled {
-                #none_polled
-            } else {
-                #futures_crate::task::Poll::Pending
-            }
+                if !__any_polled {
+                    #none_polled
+                } else {
+                    #futures_crate::task::Poll::Pending
+                }
+            };
+
+            #await_select_fut
         };
 
-        #await_and_select
+        #execute_result_expr
     } })
 }

futures/tests/async_await_macros.rs

@@ -3,9 +3,10 @@
 use futures::{Poll, pending, pin_mut, poll, join, try_join, select};
 use futures::channel::{mpsc, oneshot};
 use futures::executor::block_on;
-use futures::future::{self, FutureExt};
+use futures::future::{self, FutureExt, poll_fn};
 use futures::stream::StreamExt;
 use futures::sink::SinkExt;
+use futures::task::Context;
 
 #[test]
 fn poll_and_pending() {
@@ -180,6 +181,78 @@ fn select_on_non_unpin_expressions() {
     assert_eq!(res, 5);
 }
 
+#[test]
+fn select_can_be_used_as_expression() {
+    block_on(async {
+        let res = select! {
+            x = future::ready(7) => { x },
+            y = future::ready(3) => { y + 1 },
+        };
+        assert!(res == 7 || res == 4);
+    });
+}
+
+#[test]
+fn select_with_default_can_be_used_as_expression() {
+    fn poll_always_pending<T>(_cx: &mut Context<'_>) -> Poll<T> {
+        Poll::Pending
+    }
+
+    block_on(async {
+        let res = select! {
+            x = poll_fn(poll_always_pending::<i32>).fuse() => x,
+            y = poll_fn(poll_always_pending::<i32>).fuse() => { y + 1 },
+            default => 99,
+        };
+        assert_eq!(res, 99);
+    });
+}
+
+#[test]
+fn select_with_complete_can_be_used_as_expression() {
+    block_on(async {
+        let res = select! {
+            x = future::pending::<i32>() => { x },
+            y = future::pending::<i32>() => { y + 1 },
+            default => 99,
+            complete => 237,
+        };
+        assert_eq!(res, 237);
+    });
+}
+
+async fn require_mutable(_: &mut i32) {}
+async fn async_noop() {}
+
+#[test]
+fn select_on_mutable_borrowing_future_with_same_borrow_in_block() {
+    block_on(async {
+        let mut foo = 234;
+        select! {
+            x = require_mutable(&mut foo).fuse() => { },
+            y = async_noop().fuse() => {
+                foo += 5;
+            },
+        }
+    });
+}
+
+#[test]
+fn select_on_mutable_borrowing_future_with_same_borrow_in_block_and_default() {
+    block_on(async {
+        let mut foo = 234;
+        select! {
+            x = require_mutable(&mut foo).fuse() => { },
+            y = async_noop().fuse() => {
+                foo += 5;
+            },
+            default => {
+                foo += 27;
+            },
+        }
+    });
+}
+
 #[test]
 fn join_size() {
     let fut = async {