Keep track of the start of the argument block of a closure

15 files changed, 53 insertions, 29 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index 54bd25d6471..d30cbdca29d 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -1280,8 +1280,10 @@ pub struct Closure {
     pub movability: Movability,
     pub fn_decl: P<FnDecl>,
     pub body: P<Expr>,
-    /// The span of the argument block `|...|`.
+    /// The span of the declaration block: 'move |...| -> ...'
     pub fn_decl_span: Span,
+    /// The span of the argument block `|...|`
+    pub fn_arg_span: Span,
 }
 
 /// Limit types of a range (inclusive or exclusive)
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index 4e1dcb2842f..7e1aebe1920 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -1371,6 +1371,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
             fn_decl,
             body,
             fn_decl_span,
+            fn_arg_span: _,
         }) => {
             vis.visit_closure_binder(binder);
             vis.visit_asyncness(asyncness);
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index 5c69e535212..08166c7327f 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -840,6 +840,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
             fn_decl,
             body,
             fn_decl_span: _,
+            fn_arg_span: _,
         }) => {
             visitor.visit_fn(FnKind::Closure(binder, fn_decl, body), expression.span, expression.id)
         }
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index a00100ee0a8..4385552f2cc 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -176,6 +176,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     fn_decl,
                     body,
                     fn_decl_span,
+                    fn_arg_span,
                 }) => {
                     if let Async::Yes { closure_id, .. } = asyncness {
                         self.lower_expr_async_closure(
@@ -186,6 +187,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             fn_decl,
                             body,
                             *fn_decl_span,
+                            *fn_arg_span,
                         )
                     } else {
                         self.lower_expr_closure(
@@ -196,6 +198,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             fn_decl,
                             body,
                             *fn_decl_span,
+                            *fn_arg_span,
                         )
                     }
                 }
@@ -639,6 +642,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 fn_decl,
                 body,
                 fn_decl_span: self.lower_span(span),
+                fn_arg_span: None,
                 movability: Some(hir::Movability::Static),
             });
 
@@ -887,6 +891,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         decl: &FnDecl,
         body: &Expr,
         fn_decl_span: Span,
+        fn_arg_span: Span,
     ) -> hir::ExprKind<'hir> {
         let (binder_clause, generic_params) = self.lower_closure_binder(binder);
 
@@ -917,6 +922,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             fn_decl,
             body: body_id,
             fn_decl_span: self.lower_span(fn_decl_span),
+            fn_arg_span: Some(self.lower_span(fn_arg_span)),
             movability: generator_option,
         });
 
@@ -973,6 +979,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         decl: &FnDecl,
         body: &Expr,
         fn_decl_span: Span,
+        fn_arg_span: Span,
     ) -> hir::ExprKind<'hir> {
         if let &ClosureBinder::For { span, .. } = binder {
             self.tcx.sess.emit_err(NotSupportedForLifetimeBinderAsyncClosure { span });
@@ -1027,6 +1034,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             fn_decl,
             body,
             fn_decl_span: self.lower_span(fn_decl_span),
+            fn_arg_span: Some(self.lower_span(fn_arg_span)),
             movability: None,
         });
         hir::ExprKind::Closure(c)
diff --git a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
index 1da40d2302e..8d905322e1d 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
@@ -409,6 +409,7 @@ impl<'a> State<'a> {
                 ref fn_decl,
                 ref body,
                 fn_decl_span: _,
+                fn_arg_span: _,
             }) => {
                 self.print_closure_binder(binder);
                 self.print_movability(movability);
diff --git a/compiler/rustc_expand/src/build.rs b/compiler/rustc_expand/src/build.rs
index e17cba1478a..40f5d739f67 100644
--- a/compiler/rustc_expand/src/build.rs
+++ b/compiler/rustc_expand/src/build.rs
@@ -539,6 +539,9 @@ impl<'a> ExtCtxt<'a> {
                 fn_decl,
                 body,
                 fn_decl_span: span,
+                // FIXME(SarthakSingh31): This points to the start of the declaration block and
+                // not the span of the argument block.
+                fn_arg_span: span,
             })),
         )
     }
diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
index 473a04f33a9..34836ab7330 100644
--- a/compiler/rustc_hir/src/hir.rs
+++ b/compiler/rustc_hir/src/hir.rs
@@ -928,7 +928,10 @@ pub struct Closure<'hir> {
     pub bound_generic_params: &'hir [GenericParam<'hir>],
     pub fn_decl: &'hir FnDecl<'hir>,
     pub body: BodyId,
+    /// The span of the declaration block: 'move |...| -> ...'
     pub fn_decl_span: Span,
+    /// The span of the argument block `|...|`
+    pub fn_arg_span: Option<Span>,
     pub movability: Option<Movability>,
 }
 
diff --git a/compiler/rustc_hir/src/intravisit.rs b/compiler/rustc_hir/src/intravisit.rs
index 48db93fde9d..8ced4aa23c0 100644
--- a/compiler/rustc_hir/src/intravisit.rs
+++ b/compiler/rustc_hir/src/intravisit.rs
@@ -740,6 +740,7 @@ pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expression: &'v Expr<'v>)
             body,
             capture_clause: _,
             fn_decl_span: _,
+            fn_arg_span: _,
             movability: _,
         }) => {
             walk_list!(visitor, visit_generic_param, bound_generic_params);
diff --git a/compiler/rustc_hir_pretty/src/lib.rs b/compiler/rustc_hir_pretty/src/lib.rs
index d70ec94f5b6..72f3569fc23 100644
--- a/compiler/rustc_hir_pretty/src/lib.rs
+++ b/compiler/rustc_hir_pretty/src/lib.rs
@@ -1480,6 +1480,7 @@ impl<'a> State<'a> {
                 fn_decl,
                 body,
                 fn_decl_span: _,
+                fn_arg_span: _,
                 movability: _,
                 def_id: _,
             }) => {
diff --git a/compiler/rustc_hir_typeck/src/closure.rs b/compiler/rustc_hir_typeck/src/closure.rs
index 6cf9e23b40b..e75b96da97a 100644
--- a/compiler/rustc_hir_typeck/src/closure.rs
+++ b/compiler/rustc_hir_typeck/src/closure.rs
@@ -457,10 +457,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             .iter()
             .map(|ty| ArgKind::from_expected_ty(*ty, None))
             .collect();
-        let (closure_span, found_args) = match self.get_fn_like_arguments(expr_map_node) {
-            Some((sp, args)) => (Some(sp), args),
-            None => (None, Vec::new()),
-        };
+        let (closure_span, closure_arg_span, found_args) =
+            match self.get_fn_like_arguments(expr_map_node) {
+                Some((sp, arg_sp, args)) => (Some(sp), arg_sp, args),
+                None => (None, None, Vec::new()),
+            };
         let expected_span =
             expected_sig.cause_span.unwrap_or_else(|| self.tcx.def_span(expr_def_id));
         self.report_arg_count_mismatch(
@@ -469,6 +470,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             expected_args,
             found_args,
             true,
+            closure_arg_span,
         )
         .emit();
 
diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index d4456adf201..68be92ad398 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/compiler/rustc_middle/src/hir/map/mod.rs
@@ -1022,7 +1022,7 @@ impl<'hir> Map<'hir> {
                 ..
             }) => {
                 // Ensure that the returned span has the item's SyntaxContext.
-                fn_decl_span.find_ancestor_in_same_ctxt(*span).unwrap_or(*span)
+                fn_decl_span.find_ancestor_inside(*span).unwrap_or(*span)
             }
             _ => self.span_with_body(hir_id),
         };
diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index dc914f5ea64..a76f0f46716 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -2057,7 +2057,7 @@ impl<'a> Parser<'a> {
         };
 
         let capture_clause = self.parse_capture_clause()?;
-        let fn_decl = self.parse_fn_block_decl()?;
+        let (fn_decl, fn_arg_span) = self.parse_fn_block_decl()?;
         let decl_hi = self.prev_token.span;
         let mut body = match fn_decl.output {
             FnRetTy::Default(_) => {
@@ -2098,6 +2098,7 @@ impl<'a> Parser<'a> {
                 fn_decl,
                 body,
                 fn_decl_span: lo.to(decl_hi),
+                fn_arg_span,
             })),
         );
 
