macros: fix mut patterns in select! macro (#4211)

tokio-macros/src/lib.rs

@@ -329,3 +329,11 @@ pub fn test_fail(_args: TokenStream, _item: TokenStream) -> TokenStream {
 pub fn select_priv_declare_output_enum(input: TokenStream) -> TokenStream {
     select::declare_output_enum(input)
 }
+
+/// 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_clean_pattern(input: TokenStream) -> TokenStream {
+    select::clean_pattern_macro(input)
+}

tokio-macros/src/select.rs

@@ -41,3 +41,70 @@ pub(crate) fn declare_output_enum(input: TokenStream) -> TokenStream {
         pub(super) type Mask = #mask;
     })
 }
+
+pub(crate) fn clean_pattern_macro(input: TokenStream) -> TokenStream {
+    // If this isn't a pattern, we return the token stream as-is. The select!
+    // macro is using it in a location requiring a pattern, so an error will be
+    // emitted there.
+    let mut input: syn::Pat = match syn::parse(input.clone()) {
+        Ok(it) => it,
+        Err(_) => return input,
+    };
+
+    clean_pattern(&mut input);
+    quote::ToTokens::into_token_stream(input).into()
+}
+
+// Removes any occurrences of ref or mut in the provided pattern.
+fn clean_pattern(pat: &mut syn::Pat) {
+    match pat {
+        syn::Pat::Box(_box) => {}
+        syn::Pat::Lit(_literal) => {}
+        syn::Pat::Macro(_macro) => {}
+        syn::Pat::Path(_path) => {}
+        syn::Pat::Range(_range) => {}
+        syn::Pat::Rest(_rest) => {}
+        syn::Pat::Verbatim(_tokens) => {}
+        syn::Pat::Wild(_underscore) => {}
+        syn::Pat::Ident(ident) => {
+            ident.by_ref = None;
+            ident.mutability = None;
+            if let Some((_at, pat)) = &mut ident.subpat {
+                clean_pattern(&mut *pat);
+            }
+        }
+        syn::Pat::Or(or) => {
+            for case in or.cases.iter_mut() {
+                clean_pattern(case);
+            }
+        }
+        syn::Pat::Slice(slice) => {
+            for elem in slice.elems.iter_mut() {
+                clean_pattern(elem);
+            }
+        }
+        syn::Pat::Struct(struct_pat) => {
+            for field in struct_pat.fields.iter_mut() {
+                clean_pattern(&mut field.pat);
+            }
+        }
+        syn::Pat::Tuple(tuple) => {
+            for elem in tuple.elems.iter_mut() {
+                clean_pattern(elem);
+            }
+        }
+        syn::Pat::TupleStruct(tuple) => {
+            for elem in tuple.pat.elems.iter_mut() {
+                clean_pattern(elem);
+            }
+        }
+        syn::Pat::Reference(reference) => {
+            reference.mutability = None;
+            clean_pattern(&mut *reference.pat);
+        }
+        syn::Pat::Type(type_pat) => {
+            clean_pattern(&mut *type_pat.pat);
+        }
+        _ => {}
+    }
+}

tokio/Cargo.toml

@@ -87,7 +87,7 @@ test-util = ["rt", "sync", "time"]
 time = []
 
 [dependencies]
-tokio-macros = { version = "1.1.0", path = "../tokio-macros", optional = true }
+tokio-macros = { path = "../tokio-macros", optional = true }
 
 pin-project-lite = "0.2.0"
 

tokio/src/lib.rs

@@ -497,6 +497,12 @@ cfg_macros! {
     #[doc(hidden)]
     pub use tokio_macros::select_priv_declare_output_enum;
 
+    /// 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.
+    #[doc(hidden)]
+    pub use tokio_macros::select_priv_clean_pattern;
+
     cfg_rt! {
         #[cfg(feature = "rt-multi-thread")]
         #[cfg(not(test))] // Work around for rust-lang/rust#62127

tokio/src/macros/select.rs

@@ -520,7 +520,7 @@ macro_rules! select {
                                 #[allow(unused_variables)]
                                 #[allow(unused_mut)]
                                 match &out {
-                                    $bind => {}
+                                    $crate::select_priv_clean_pattern!($bind) => {}
                                     _ => continue,
                                 }
 

tokio/tests/macros_select.rs

@@ -558,3 +558,29 @@ pub async fn default_numeric_fallback() {
         else => (),
     }
 }
+
+// https://github.com/tokio-rs/tokio/issues/4182
+#[tokio::test]
+async fn mut_ref_patterns() {
+    tokio::select! {
+        Some(mut foo) = async { Some("1".to_string()) } => {
+            assert_eq!(foo, "1");
+            foo = "2".to_string();
+            assert_eq!(foo, "2");
+        },
+    };
+
+    tokio::select! {
+        Some(ref foo) = async { Some("1".to_string()) } => {
+            assert_eq!(*foo, "1");
+        },
+    };
+
+    tokio::select! {
+        Some(ref mut foo) = async { Some("1".to_string()) } => {
+            assert_eq!(*foo, "1");
+            *foo = "2".to_string();
+            assert_eq!(*foo, "2");
+        },
+    };
+}