Rollup merge of #155877 - qaijuang:fix-155727-fnmut-diagnostic, r=wesleywiser

Avoid misleading return-type note for foreign `Fn` callees

Fixes #155727.

The issue occurred because the code that emitted the `Fn`/`FnMut` suggestion only avoided the return-type fallback when it could point at a local callee argument. The local case from #125325 was fixed in #126226, but nested cross-crate cases could still suggest changing the enclosing function's return type even though the relevant `Fn` requirement came from the callee argument.

This is also the broader diagnostic shape discussed in #119985.

This PR checks the instantiated callee predicates for an exact `Fn` bound
on the closure argument, and extends `wrong-closure-arg-suggestion-125325` with cross-crate function and method cases for that path.

Author: JonathanBrouwer

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -4,6 +4,7 @@ use either::Either;
 use hir::{ExprKind, Param};
 use rustc_abi::FieldIdx;
 use rustc_errors::{Applicability, Diag};
+use rustc_hir::def_id::DefId;
 use rustc_hir::intravisit::Visitor;
 use rustc_hir::{self as hir, BindingMode, ByRef, Node};
 use rustc_middle::bug;
@@ -1092,42 +1093,65 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
         let mut look_at_return = true;
 
         err.span_label(closure_span, "in this closure");
-        // If the HIR node is a function or method call, get the DefId
-        // of the callee function or method, 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 closure_arg_has_fn_trait_bound =
+            |callee_def_id, input_index, generic_args: ty::GenericArgsRef<'tcx>| {
+                let sig = tcx.fn_sig(callee_def_id).instantiate(tcx, generic_args).skip_binder();
+                let Some(input_ty): Option<Ty<'tcx>> = sig.inputs().get(input_index).copied()
+                else {
+                    return false;
+                };
+
+                tcx.predicates_of(callee_def_id)
+                    .instantiate(tcx, generic_args)
+                    .predicates
+                    .iter()
+                    .any(|predicate| {
+                        predicate.as_trait_clause().is_some_and(|trait_pred| {
+                            trait_pred.polarity() == ty::PredicatePolarity::Positive
+                                && tcx.fn_trait_kind_from_def_id(trait_pred.def_id())
+                                    == Some(ty::ClosureKind::Fn)
+                                && trait_pred.self_ty().skip_binder().peel_refs()
+                                    == input_ty.peel_refs()
+                        })
+                    })
             };
 
-            let typeck_results = tcx.typeck(def_id);
+        // If the HIR node is a function or method call, get the DefId
+        // of the callee function or method, the span, and argument info for the call expr.
+        let get_call_details =
+            || -> Option<(DefId, Span, usize, usize, ty::GenericArgsRef<'tcx>)> {
+                let hir::Node::Expr(hir::Expr { hir_id, kind, .. }) = node else {
+                    return None;
+                };
 
-            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
+                let typeck_results = tcx.typeck(def_id);
+
+                match kind {
+                    hir::ExprKind::Call(expr, args) => {
+                        if let Some(ty::FnDef(def_id, generic_args)) =
+                            typeck_results.node_type_opt(expr.hir_id).as_ref().map(|ty| ty.kind())
+                        {
+                            let arg_pos = args.iter().position(|arg| arg.hir_id == closure_id)?;
+                            Some((*def_id, expr.span, arg_pos, arg_pos, generic_args))
+                        } else {
+                            None
+                        }
                     }
+                    hir::ExprKind::MethodCall(_, _, args, span) => {
+                        let arg_pos = args.iter().position(|arg| arg.hir_id == closure_id)?;
+                        let def_id = typeck_results.type_dependent_def_id(*hir_id)?;
+                        let generic_args = typeck_results.node_args_opt(*hir_id)?;
+                        Some((def_id, *span, arg_pos, arg_pos + 1, generic_args))
+                    }
+                    _ => None,
                 }
-                hir::ExprKind::MethodCall(_, _, args, span) => typeck_results
-                    .type_dependent_def_id(*hir_id)
-                    .map(|def_id| (def_id, *span, *args)),
-                _ => None,
-            }
-        };
+            };
 
         // If we can detect the expression to be a 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();