@@ -2126,7 +2127,9 @@ impl<'a> Parser<'a> {
     }
 
     /// Parses the `|arg, arg|` header of a closure.
-    fn parse_fn_block_decl(&mut self) -> PResult<'a, P<FnDecl>> {
+    fn parse_fn_block_decl(&mut self) -> PResult<'a, (P<FnDecl>, Span)> {
+        let arg_start = self.token.span.lo();
+
         let inputs = if self.eat(&token::OrOr) {
             Vec::new()
         } else {
@@ -2142,10 +2145,11 @@ impl<'a> Parser<'a> {
             self.expect_or()?;
             args
         };
+        let arg_span = self.prev_token.span.with_lo(arg_start);
         let output =
             self.parse_ret_ty(AllowPlus::Yes, RecoverQPath::Yes, RecoverReturnSign::Yes)?;
 
-        Ok(P(FnDecl { inputs, output }))
+        Ok((P(FnDecl { inputs, output }), arg_span))
     }
 
     /// Parses a parameter in a closure header (e.g., `|arg, arg|`).
diff --git a/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs b/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs
index 59d017545c0..53ce8f368e2 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs
@@ -70,7 +70,7 @@ pub trait InferCtxtExt<'tcx> {
     /// returns a span and `ArgKind` information that describes the
     /// arguments it expects. This can be supplied to
     /// `report_arg_count_mismatch`.
