Enhance mutable closure suggestions with as_mut() support

compiler/rustc_hir_typeck/src/fn_ctxt/suggestions.rs

@@ -2855,7 +2855,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         true
     }
 
-    /// Identify some cases where `as_ref()` would be appropriate and suggest it.
+    /// Identify some cases where `as_ref()` or `as_mut()` would be appropriate and suggest it.
     ///
     /// Given the following code:
     /// ```compile_fail,E0308
@@ -2871,7 +2871,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     /// ```ignore (illustrative)
     /// opt.map(|param| { takes_ref(param) });
     /// ```
-    fn can_use_as_ref(&self, expr: &hir::Expr<'_>) -> Option<(Vec<(Span, String)>, &'static str)> {
+    fn can_use_as_ref_or_mut(
+        &self,
+        expr: &hir::Expr<'_>,
+        mutability: hir::Mutability,
+    ) -> Option<(Vec<(Span, String)>, &'static str)> {
         let hir::ExprKind::Path(hir::QPath::Resolved(_, path)) = expr.kind else {
             return None;
         };
@@ -2908,21 +2912,30 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             return None;
         };
 
-        let self_ty = self.typeck_results.borrow().expr_ty_opt(receiver)?;
+        let mut self_ty = self.typeck_results.borrow().expr_ty_opt(receiver)?;
+        while let ty::Ref(_, inner, ref_mutability) = self_ty.kind() {
+            // `as_mut()` cannot borrow through a shared reference,
+            // also we cannot suggest `as_ref()` either when the reference is shared
+            if mutability.is_mut() && ref_mutability.is_not() {
+                return None;
+            }
+            self_ty = *inner;
+        }
         let name = method_path.ident.name;
-        let is_as_ref_able = match self_ty.peel_refs().kind() {
+        let can_borrow = match self_ty.kind() {
             ty::Adt(def, _) => {
                 (self.tcx.is_diagnostic_item(sym::Option, def.did())
                     || self.tcx.is_diagnostic_item(sym::Result, def.did()))
                     && (name == sym::map || name == sym::and_then)
             }
             _ => false,
         };
-        if is_as_ref_able {
-            Some((
-                vec![(method_path.ident.span.shrink_to_lo(), "as_ref().".to_string())],
-                "consider using `as_ref` instead",
-            ))
+        if can_borrow {
+            let (suggestion, message) = match mutability {
+                hir::Mutability::Not => ("as_ref().", "consider using `as_ref` instead"),
+                hir::Mutability::Mut => ("as_mut().", "consider using `as_mut` instead"),
+            };
+            Some((vec![(method_path.ident.span.shrink_to_lo(), suggestion.to_string())], message))
         } else {
             None
         }
@@ -3113,7 +3126,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         return Some((suggs, help, app, mutref));
                     }
 
-                    if let Some((sugg, msg)) = self.can_use_as_ref(expr) {
+                    if let Some((sugg, msg)) = self.can_use_as_ref_or_mut(expr, mutability) {
                         return Some((
                             sugg,
                             msg.to_string(),

tests/ui/suggestions/as-mut-closure.fixed

@@ -0,0 +1,34 @@
+// Borrowing a container for a closure argument must preserve the expected mutability.
+//@ run-rustfix
+
+#![allow(unused_mut)]
+
+fn append(value: &mut String) {
+    value.push('!');
+}
+
+fn increment(value: &mut i32) {
+    *value += 1;
+}
+
+fn main() {
+    let mut option = Some(String::new());
+    let _ = option.as_mut().map(|arg| append(arg));
+    //~^ ERROR mismatched types
+    let _ = (&mut option).as_mut().and_then(|arg| Some(append(arg)));
+    //~^ ERROR mismatched types
+
+    let mut result: Result<_, ()> = Ok(String::new());
+    let _ = result.as_mut().map(|arg| append(arg));
+    //~^ ERROR mismatched types
+    let _ = (&mut result).as_mut().and_then(|arg| Ok(append(arg)));
+    //~^ ERROR mismatched types
+
+    // A shared reference cannot supply `as_mut()`. Borrow the copied argument instead.
+    let shared = &Some(0);
+    let _ = shared.map(|mut arg| increment(&mut arg));
+    //~^ ERROR mismatched types
+    let nested = &mut &Some(0);
+    let _ = nested.map(|mut arg| increment(&mut arg));
+    //~^ ERROR mismatched types
+}

tests/ui/suggestions/as-mut-closure.rs

@@ -0,0 +1,34 @@
+// Borrowing a container for a closure argument must preserve the expected mutability.
+//@ run-rustfix
+
+#![allow(unused_mut)]
+
+fn append(value: &mut String) {
+    value.push('!');
+}
+
+fn increment(value: &mut i32) {
+    *value += 1;
+}
+
+fn main() {
+    let mut option = Some(String::new());
+    let _ = option.map(|arg| append(arg));
+    //~^ ERROR mismatched types
+    let _ = (&mut option).and_then(|arg| Some(append(arg)));
+    //~^ ERROR mismatched types
+
+    let mut result: Result<_, ()> = Ok(String::new());
+    let _ = result.map(|arg| append(arg));
+    //~^ ERROR mismatched types
+    let _ = (&mut result).and_then(|arg| Ok(append(arg)));
+    //~^ ERROR mismatched types
+
+    // A shared reference cannot supply `as_mut()`. Borrow the copied argument instead.
+    let shared = &Some(0);
+    let _ = shared.map(|mut arg| increment(arg));
+    //~^ ERROR mismatched types
+    let nested = &mut &Some(0);
+    let _ = nested.map(|mut arg| increment(arg));
+    //~^ ERROR mismatched types
+}

tests/ui/suggestions/as-mut-closure.stderr

@@ -0,0 +1,111 @@
+error[E0308]: mismatched types
+  --> $DIR/as-mut-closure.rs:16:37
+   |
+LL |     let _ = option.map(|arg| append(arg));
+   |                              ------ ^^^ expected `&mut String`, found `String`
+   |                              |
+   |                              arguments to this function are incorrect
+   |
+note: function defined here
+  --> $DIR/as-mut-closure.rs:6:4
+   |
+LL | fn append(value: &mut String) {
+   |    ^^^^^^ ------------------
+help: consider using `as_mut` instead
+   |
+LL |     let _ = option.as_mut().map(|arg| append(arg));
+   |                    +++++++++
+
+error[E0308]: mismatched types
+  --> $DIR/as-mut-closure.rs:18:54
+   |
+LL |     let _ = (&mut option).and_then(|arg| Some(append(arg)));
+   |                                               ------ ^^^ expected `&mut String`, found `String`
+   |                                               |
+   |                                               arguments to this function are incorrect
+   |
+note: function defined here
+  --> $DIR/as-mut-closure.rs:6:4
+   |
+LL | fn append(value: &mut String) {
+   |    ^^^^^^ ------------------
+help: consider using `as_mut` instead
+   |
+LL |     let _ = (&mut option).as_mut().and_then(|arg| Some(append(arg)));
+   |                           +++++++++
+
+error[E0308]: mismatched types
+  --> $DIR/as-mut-closure.rs:22:37
+   |
+LL |     let _ = result.map(|arg| append(arg));
+   |                              ------ ^^^ expected `&mut String`, found `String`
+   |                              |
+   |                              arguments to this function are incorrect
+   |
+note: function defined here
+  --> $DIR/as-mut-closure.rs:6:4
+   |
+LL | fn append(value: &mut String) {
+   |    ^^^^^^ ------------------
+help: consider using `as_mut` instead
+   |
+LL |     let _ = result.as_mut().map(|arg| append(arg));
+   |                    +++++++++
+
+error[E0308]: mismatched types
+  --> $DIR/as-mut-closure.rs:24:52
+   |
+LL |     let _ = (&mut result).and_then(|arg| Ok(append(arg)));
+   |                                             ------ ^^^ expected `&mut String`, found `String`
+   |                                             |
+   |                                             arguments to this function are incorrect
+   |
+note: function defined here
+  --> $DIR/as-mut-closure.rs:6:4
+   |
+LL | fn append(value: &mut String) {
+   |    ^^^^^^ ------------------
+help: consider using `as_mut` instead
+   |
+LL |     let _ = (&mut result).as_mut().and_then(|arg| Ok(append(arg)));
+   |                           +++++++++
+
+error[E0308]: mismatched types
+  --> $DIR/as-mut-closure.rs:29:44
+   |
+LL |     let _ = shared.map(|mut arg| increment(arg));
+   |                                  --------- ^^^ expected `&mut i32`, found integer
+   |                                  |
+   |                                  arguments to this function are incorrect
+   |
+note: function defined here
+  --> $DIR/as-mut-closure.rs:10:4
+   |
+LL | fn increment(value: &mut i32) {
+   |    ^^^^^^^^^ ---------------
+help: consider mutably borrowing here
+   |
+LL |     let _ = shared.map(|mut arg| increment(&mut arg));
+   |                                            ++++
+
+error[E0308]: mismatched types
+  --> $DIR/as-mut-closure.rs:32:44
+   |
+LL |     let _ = nested.map(|mut arg| increment(arg));
+   |                                  --------- ^^^ expected `&mut i32`, found integer
+   |                                  |
+   |                                  arguments to this function are incorrect
+   |
+note: function defined here
+  --> $DIR/as-mut-closure.rs:10:4
+   |
+LL | fn increment(value: &mut i32) {
+   |    ^^^^^^^^^ ---------------
+help: consider mutably borrowing here
+   |
+LL |     let _ = nested.map(|mut arg| increment(&mut arg));
+   |                                            ++++
+
+error: aborting due to 6 previous errors
+
+For more information about this error, try `rustc --explain E0308`.