Display information about captured variable in FnMut error

Fixes #69446

When we encounter a region error involving an `FnMut` closure, we
display a specialized error message. However, we currently do not
tell the user which upvar was captured. This makes it difficult to
determine the cause of the error, especially when the closure is large.

This commit records marks constraints involving closure upvars
with `ConstraintCategory::ClosureUpvar`. When we decide to 'blame'
a `ConstraintCategory::Return`, we additionall store
the captured upvar if we found a `ConstraintCategory::ClosureUpvar` in
the path.

When generating an error message, we point to relevant spans if we have
closure upvar information available. We further customize the message if
an `async` closure is being returned, to make it clear that the captured
variable is being returned indirectly.

Author: Aaron1011

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];
         }