Add suggestion to use closure argument instead of a capture on borrowck error

compiler/rustc_borrowck/src/borrowck_errors.rs

@@ -469,6 +469,7 @@ impl<'cx, 'tcx> crate::MirBorrowckCtxt<'cx, 'tcx> {
     }
 
     #[rustc_lint_diagnostics]
+    #[track_caller]
     pub(crate) fn struct_span_err_with_code<S: Into<MultiSpan>>(
         &self,
         sp: S,

compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs

@@ -1,3 +1,5 @@
+use std::iter;
+
 use either::Either;
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxIndexSet;
@@ -10,27 +12,26 @@ use rustc_hir::intravisit::{walk_block, walk_expr, Visitor};
 use rustc_hir::{AsyncGeneratorKind, GeneratorKind, LangItem};
 use rustc_infer::infer::TyCtxtInferExt;
 use rustc_infer::traits::ObligationCause;
+use rustc_middle::hir::nested_filter::OnlyBodies;
 use rustc_middle::mir::tcx::PlaceTy;
 use rustc_middle::mir::{
     self, AggregateKind, BindingForm, BorrowKind, ClearCrossCrate, ConstraintCategory,
     FakeReadCause, LocalDecl, LocalInfo, LocalKind, Location, Operand, Place, PlaceRef,
     ProjectionElem, Rvalue, Statement, StatementKind, Terminator, TerminatorKind, VarBindingForm,
 };
-use rustc_middle::ty::{self, suggest_constraining_type_params, PredicateKind, Ty};
+use rustc_middle::ty::{self, suggest_constraining_type_params, PredicateKind, Ty, TypeckResults};
 use rustc_middle::util::CallKind;
 use rustc_mir_dataflow::move_paths::{InitKind, MoveOutIndex, MovePathIndex};
 use rustc_span::def_id::LocalDefId;
 use rustc_span::hygiene::DesugaringKind;
-use rustc_span::symbol::{kw, sym};
+use rustc_span::symbol::{kw, sym, Ident};
 use rustc_span::{BytePos, Span, Symbol};
 use rustc_trait_selection::infer::InferCtxtExt;
 use rustc_trait_selection::traits::ObligationCtxt;
 
 use crate::borrow_set::TwoPhaseActivation;
 use crate::borrowck_errors;
-
 use crate::diagnostics::conflict_errors::StorageDeadOrDrop::LocalStorageDead;
-use crate::diagnostics::mutability_errors::mut_borrow_of_mutable_ref;
 use crate::diagnostics::{find_all_local_uses, CapturedMessageOpt};
 use crate::{
     borrow_set::BorrowData, diagnostics::Instance, prefixes::IsPrefixOf,
@@ -959,7 +960,8 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     &msg_borrow,
                     None,
                 );
-                self.suggest_binding_for_closure_capture_self(
+                self.suggest_binding_for_closure_capture_self(&mut err, &issued_spans);
+                self.suggest_using_closure_argument_instead_of_capture(
                     &mut err,
                     issued_borrow.borrowed_place,
                     &issued_spans,
@@ -982,6 +984,11 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     place,
                     issued_borrow.borrowed_place,
                 );
+                self.suggest_using_closure_argument_instead_of_capture(
+                    &mut err,
+                    issued_borrow.borrowed_place,
+                    &issued_spans,
+                );
                 err
             }
 
@@ -1268,22 +1275,161 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         }
     }
 