-
+        if let Some((callee_def_id, call_span, arg_pos, input_index, generic_args)) =
+            get_call_details()
+        {
             let arg = match tcx.hir_get_if_local(callee_def_id) {
                 Some(
                     hir::Node::Item(hir::Item {
@@ -1144,16 +1168,12 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                         ..
                     }),
                 ) => Some(
-                    arg_pos
-                        .and_then(|pos| {
-                            sig.decl.inputs.get(
-                                pos + if sig.decl.implicit_self().has_implicit_self() {
-                                    1
-                                } else {
-                                    0
-                                },
-                            )
-                        })
+                    sig.decl
+                        .inputs
+                        .get(
+                            arg_pos
+                                + if sig.decl.implicit_self().has_implicit_self() { 1 } else { 0 },
+                        )
                         .map(|arg| arg.span)
                         .unwrap_or(ident.span),
                 ),
@@ -1163,6 +1183,13 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                 err.span_label(span, "change this to accept `FnMut` instead of `Fn`");
                 err.span_label(call_span, "expects `Fn` instead of `FnMut`");
                 look_at_return = false;
+            } else if closure_arg_has_fn_trait_bound(callee_def_id, input_index, generic_args) {
+                // The callee is not local, so we cannot point at its argument declaration, but we
+                // can still explain that this call site expects an `Fn` closure. Avoid falling
+                // through to the enclosing function's return type, which is misleading in cases
+                // like `flat_map(|_| external::map(|_| ...))`.
+                err.span_label(call_span, "expects `Fn` instead of `FnMut`");
+                look_at_return = false;
             }
         }
 

tests/ui/closures/auxiliary/wrong-closure-arg-suggestion-aux.rs

@@ -0,0 +1,34 @@
+pub struct P;
+pub struct Map<F>(pub F);
+
+pub trait PIter: Sized {
+    fn map<F>(self, f: F) -> Map<F>
+    where
+        F: Fn(i32);
+
+    fn flat_map<F, U>(self, f: F)
+    where
+        F: Fn(i32) -> U;
+}
+
+impl PIter for P {
+    fn map<F>(self, f: F) -> Map<F>
+    where
+        F: Fn(i32),
+    {
+        Map(f)
+    }
+
+    fn flat_map<F, U>(self, _: F)
+    where
+        F: Fn(i32) -> U,
+    {
+    }
+}
+
+pub fn to_fn<F>(f: F) -> Map<F>
+where
+    F: Fn(i32),
+{
+    Map(f)
+}

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

@@ -1,3 +1,5 @@
+//@ aux-build:wrong-closure-arg-suggestion-aux.rs
+
 // Regression test for #125325
 
 // Tests that we suggest changing an `impl Fn` param
@@ -6,6 +8,10 @@
 // Ensures that it works that way for both
 // functions and methods
 