-    fn get_fn_like_arguments(&self, node: Node<'_>) -> Option<(Span, Vec<ArgKind>)>;
+    fn get_fn_like_arguments(&self, node: Node<'_>) -> Option<(Span, Option<Span>, Vec<ArgKind>)>;
 
     /// Reports an error when the number of arguments needed by a
     /// trait match doesn't match the number that the expression
@@ -82,6 +82,7 @@ pub trait InferCtxtExt<'tcx> {
         expected_args: Vec<ArgKind>,
         found_args: Vec<ArgKind>,
         is_closure: bool,
+        closure_pipe_span: Option<Span>,
     ) -> DiagnosticBuilder<'tcx, ErrorGuaranteed>;
 
     /// Checks if the type implements one of `Fn`, `FnMut`, or `FnOnce`
@@ -134,15 +135,16 @@ impl<'tcx> InferCtxtExt<'tcx> for InferCtxt<'tcx> {
     /// returns a span and `ArgKind` information that describes the
     /// arguments it expects. This can be supplied to
     /// `report_arg_count_mismatch`.
-    fn get_fn_like_arguments(&self, node: Node<'_>) -> Option<(Span, Vec<ArgKind>)> {
+    fn get_fn_like_arguments(&self, node: Node<'_>) -> Option<(Span, Option<Span>, Vec<ArgKind>)> {
         let sm = self.tcx.sess.source_map();
         let hir = self.tcx.hir();
         Some(match node {
             Node::Expr(&hir::Expr {
-                kind: hir::ExprKind::Closure(&hir::Closure { body, fn_decl_span, .. }),
+                kind: hir::ExprKind::Closure(&hir::Closure { body, fn_decl_span, fn_arg_span, .. }),
                 ..
             }) => (
                 fn_decl_span,
+                fn_arg_span,
                 hir.body(body)
                     .params
                     .iter()
@@ -173,6 +175,7 @@ impl<'tcx> InferCtxtExt<'tcx> for InferCtxt<'tcx> {
                 kind: hir::TraitItemKind::Fn(ref sig, _), ..
             }) => (
                 sig.span,
+                None,
                 sig.decl
                     .inputs
                     .iter()
@@ -187,7 +190,7 @@ impl<'tcx> InferCtxtExt<'tcx> for InferCtxt<'tcx> {
             ),
             Node::Ctor(ref variant_data) => {
                 let span = variant_data.ctor_hir_id().map_or(DUMMY_SP, |id| hir.span(id));
-                (span, vec![ArgKind::empty(); variant_data.fields().len()])
+                (span, None, vec![ArgKind::empty(); variant_data.fields().len()])
             }
             _ => panic!("non-FnLike node found: {:?}", node),
         })
@@ -203,6 +206,7 @@ impl<'tcx> InferCtxtExt<'tcx> for InferCtxt<'tcx> {
         expected_args: Vec<ArgKind>,
         found_args: Vec<ArgKind>,
         is_closure: bool,
+        closure_arg_span: Option<Span>,
     ) -> DiagnosticBuilder<'tcx, ErrorGuaranteed> {
         let kind = if is_closure { "closure" } else { "function" };
 
@@ -240,24 +244,13 @@ impl<'tcx> InferCtxtExt<'tcx> for InferCtxt<'tcx> {
         if let Some(found_span) = found_span {
             err.span_label(found_span, format!("takes {}", found_str));
 
-            // move |_| { ... }
-            // ^^^^^^^^-- def_span
-            //
-            // move |_| { ... }
-            // ^^^^^-- prefix
-            let prefix_span = self.tcx.sess.source_map().span_until_non_whitespace(found_span);
-            // move |_| { ... }
-            //      ^^^-- pipe_span
-            let pipe_span =
-                if let Some(span) = found_span.trim_start(prefix_span) { span } else { found_span };
-
             // Suggest to take and ignore the arguments with expected_args_length `_`s if
             // found arguments is empty (assume the user just wants to ignore args in this case).
             // For example, if `expected_args_length` is 2, suggest `|_, _|`.
             if found_args.is_empty() && is_closure {
                 let underscores = vec!["_"; expected_args.len()].join(", ");
                 err.span_suggestion_verbose(
-                    pipe_span,
+                    closure_arg_span.unwrap_or(found_span),
                     &format!(
                         "consider changing the closure to take and ignore the expected argument{}",
                         pluralize!(expected_args.len())
@@ -1251,13 +1244,14 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
                         obligation.cause.code(),
                     )
                 } else {
-                    let (closure_span, found) = found_did
+                    let (closure_span, closure_arg_span, found) = found_did
                         .and_then(|did| {
                             let node = self.tcx.hir().get_if_local(did)?;
-                            let (found_span, found) = self.get_fn_like_arguments(node)?;
-                            Some((Some(found_span), found))
+                            let (found_span, closure_arg_span, found) =
+                                self.get_fn_like_arguments(node)?;
+                            Some((Some(found_span), closure_arg_span, found))
                         })
-                        .unwrap_or((found_span, found));
+                        .unwrap_or((found_span, None, found));
 
                     self.report_arg_count_mismatch(
                         span,
@@ -1265,6 +1259,7 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
                         expected,
                         found,
                         found_trait_ty.is_closure(),
+                        closure_arg_span,
                     )
                 }
             }
diff --git a/src/test/ui-fulldeps/pprust-expr-roundtrip.rs b/src/test/ui-fulldeps/pprust-expr-roundtrip.rs
index d6dc179da7f..a93ba87470a 100644
--- a/src/test/ui-fulldeps/pprust-expr-roundtrip.rs
+++ b/src/test/ui-fulldeps/pprust-expr-roundtrip.rs
@@ -126,6 +126,7 @@ fn iter_exprs(depth: usize, f: &mut dyn FnMut(P<Expr>)) {
                         fn_decl: decl.clone(),
                         body: e,
                         fn_decl_span: DUMMY_SP,
+                        fn_arg_span: DUMMY_SP,
                     })))
                 });
             }
diff --git a/src/tools/rustfmt/src/closures.rs b/src/tools/rustfmt/src/closures.rs
index 423c3a997f5..244d4427c56 100644
--- a/src/tools/rustfmt/src/closures.rs
+++ b/src/tools/rustfmt/src/closures.rs
@@ -335,6 +335,7 @@ pub(crate) fn rewrite_last_closure(
             ref fn_decl,
             ref body,
             fn_decl_span: _,
+            fn_arg_span: _,
         } = **closure;
         let body = match body.kind {
             ast::ExprKind::Block(ref block, _)