Display information about captured variable in FnMut error

src/libcore/future/mod.rs

@@ -56,6 +56,7 @@ pub const fn from_generator<T>(gen: T) -> impl Future<Output = T::Return>
 where
     T: Generator<ResumeTy, Yield = ()>,
 {
+    #[rustc_diagnostic_item = "gen_future"]
     struct GenFuture<T: Generator<ResumeTy, Yield = ()>>(T);
 
     // We rely on the fact that async/await futures are immovable in order to create

src/librustc_middle/mir/query.rs

@@ -171,7 +171,7 @@ pub struct ClosureOutlivesRequirement<'tcx> {
 #[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash)]
 #[derive(RustcEncodable, RustcDecodable, HashStable)]
 pub enum ConstraintCategory {
-    Return,
+    Return(ReturnConstraint),
     Yield,
     UseAsConst,
     UseAsStatic,
@@ -187,6 +187,7 @@ pub enum ConstraintCategory {
     SizedBound,
     Assignment,
     OpaqueType,
+    ClosureUpvar(hir::HirId),
 
     /// A "boring" constraint (caused by the given location) is one that
     /// the user probably doesn't want to see described in diagnostics,
@@ -204,6 +205,13 @@ pub enum ConstraintCategory {
     Internal,
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash)]
+#[derive(RustcEncodable, RustcDecodable, HashStable)]
+pub enum ReturnConstraint {
+    Normal,
+    ClosureUpvar(hir::HirId),
+}
+
 /// The subject of a `ClosureOutlivesRequirement` -- that is, the thing
 /// that must outlive some region.
 #[derive(Copy, Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]

src/librustc_mir/borrow_check/diagnostics/conflict_errors.rs

@@ -766,7 +766,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     category:
                         category
                         @
-                        (ConstraintCategory::Return
+                        (ConstraintCategory::Return(_)
                         | ConstraintCategory::CallArgument
                         | ConstraintCategory::OpaqueType),
                     from_closure: false,
@@ -1096,7 +1096,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         opt_place_desc: Option<&String>,
     ) -> Option<DiagnosticBuilder<'cx>> {
         let return_kind = match category {
-            ConstraintCategory::Return => "return",
+            ConstraintCategory::Return(_) => "return",
             ConstraintCategory::Yield => "yield",
             _ => return None,
         };
@@ -1210,7 +1210,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         );
 
         let msg = match category {
-            ConstraintCategory::Return | ConstraintCategory::OpaqueType => {
+            ConstraintCategory::Return(_) | ConstraintCategory::OpaqueType => {
                 format!("{} is returned here", kind)
             }
             ConstraintCategory::CallArgument => {

src/librustc_mir/borrow_check/diagnostics/region_errors.rs

@@ -5,9 +5,9 @@ use rustc_infer::infer::{
     error_reporting::nice_region_error::NiceRegionError,
     error_reporting::unexpected_hidden_region_diagnostic, NLLRegionVariableOrigin,
 };
-use rustc_middle::mir::ConstraintCategory;
+use rustc_middle::mir::{ConstraintCategory, ReturnConstraint};
 use rustc_middle::ty::{self, RegionVid, Ty};
-use rustc_span::symbol::kw;
+use rustc_span::symbol::{kw, sym};
 use rustc_span::Span;
 
 use crate::util::borrowck_errors;
@@ -26,7 +26,7 @@ impl ConstraintDescription for ConstraintCategory {
         // Must end with a space. Allows for empty names to be provided.
         match self {
             ConstraintCategory::Assignment => "assignment ",
-            ConstraintCategory::Return => "returning this value ",
+            ConstraintCategory::Return(_) => "returning this value ",
             ConstraintCategory::Yield => "yielding this value ",
             ConstraintCategory::UseAsConst => "using this value as a constant ",
             ConstraintCategory::UseAsStatic => "using this value as a static ",
@@ -37,6 +37,7 @@ impl ConstraintDescription for ConstraintCategory {
             ConstraintCategory::SizedBound => "proving this value is `Sized` ",
             ConstraintCategory::CopyBound => "copying this value ",
             ConstraintCategory::OpaqueType => "opaque type ",
+            ConstraintCategory::ClosureUpvar(_) => "closure capture ",
             ConstraintCategory::Boring
             | ConstraintCategory::BoringNoLocation
             | ConstraintCategory::Internal => "",
@@ -306,8 +307,8 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
         };
 
         let diag = match (category, fr_is_local, outlived_fr_is_local) {
-            (ConstraintCategory::Return, true, false) if self.is_closure_fn_mut(fr) => {
-                self.report_fnmut_error(&errci)
+            (ConstraintCategory::Return(kind), true, false) if self.is_closure_fn_mut(fr) => {
+                self.report_fnmut_error(&errci, kind)
             }
             (ConstraintCategory::Assignment, true, false)
             | (ConstraintCategory::CallArgument, true, false) => {
@@ -347,7 +348,11 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
     ///            executing...
     ///    = note: ...therefore, returned references to captured variables will escape the closure
     /// ```
-    fn report_fnmut_error(&self, errci: &ErrorConstraintInfo) -> DiagnosticBuilder<'tcx> {
+    fn report_fnmut_error(
+        &self,
+        errci: &ErrorConstraintInfo,
+        kind: ReturnConstraint,
+    ) -> DiagnosticBuilder<'tcx> {
         let ErrorConstraintInfo { outlived_fr, span, .. } = errci;
 
         let mut diag = self
@@ -356,19 +361,39 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
             .sess
             .struct_span_err(*span, "captured variable cannot escape `FnMut` closure body");
 
-        // We should check if the return type of this closure is in fact a closure - in that
-        // case, we can special case the error further.
-        let return_type_is_closure =
-            self.regioncx.universal_regions().unnormalized_output_ty.is_closure();
-        let message = if return_type_is_closure {
-            "returns a closure that contains a reference to a captured variable, which then \
-             escapes the closure body"
-        } else {
-            "returns a reference to a captured variable which escapes the closure body"
+        let mut output_ty = self.regioncx.universal_regions().unnormalized_output_ty;
+        if let ty::Opaque(def_id, _) = output_ty.kind {
+            output_ty = self.infcx.tcx.type_of(def_id)
+        };
+
+        debug!("report_fnmut_error: output_ty={:?}", output_ty);
+
+        let message = match output_ty.kind {
+            ty::Closure(_, _) => {
+                "returns a closure that contains a reference to a captured variable, which then \
+                 escapes the closure body"
+            }
+            ty::Adt(def, _) if self.infcx.tcx.is_diagnostic_item(sym::gen_future, def.did) => {
+                "returns an `async` block that contains a reference to a captured variable, which then \
+                 escapes the closure body"
+            }
+            _ => "returns a reference to a captured variable which escapes the closure body",
         };
 
         diag.span_label(*span, message);
 
+        if let ReturnConstraint::ClosureUpvar(upvar) = kind {
+            let def_id = match self.regioncx.universal_regions().defining_ty {
+                DefiningTy::Closure(def_id, _) => def_id,
+                ty @ _ => bug!("unexpected DefiningTy {:?}", ty),
+            };
+
+            let upvar_def_span = self.infcx.tcx.hir().span(upvar);
+            let upvar_span = self.infcx.tcx.upvars_mentioned(def_id).unwrap()[&upvar].span;
+            diag.span_label(upvar_def_span, "variable defined here");
+            diag.span_label(upvar_span, "variable captured here");
+        }
+
         match self.give_region_a_name(*outlived_fr).unwrap().source {
             RegionNameSource::NamedEarlyBoundRegion(fr_span)
             | RegionNameSource::NamedFreeRegion(fr_span)
@@ -506,7 +531,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
         outlived_fr_name.highlight_region_name(&mut diag);
 
         match (category, outlived_fr_is_local, fr_is_local) {
-            (ConstraintCategory::Return, true, _) => {
+            (ConstraintCategory::Return(_), true, _) => {
                 diag.span_label(
                     *span,
                     format!(

src/librustc_mir/borrow_check/mod.rs

@@ -218,6 +218,7 @@ fn do_mir_borrowck<'a, 'tcx>(
         &mut flow_inits,
         &mdpe.move_data,
         &borrow_set,
+        &upvars,
     );
 
     // Dump MIR results into a file, if that is enabled. This let us
@@ -2301,30 +2302,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
     /// be `self` in the current MIR, because that is the only time we directly access the fields
     /// of a closure type.
     pub fn is_upvar_field_projection(&self, place_ref: PlaceRef<'tcx>) -> Option<Field> {
-        let mut place_projection = place_ref.projection;
-        let mut by_ref = false;
-
-        if let [proj_base @ .., ProjectionElem::Deref] = place_projection {
-            place_projection = proj_base;
-            by_ref = true;
-        }
-
-        match place_projection {
-            [base @ .., ProjectionElem::Field(field, _ty)] => {
-                let tcx = self.infcx.tcx;
-                let base_ty = Place::ty_from(place_ref.local, base, self.body(), tcx).ty;
-
-                if (base_ty.is_closure() || base_ty.is_generator())
-                    && (!by_ref || self.upvars[field.index()].by_ref)
-                {
-                    Some(*field)
-                } else {
-                    None
-                }
-            }
-
-            _ => None,
-        }
+        path_utils::is_upvar_field_projection(self.infcx.tcx, &self.upvars, place_ref, self.body())
     }
 }
 

src/librustc_mir/borrow_check/nll.rs

@@ -39,6 +39,7 @@ use crate::borrow_check::{
     renumber,
     type_check::{self, MirTypeckRegionConstraints, MirTypeckResults},
     universal_regions::UniversalRegions,
+    Upvar,
 };
 
 crate type PoloniusOutput = Output<RustcFacts>;
@@ -166,6 +167,7 @@ pub(in crate::borrow_check) fn compute_regions<'cx, 'tcx>(
     flow_inits: &mut ResultsCursor<'cx, 'tcx, MaybeInitializedPlaces<'cx, 'tcx>>,
     move_data: &MoveData<'tcx>,
     borrow_set: &BorrowSet<'tcx>,
+    upvars: &[Upvar],
 ) -> NllOutput<'tcx> {
     let mut all_facts = AllFacts::enabled(infcx.tcx).then_some(AllFacts::default());
 
@@ -188,6 +190,7 @@ pub(in crate::borrow_check) fn compute_regions<'cx, 'tcx>(
             flow_inits,
             move_data,
             elements,
+            upvars,
         );
 
     if let Some(all_facts) = &mut all_facts {

src/librustc_mir/borrow_check/path_utils.rs

@@ -1,10 +1,11 @@
 use crate::borrow_check::borrow_set::{BorrowData, BorrowSet, TwoPhaseActivation};
 use crate::borrow_check::places_conflict;
 use crate::borrow_check::AccessDepth;
+use crate::borrow_check::Upvar;
 use crate::dataflow::indexes::BorrowIndex;
 use rustc_data_structures::graph::dominators::Dominators;
 use rustc_middle::mir::BorrowKind;
-use rustc_middle::mir::{BasicBlock, Body, Location, Place};
+use rustc_middle::mir::{BasicBlock, Body, Field, Location, Place, PlaceRef, ProjectionElem};
 use rustc_middle::ty::TyCtxt;
 
 /// Returns `true` if the borrow represented by `kind` is
@@ -135,3 +136,38 @@ pub(super) fn borrow_of_local_data(place: Place<'_>) -> bool {
     // Any errors will be caught on the initial borrow
     !place.is_indirect()
 }
+
+/// If `place` is a field projection, and the field is being projected from a closure type,
+/// then returns the index of the field being projected. Note that this closure will always
+/// be `self` in the current MIR, because that is the only time we directly access the fields
+/// of a closure type.
+pub(crate) fn is_upvar_field_projection(
+    tcx: TyCtxt<'tcx>,
+    upvars: &[Upvar],
+    place_ref: PlaceRef<'tcx>,
+    body: &Body<'tcx>,
+) -> Option<Field> {
+    let mut place_projection = place_ref.projection;
+    let mut by_ref = false;
+
+    if let [proj_base @ .., ProjectionElem::Deref] = place_projection {
+        place_projection = proj_base;
+        by_ref = true;
+    }
+
+    match place_projection {
+        [base @ .., ProjectionElem::Field(field, _ty)] => {
+            let base_ty = Place::ty_from(place_ref.local, base, body, tcx).ty;
+
+            if (base_ty.is_closure() || base_ty.is_generator())
+                && (!by_ref || upvars[field.index()].by_ref)
+            {
+                Some(*field)
+            } else {
+                None
+            }
+        }
+
+        _ => None,
+    }
+}

src/librustc_mir/borrow_check/region_infer/mod.rs

@@ -12,7 +12,7 @@ use rustc_infer::infer::region_constraints::{GenericKind, VarInfos, VerifyBound}
 use rustc_infer::infer::{InferCtxt, NLLRegionVariableOrigin, RegionVariableOrigin};
 use rustc_middle::mir::{
     Body, ClosureOutlivesRequirement, ClosureOutlivesSubject, ClosureRegionRequirements,
-    ConstraintCategory, Local, Location,
+    ConstraintCategory, Local, Location, ReturnConstraint,
 };
 use rustc_middle::ty::{self, subst::SubstsRef, RegionVid, Ty, TyCtxt, TypeFoldable};
 use rustc_span::Span;
@@ -2017,7 +2017,7 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                     | ConstraintCategory::BoringNoLocation
                     | ConstraintCategory::Internal => false,
                     ConstraintCategory::TypeAnnotation
-                    | ConstraintCategory::Return
+                    | ConstraintCategory::Return(_)
                     | ConstraintCategory::Yield => true,
                     _ => constraint_sup_scc != target_scc,
                 }
@@ -2042,14 +2042,26 @@ impl<'tcx> RegionInferenceContext<'tcx> {
 
         if let Some(i) = best_choice {
             if let Some(next) = categorized_path.get(i + 1) {
-                if categorized_path[i].0 == ConstraintCategory::Return
+                if matches!(categorized_path[i].0, ConstraintCategory::Return(_))
                     && next.0 == ConstraintCategory::OpaqueType
                 {
                     // The return expression is being influenced by the return type being
                     // impl Trait, point at the return type and not the return expr.
                     return *next;
                 }
             }
+
+            if categorized_path[i].0 == ConstraintCategory::Return(ReturnConstraint::Normal) {
+                let field = categorized_path.iter().find_map(|p| {
+                    if let ConstraintCategory::ClosureUpvar(f) = p.0 { Some(f) } else { None }
+                });
+
+                if let Some(field) = field {
+                    categorized_path[i].0 =
+                        ConstraintCategory::Return(ReturnConstraint::ClosureUpvar(field));
+                }
+            }
+
             return categorized_path[i];
         }