-    fn suggest_binding_for_closure_capture_self(
+    /// Suggest using closure argument instead of capture.
+    ///
+    /// For example:
+    /// ```ignore (illustrative)
+    /// struct S;
+    ///
+    /// impl S {
+    ///     fn call(&mut self, f: impl Fn(&mut Self)) { /* ... */ }
+    ///     fn x(&self) {}
+    /// }
+    ///
+    ///     let mut v = S;
+    ///     v.call(|this: &mut S| v.x());
+    /// //  ^\                    ^-- help: try using the closure argument: `this`
+    /// //    *-- error: cannot borrow `v` as mutable because it is also borrowed as immutable
+    /// ```
+    fn suggest_using_closure_argument_instead_of_capture(
         &self,
         err: &mut Diagnostic,
         borrowed_place: Place<'tcx>,
         issued_spans: &UseSpans<'tcx>,
     ) {
-        let UseSpans::ClosureUse { capture_kind_span, .. } = issued_spans else { return };
-        let hir = self.infcx.tcx.hir();
+        let &UseSpans::ClosureUse { capture_kind_span, .. } = issued_spans else { return };
+        let tcx = self.infcx.tcx;
+        let hir = tcx.hir();
 
-        // check whether the borrowed place is capturing `self` by mut reference
+        // Get the type of the local that we are trying to borrow
         let local = borrowed_place.local;
-        let Some(_) = self
-            .body
-            .local_decls
-            .get(local)
-            .map(|l| mut_borrow_of_mutable_ref(l, self.local_names[local])) else { return };
+        let local_ty = self.body.local_decls[local].ty;
+
+        // Get the body the error happens in
+        let Some(body_id) = hir.get(self.mir_hir_id()).body_id() else { return };
+
+        let body_expr = hir.body(body_id).value;
+
+        struct ClosureFinder<'hir> {
+            hir: rustc_middle::hir::map::Map<'hir>,
+            borrow_span: Span,
+            res: Option<(&'hir hir::Expr<'hir>, &'hir hir::Closure<'hir>)>,
+            /// The path expression with the `borrow_span` span
+            error_path: Option<(&'hir hir::Expr<'hir>, &'hir hir::QPath<'hir>)>,
+        }
+        impl<'hir> Visitor<'hir> for ClosureFinder<'hir> {
+            type NestedFilter = OnlyBodies;
+
+            fn nested_visit_map(&mut self) -> Self::Map {
+                self.hir
+            }
+
+            fn visit_expr(&mut self, ex: &'hir hir::Expr<'hir>) {
+                if let hir::ExprKind::Path(qpath) = &ex.kind
+                    && ex.span == self.borrow_span
+                {
+                    self.error_path = Some((ex, qpath));
+                }
+
+                if let hir::ExprKind::Closure(closure) = ex.kind
+                    && ex.span.contains(self.borrow_span)
+                    // To support cases like `|| { v.call(|this| v.get()) }`
+                    // FIXME: actually support such cases (need to figure out how to move from the capture place to original local)
+                    && self.res.as_ref().map_or(true, |(prev_res, _)| prev_res.span.contains(ex.span))
+                {
+                    self.res = Some((ex, closure));
+                }
+
+                hir::intravisit::walk_expr(self, ex);
+            }
+        }
+
+        // Find the closure that most tightly wraps `capture_kind_span`
+        let mut finder =
+            ClosureFinder { hir, borrow_span: capture_kind_span, res: None, error_path: None };
+        finder.visit_expr(body_expr);
+        let Some((closure_expr, closure)) = finder.res else { return };
+
+        let typeck_results: &TypeckResults<'_> =
+            tcx.typeck_opt_const_arg(self.body.source.with_opt_param().as_local().unwrap());
+
+        // Check that the parent of the closure is a method call,
+        // with receiver matching with local's type (modulo refs)
+        let parent = hir.parent_id(closure_expr.hir_id);
+        if let hir::Node::Expr(parent) = hir.get(parent) {
+            if let hir::ExprKind::MethodCall(_, recv, ..) = parent.kind {
+                let recv_ty = typeck_results.expr_ty(recv);
+
+                if recv_ty.peel_refs() != local_ty {
+                    return;
+                }
+            }
+        }
+
+        // Get closure's arguments
+        let ty::Closure(_, substs) = typeck_results.expr_ty(closure_expr).kind() else { unreachable!() };
+        let sig = substs.as_closure().sig();
+        let tupled_params =
+            tcx.erase_late_bound_regions(sig.inputs().iter().next().unwrap().map_bound(|&b| b));
+        let ty::Tuple(params) = tupled_params.kind() else { return };
+
+        // Find the first argument with a matching type, get its name
+        let Some((_, this_name)) = params
+            .iter()
+            .zip(hir.body_param_names(closure.body))
+            .find(|(param_ty, name)|{
+                // FIXME: also support deref for stuff like `Rc` arguments
+                param_ty.peel_refs() == local_ty && name != &Ident::empty()
+            })
+            else { return };
+
+        let spans;
+        if let Some((_path_expr, qpath)) = finder.error_path
+            && let hir::QPath::Resolved(_, path) = qpath
+            && let hir::def::Res::Local(local_id) = path.res
+        {
+            // Find all references to the problematic variable in this closure body
+
+            struct VariableUseFinder {
+                local_id: hir::HirId,
+                spans: Vec<Span>,
+            }
+            impl<'hir> Visitor<'hir> for VariableUseFinder {
+                fn visit_expr(&mut self, ex: &'hir hir::Expr<'hir>) {
+                    if let hir::ExprKind::Path(qpath) = &ex.kind
+                        && let hir::QPath::Resolved(_, path) = qpath
+                        && let hir::def::Res::Local(local_id) = path.res
+                        && local_id == self.local_id
+                    {
+                        self.spans.push(ex.span);
+                    }
+
+                    hir::intravisit::walk_expr(self, ex);
+                }
+            }
+
+            let mut finder = VariableUseFinder { local_id, spans: Vec::new() };
+            finder.visit_expr(hir.body(closure.body).value);
+
+            spans = finder.spans;
+        } else {
+            spans = vec![capture_kind_span];
+        }
+
+        err.multipart_suggestion(
+            "try using the closure argument",
+            iter::zip(spans, iter::repeat(this_name.to_string())).collect(),
+            Applicability::MaybeIncorrect,
+        );
+    }
+
+    fn suggest_binding_for_closure_capture_self(
+        &self,
+        err: &mut Diagnostic,
+        issued_spans: &UseSpans<'tcx>,
+    ) {
+        let UseSpans::ClosureUse { capture_kind_span, .. } = issued_spans else { return };
+        let hir = self.infcx.tcx.hir();
 
         struct ExpressionFinder<'hir> {
             capture_span: Span,

compiler/rustc_hir/src/hir.rs

@@ -3529,12 +3529,20 @@ impl<'hir> OwnerNode<'hir> {
 
     pub fn body_id(&self) -> Option<BodyId> {
         match self {
-            OwnerNode::TraitItem(TraitItem {
-                kind: TraitItemKind::Fn(_, TraitFn::Provided(body_id)),
+            OwnerNode::Item(Item {
+                kind:
+                    ItemKind::Static(_, _, body) | ItemKind::Const(_, body) | ItemKind::Fn(_, _, body),
                 ..
             })
-            | OwnerNode::ImplItem(ImplItem { kind: ImplItemKind::Fn(_, body_id), .. })
-            | OwnerNode::Item(Item { kind: ItemKind::Fn(.., body_id), .. }) => Some(*body_id),
+            | OwnerNode::TraitItem(TraitItem {
+                kind:
+                    TraitItemKind::Fn(_, TraitFn::Provided(body)) | TraitItemKind::Const(_, Some(body)),
+                ..
+            })
+            | OwnerNode::ImplItem(ImplItem {
+                kind: ImplItemKind::Fn(_, body) | ImplItemKind::Const(_, body),
+                ..
+            }) => Some(*body),
             _ => None,
         }
     }
@@ -3729,12 +3737,27 @@ impl<'hir> Node<'hir> {
 
     pub fn body_id(&self) -> Option<BodyId> {
         match self {
-            Node::TraitItem(TraitItem {
-                kind: TraitItemKind::Fn(_, TraitFn::Provided(body_id)),
+            Node::Item(Item {
+                kind:
+                    ItemKind::Static(_, _, body) | ItemKind::Const(_, body) | ItemKind::Fn(_, _, body),
+                ..
+            })
+            | Node::TraitItem(TraitItem {
+                kind:
+                    TraitItemKind::Fn(_, TraitFn::Provided(body)) | TraitItemKind::Const(_, Some(body)),
+                ..
+            })
+            | Node::ImplItem(ImplItem {
+                kind: ImplItemKind::Fn(_, body) | ImplItemKind::Const(_, body),
                 ..
             })
-            | Node::ImplItem(ImplItem { kind: ImplItemKind::Fn(_, body_id), .. })
-            | Node::Item(Item { kind: ItemKind::Fn(.., body_id), .. }) => Some(*body_id),
+            | Node::Expr(Expr {
+                kind:
+                    ExprKind::ConstBlock(AnonConst { body, .. })
+                    | ExprKind::Closure(Closure { body, .. })
+                    | ExprKind::Repeat(_, ArrayLen::Body(AnonConst { body, .. })),
+                ..
+            }) => Some(*body),
             _ => None,
         }
     }

tests/ui/async-await/feature-self-return-type.stderr

@@ -4,6 +4,7 @@ error[E0597]: `bar` does not live long enough
 LL |     let x = {
    |         - borrow later stored here
 LL |         let bar = 22;
+   |             --- binding `bar` declared here
 LL |         Foo::new(&bar).await
    |                  ^^^^ borrowed value does not live long enough
 LL |

tests/ui/async-await/issue-61949-self-return-type.stderr

@@ -13,6 +13,7 @@ error[E0597]: `bar` does not live long enough
 LL |     let x = {
    |         - borrow later stored here
 LL |         let bar = 22;
+   |             --- binding `bar` declared here
 LL |         Foo::new(&bar).await
    |                  ^^^^ borrowed value does not live long enough
 LL |

tests/ui/borrowck/issue-109271-pass-self-into-closure.fixed

@@ -0,0 +1,39 @@
+// run-rustfix
+#![allow(unused)]
+struct S;
+
+impl S {
+    fn call(&mut self, f: impl FnOnce((), &mut Self)) {
+        // change state or something ...
+        f((), self);
+        // change state or something ...
+    }
+
+    fn get(&self) {}
+    fn set(&mut self) {}
+}
+
+fn main() {
+    let mut v = S;
+
+    v.call(|(), this: &mut S| this.get());
+    //~^ error: cannot borrow `v` as mutable because it is also borrowed as immutable
+    v.call(|(), this: &mut S| this.set());
+    //~^ error: cannot borrow `v` as mutable more than once at a time
+    //~| error: cannot borrow `v` as mutable more than once at a time
+
+    v.call(|(), this: &mut S| {
+        //~^ error: cannot borrow `v` as mutable more than once at a time
+        //~| error: cannot borrow `v` as mutable more than once at a time
+
+        _ = this;
+        this.set();
+        this.get();
+        S::get(&this);
+
+        use std::ops::Add;
+        let v = 0u32;
+        _ = v + v;
+        _ = v.add(3);
+    });
+}

tests/ui/borrowck/issue-109271-pass-self-into-closure.rs

@@ -0,0 +1,39 @@
+// run-rustfix
+#![allow(unused)]
+struct S;
+
+impl S {
+    fn call(&mut self, f: impl FnOnce((), &mut Self)) {
+        // change state or something ...
+        f((), self);
+        // change state or something ...
+    }
+
+    fn get(&self) {}
+    fn set(&mut self) {}
+}
+
+fn main() {
+    let mut v = S;
+
+    v.call(|(), this: &mut S| v.get());
+    //~^ error: cannot borrow `v` as mutable because it is also borrowed as immutable
+    v.call(|(), this: &mut S| v.set());
+    //~^ error: cannot borrow `v` as mutable more than once at a time
+    //~| error: cannot borrow `v` as mutable more than once at a time
+
+    v.call(|(), this: &mut S| {
+        //~^ error: cannot borrow `v` as mutable more than once at a time
+        //~| error: cannot borrow `v` as mutable more than once at a time
+
+        _ = v;
+        v.set();
+        v.get();
+        S::get(&v);
+
+        use std::ops::Add;
+        let v = 0u32;
+        _ = v + v;
+        _ = v.add(3);
+    });
+}