Rollup merge of #126226 - gurry:125325-improve-closure-arg-sugg, r=oli-obk

Make suggestion to change `Fn` to `FnMut` work with methods as well

Fixes #125325

The issue occurred because the code that emitted the suggestion to change `Fn` to `FnMut` worked only for function calls  and not method calls. This PR makes it work with methods as well.

Author: GuillaumeGomez

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -937,56 +937,81 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
         let node = self.infcx.tcx.hir_node(fn_call_id);
         let def_id = hir.enclosing_body_owner(fn_call_id);
         let mut look_at_return = true;
-        // If we can detect the expression to be an `fn` call where the closure was an argument,
-        // we point at the `fn` definition argument...
-        if let hir::Node::Expr(hir::Expr { kind: hir::ExprKind::Call(func, args), .. }) = node {
-            let arg_pos = args
+
+        // If the HIR node is a function or method call gets the def ID
+        // of the called function or method and the span and args of the call expr
+        let get_call_details = || {
+            let hir::Node::Expr(hir::Expr { hir_id, kind, .. }) = node else {
+                return None;
+            };
+
+            let typeck_results = self.infcx.tcx.typeck(def_id);
+
+            match kind {
+                hir::ExprKind::Call(expr, args) => {
+                    if let Some(ty::FnDef(def_id, _)) =
+                        typeck_results.node_type_opt(expr.hir_id).as_ref().map(|ty| ty.kind())
+                    {
+                        Some((*def_id, expr.span, *args))
+                    } else {
+                        None
+                    }
+                }
+                hir::ExprKind::MethodCall(_, _, args, span) => {
+                    if let Some(def_id) = typeck_results.type_dependent_def_id(*hir_id) {
+                        Some((def_id, *span, *args))
+                    } else {
+                        None
+                    }
+                }
+                _ => None,
+            }
+        };
+
+        // If we can detect the expression to be an function or method call where the closure was an argument,
+        // we point at the function or method definition argument...
+        if let Some((callee_def_id, call_span, call_args)) = get_call_details() {
+            let arg_pos = call_args
                 .iter()
                 .enumerate()
                 .filter(|(_, arg)| arg.hir_id == closure_id)
                 .map(|(pos, _)| pos)
                 .next();
-            let tables = self.infcx.tcx.typeck(def_id);
-            if let Some(ty::FnDef(def_id, _)) =
-                tables.node_type_opt(func.hir_id).as_ref().map(|ty| ty.kind())
-            {
-                let arg = match hir.get_if_local(*def_id) {
-                    Some(
-                        hir::Node::Item(hir::Item {
-                            ident, kind: hir::ItemKind::Fn(sig, ..), ..
-                        })
-                        | hir::Node::TraitItem(hir::TraitItem {
-                            ident,
-                            kind: hir::TraitItemKind::Fn(sig, _),
-                            ..
+
+            let arg = match hir.get_if_local(callee_def_id) {
+                Some(
+                    hir::Node::Item(hir::Item { ident, kind: hir::ItemKind::Fn(sig, ..), .. })
+                    | hir::Node::TraitItem(hir::TraitItem {
+                        ident,
+                        kind: hir::TraitItemKind::Fn(sig, _),
+                        ..
+                    })
+                    | hir::Node::ImplItem(hir::ImplItem {
+                        ident,
+                        kind: hir::ImplItemKind::Fn(sig, _),
+                        ..
+                    }),
+                ) => Some(
+                    arg_pos
+                        .and_then(|pos| {
+                            sig.decl.inputs.get(
+                                pos + if sig.decl.implicit_self.has_implicit_self() {
+                                    1
+                                } else {
+                                    0
+                                },
+                            )
                         })
-                        | hir::Node::ImplItem(hir::ImplItem {
-                            ident,
-                            kind: hir::ImplItemKind::Fn(sig, _),
-                            ..
-                        }),
-                    ) => Some(
-                        arg_pos
-                            .and_then(|pos| {
-                                sig.decl.inputs.get(
-                                    pos + if sig.decl.implicit_self.has_implicit_self() {
-                                        1
-                                    } else {
-                                        0
-                                    },
-                                )
-                            })
-                            .map(|arg| arg.span)
-                            .unwrap_or(ident.span),
-                    ),
-                    _ => None,
-                };
-                if let Some(span) = arg {
-                    err.span_label(span, "change this to accept `FnMut` instead of `Fn`");
-                    err.span_label(func.span, "expects `Fn` instead of `FnMut`");
-                    err.span_label(closure_span, "in this closure");
-                    look_at_return = false;
-                }
+                        .map(|arg| arg.span)
+                        .unwrap_or(ident.span),
+                ),
+                _ => None,
+            };
+            if let Some(span) = arg {
+                err.span_label(span, "change this to accept `FnMut` instead of `Fn`");
+                err.span_label(call_span, "expects `Fn` instead of `FnMut`");
+                err.span_label(closure_span, "in this closure");
+                look_at_return = false;
             }
         }
 

tests/ui/closures/wrong-closure-arg-suggestion-125325.rs

@@ -0,0 +1,29 @@
+// Regression test for #125325
+
+// Tests that we suggest changing an `impl Fn` param
+// to `impl FnMut` when the provided closure arg
+// is trying to mutate the closure env.
+// Ensures that it works that way for both
+// functions and methods
+
+struct S;
+
+impl S {
+    fn assoc_func(&self, _f: impl Fn()) -> usize {
+        0
+    }
+}
+
+fn func(_f: impl Fn()) -> usize {
+    0
+}
+
+fn test_func(s: &S) -> usize {
+    let mut x = ();
+    s.assoc_func(|| x = ());
+    //~^ cannot assign to `x`, as it is a captured variable in a `Fn` closure
+    func(|| x = ())
+    //~^ cannot assign to `x`, as it is a captured variable in a `Fn` closure
+}
+
+fn main() {}

tests/ui/closures/wrong-closure-arg-suggestion-125325.stderr

@@ -0,0 +1,28 @@
+error[E0594]: cannot assign to `x`, as it is a captured variable in a `Fn` closure
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:23:21
+   |
+LL |     fn assoc_func(&self, _f: impl Fn()) -> usize {
+   |                              --------- change this to accept `FnMut` instead of `Fn`
+...
+LL |     s.assoc_func(|| x = ());
+   |       --------------^^^^^^-
+   |       |          |  |
+   |       |          |  cannot assign
+   |       |          in this closure
+   |       expects `Fn` instead of `FnMut`
+
+error[E0594]: cannot assign to `x`, as it is a captured variable in a `Fn` closure
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:25:13
+   |
+LL | fn func(_f: impl Fn()) -> usize {
+   |             --------- change this to accept `FnMut` instead of `Fn`
+...
+LL |     func(|| x = ())
+   |     ---- -- ^^^^^^ cannot assign
+   |     |    |
+   |     |    in this closure
+   |     expects `Fn` instead of `FnMut`
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0594`.