+extern crate wrong_closure_arg_suggestion_aux as aux;
+
+use aux::{P, PIter, to_fn};
+
 struct S;
 
 impl S {
@@ -26,4 +32,37 @@ fn test_func(s: &S) -> usize {
     //~^ ERROR cannot assign to `x`, as it is a captured variable in a `Fn` closure
 }
 
+// Regression test for <https://github.com/rust-lang/rust/issues/155727>.
+//
+// When the relevant `Fn` bound comes from a non-local callee, we should still
+// explain the call-site expectation instead of falling back to the enclosing
+// function's return type.
+struct Counter {
+    counter: i32,
+}
+
+impl Counter {
+    fn external_fn(mut self) -> i32 {
+        take(|_| to_fn(|_| self.counter += 1));
+        //~^ ERROR cannot assign to `self.counter`, as `Fn` closures cannot mutate their captured variables [E0594]
+        //~| ERROR lifetime may not live long enough
+        //~| ERROR cannot borrow `self` as mutable, as it is a captured variable in a `Fn` closure [E0596]
+        self.counter
+    }
+
+    fn external_method(mut self) -> i32 {
+        P.flat_map(|_| P.map(|_| self.counter += 1));
+        //~^ ERROR cannot assign to `self.counter`, as `Fn` closures cannot mutate their captured variables [E0594]
+        //~| ERROR lifetime may not live long enough
+        //~| ERROR cannot borrow `self` as mutable, as it is a captured variable in a `Fn` closure [E0596]
+        self.counter
+    }
+}
+
+fn take<F, U>(_: F)
+where
+    F: Fn(i32) -> U,
+{
+}
+
 fn main() {}

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

@@ -1,5 +1,5 @@
 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
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:29:21
    |
 LL |     fn assoc_func(&self, _f: impl Fn()) -> usize {
    |                              --------- change this to accept `FnMut` instead of `Fn`
@@ -14,7 +14,7 @@ LL |     s.assoc_func(|| x = ());
    |       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
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:31:13
    |
 LL | fn func(_f: impl Fn()) -> usize {
    |             --------- change this to accept `FnMut` instead of `Fn`
@@ -28,6 +28,76 @@ LL |     func(|| x = ())
    |     |    in this closure
    |     expects `Fn` instead of `FnMut`
 
-error: aborting due to 2 previous errors
+error[E0594]: cannot assign to `self.counter`, as `Fn` closures cannot mutate their captured variables
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:46:28
+   |
+LL |         take(|_| to_fn(|_| self.counter += 1));
+   |                  ----- --- ^^^^^^^^^^^^^^^^^ cannot assign
+   |                  |     |
+   |                  |     in this closure
+   |                  expects `Fn` instead of `FnMut`
+
+error: lifetime may not live long enough
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:46:18
+   |
+LL |         take(|_| to_fn(|_| self.counter += 1));
+   |              --- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ returning this value requires that `'1` must outlive `'2`
+   |              | |
+   |              | return type of closure `aux::Map<{closure@$DIR/wrong-closure-arg-suggestion-125325.rs:46:24: 46:27}>` contains a lifetime `'2`
+   |              lifetime `'1` represents this closure's body
+   |
+   = note: closure implements `Fn`, so references to captured variables can't escape the closure
+
+error[E0596]: cannot borrow `self` as mutable, as it is a captured variable in a `Fn` closure
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:46:24
+   |
+LL |         take(|_| to_fn(|_| self.counter += 1));
+   |         ---- ---       ^^^ ------------ mutable borrow occurs due to use of `self` in closure
+   |         |    |         |
+   |         |    |         cannot borrow as mutable
+   |         |    in this closure
+   |         expects `Fn` instead of `FnMut`
+...
+LL | fn take<F, U>(_: F)
+   |                  - change this to accept `FnMut` instead of `Fn`
+
+error[E0594]: cannot assign to `self.counter`, as `Fn` closures cannot mutate their captured variables
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:54:34
+   |
+LL |         P.flat_map(|_| P.map(|_| self.counter += 1));
+   |                          --------^^^^^^^^^^^^^^^^^-
+   |                          |   |   |
+   |                          |   |   cannot assign
+   |                          |   in this closure
+   |                          expects `Fn` instead of `FnMut`
+
+error: lifetime may not live long enough
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:54:24
+   |
+LL |         P.flat_map(|_| P.map(|_| self.counter += 1));
+   |                    --- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ returning this value requires that `'1` must outlive `'2`
+   |                    | |
+   |                    | return type of closure `aux::Map<{closure@$DIR/wrong-closure-arg-suggestion-125325.rs:54:30: 54:33}>` contains a lifetime `'2`
+   |                    lifetime `'1` represents this closure's body
+   |
+   = note: closure implements `Fn`, so references to captured variables can't escape the closure
+help: consider adding 'move' keyword before the nested closure
+   |
+LL |         P.flat_map(|_| P.map(move |_| self.counter += 1));
+   |                              ++++
+
+error[E0596]: cannot borrow `self` as mutable, as it is a captured variable in a `Fn` closure
+  --> $DIR/wrong-closure-arg-suggestion-125325.rs:54:30
+   |
+LL |         P.flat_map(|_| P.map(|_| self.counter += 1));
+   |           -------------------^^^--------------------
+   |           |        |         |   |
+   |           |        |         |   mutable borrow occurs due to use of `self` in closure
+   |           |        |         cannot borrow as mutable
+   |           |        in this closure
+   |           expects `Fn` instead of `FnMut`
+
+error: aborting due to 8 previous errors
 
-For more information about this error, try `rustc --explain E0594`.
+Some errors have detailed explanations: E0594, E0596.
+For more information about an error, try `rustc --explain E0594`.