Track span of function in method calls, and use this in #[track_caller]

Fixes #69977

When we parse a chain of method calls like `foo.a().b().c()`, each
`MethodCallExpr` gets assigned a span that starts at the beginning of
the call chain (`foo`). While this is useful for diagnostics, it means
that `Location::caller` will return the same location for every call
in a call chain.

This PR makes us separately record the span of the function name and
arguments for a method call (e.g. `b()` in `foo.a().b().c()`). This
`Span` is passed through HIR lowering and MIR building to
`TerminatorKind::Call`, where it is used in preference to
`Terminator.source_info.span` when determining `Location::caller`.

This new span is also useful for diagnostics where we want to emphasize
a particular method call - for an example, see
#72389 (comment)

src/librustc_ast/ast.rs

@@ -1165,7 +1165,9 @@ pub enum ExprKind {
     /// and the remaining elements are the rest of the arguments.
     /// Thus, `x.foo::<Bar, Baz>(a, b, c, d)` is represented as
     /// `ExprKind::MethodCall(PathSegment { foo, [Bar, Baz] }, [x, a, b, c, d])`.
-    MethodCall(PathSegment, Vec<P<Expr>>),
+    /// The `Span` is the span of the function, without the dot and receiver
+    /// (e.g. `foo(a, b)` in `x.foo(a, b)`
+    MethodCall(PathSegment, Vec<P<Expr>>, Span),
     /// A tuple (e.g., `(a, b, c, d)`).
     Tup(Vec<P<Expr>>),
     /// A binary operation (e.g., `a + b`, `a * b`).

src/librustc_ast/mut_visit.rs

@@ -1111,11 +1111,12 @@ pub fn noop_visit_expr<T: MutVisitor>(
             vis.visit_expr(f);
             visit_exprs(args, vis);
         }
-        ExprKind::MethodCall(PathSegment { ident, id, args }, exprs) => {
+        ExprKind::MethodCall(PathSegment { ident, id, args }, exprs, span) => {
             vis.visit_ident(ident);
             vis.visit_id(id);
             visit_opt(args, |args| vis.visit_generic_args(args));
             visit_exprs(exprs, vis);
+            vis.visit_span(span);
         }
         ExprKind::Binary(_binop, lhs, rhs) => {
             vis.visit_expr(lhs);

src/librustc_ast/util/parser.rs

@@ -394,7 +394,7 @@ pub fn contains_exterior_struct_lit(value: &ast::Expr) -> bool {
             contains_exterior_struct_lit(&x)
         }
 
-        ast::ExprKind::MethodCall(.., ref exprs) => {
+        ast::ExprKind::MethodCall(.., ref exprs, _) => {
             // X { y: 1 }.bar(...)
             contains_exterior_struct_lit(&exprs[0])
         }

src/librustc_ast/visit.rs

@@ -726,7 +726,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
             visitor.visit_expr(callee_expression);
             walk_list!(visitor, visit_expr, arguments);
         }
-        ExprKind::MethodCall(ref segment, ref arguments) => {
+        ExprKind::MethodCall(ref segment, ref arguments, _span) => {
             visitor.visit_path_segment(expression.span, segment);
             walk_list!(visitor, visit_expr, arguments);
         }

src/librustc_ast_lowering/expr.rs

@@ -39,7 +39,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     let f = self.lower_expr(f);
                     hir::ExprKind::Call(f, self.lower_exprs(args))
                 }
-                ExprKind::MethodCall(ref seg, ref args) => {
+                ExprKind::MethodCall(ref seg, ref args, span) => {
                     let hir_seg = self.arena.alloc(self.lower_path_segment(
                         e.span,
                         seg,
@@ -50,7 +50,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         None,
                     ));
                     let args = self.lower_exprs(args);
-                    hir::ExprKind::MethodCall(hir_seg, seg.ident.span, args)
+                    hir::ExprKind::MethodCall(hir_seg, seg.ident.span, args, span)
                 }
                 ExprKind::Binary(binop, ref lhs, ref rhs) => {
                     let binop = self.lower_binop(binop);

src/librustc_ast_pretty/pprust.rs

@@ -1818,7 +1818,7 @@ impl<'a> State<'a> {
             ast::ExprKind::Call(ref func, ref args) => {
                 self.print_expr_call(func, &args[..]);
             }
-            ast::ExprKind::MethodCall(ref segment, ref args) => {
+            ast::ExprKind::MethodCall(ref segment, ref args, _) => {
                 self.print_expr_method_call(segment, &args[..]);
             }
             ast::ExprKind::Binary(op, ref lhs, ref rhs) => {

src/librustc_codegen_ssa/mir/block.rs

@@ -530,6 +530,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         args: &Vec<mir::Operand<'tcx>>,
         destination: &Option<(mir::Place<'tcx>, mir::BasicBlock)>,
         cleanup: Option<mir::BasicBlock>,
+        fn_span: Span,
     ) {
         let span = terminator.source_info.span;
         // Create the callee. This is a fn ptr or zero-sized and hence a kind of scalar.
@@ -634,7 +635,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
         if intrinsic == Some("caller_location") {
             if let Some((_, target)) = destination.as_ref() {
-                let location = self.get_caller_location(&mut bx, span);
+                let location = self.get_caller_location(&mut bx, fn_span);
 
                 if let ReturnDest::IndirectOperand(tmp, _) = ret_dest {
                     location.val.store(&mut bx, tmp);
@@ -798,7 +799,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 args.len() + 1,
                 "#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
             );
-            let location = self.get_caller_location(&mut bx, span);
+            let location = self.get_caller_location(&mut bx, fn_span);
+            debug!(
+                "codegen_call_terminator({:?}): location={:?} (fn_span {:?})",
+                terminator, location, fn_span
+            );
+
             let last_arg = fn_abi.args.last().unwrap();
             self.codegen_argument(&mut bx, location, &mut llargs, last_arg);
         }
@@ -1016,6 +1022,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 ref destination,
                 cleanup,
                 from_hir_call: _,
+                fn_span,
             } => {
                 self.codegen_call_terminator(
                     helper,
@@ -1025,6 +1032,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     args,
                     destination,
                     cleanup,
+                    fn_span,
                 );
             }
             mir::TerminatorKind::GeneratorDrop | mir::TerminatorKind::Yield { .. } => {

src/librustc_expand/build.rs

@@ -272,7 +272,7 @@ impl<'a> ExtCtxt<'a> {
     ) -> P<ast::Expr> {
         args.insert(0, expr);
         let segment = ast::PathSegment::from_ident(ident.with_span_pos(span));
-        self.expr(span, ast::ExprKind::MethodCall(segment, args))
+        self.expr(span, ast::ExprKind::MethodCall(segment, args, span))
     }
     pub fn expr_block(&self, b: P<ast::Block>) -> P<ast::Expr> {
         self.expr(b.span, ast::ExprKind::Block(b, None))

src/librustc_hir/hir.rs

@@ -1371,7 +1371,7 @@ pub struct Expr<'hir> {
 
 // `Expr` is used a lot. Make sure it doesn't unintentionally get bigger.
 #[cfg(target_arch = "x86_64")]
-rustc_data_structures::static_assert_size!(Expr<'static>, 64);
+rustc_data_structures::static_assert_size!(Expr<'static>, 72);
 
 impl Expr<'_> {
     pub fn precedence(&self) -> ExprPrecedence {
@@ -1568,12 +1568,14 @@ pub enum ExprKind<'hir> {
     /// and the remaining elements are the rest of the arguments.
     /// Thus, `x.foo::<Bar, Baz>(a, b, c, d)` is represented as
     /// `ExprKind::MethodCall(PathSegment { foo, [Bar, Baz] }, [x, a, b, c, d])`.
+    /// The final `Span` represents the span of the function and arguments
+    /// (e.g. `foo::<Bar, Baz>(a, b, c, d)` in `x.foo::<Bar, Baz>(a, b, c, d)`
     ///
     /// To resolve the called method to a `DefId`, call [`type_dependent_def_id`] with
     /// the `hir_id` of the `MethodCall` node itself.
     ///
     /// [`type_dependent_def_id`]: ../ty/struct.TypeckTables.html#method.type_dependent_def_id
-    MethodCall(&'hir PathSegment<'hir>, Span, &'hir [Expr<'hir>]),
+    MethodCall(&'hir PathSegment<'hir>, Span, &'hir [Expr<'hir>], Span),
     /// A tuple (e.g., `(a, b, c, d)`).
     Tup(&'hir [Expr<'hir>]),
     /// A binary operation (e.g., `a + b`, `a * b`).

src/librustc_hir/intravisit.rs

@@ -1090,7 +1090,7 @@ pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expression: &'v Expr<'v>)
             visitor.visit_expr(callee_expression);
             walk_list!(visitor, visit_expr, arguments);
         }
-        ExprKind::MethodCall(ref segment, _, arguments) => {
+        ExprKind::MethodCall(ref segment, _, arguments, _) => {
             visitor.visit_path_segment(expression.span, segment);
             walk_list!(visitor, visit_expr, arguments);
         }

src/librustc_hir_pretty/lib.rs

@@ -1286,7 +1286,7 @@ impl<'a> State<'a> {
             hir::ExprKind::Call(ref func, ref args) => {
                 self.print_expr_call(&func, args);
             }
-            hir::ExprKind::MethodCall(ref segment, _, ref args) => {
+            hir::ExprKind::MethodCall(ref segment, _, ref args, _) => {
                 self.print_expr_method_call(segment, args);
             }
             hir::ExprKind::Binary(op, ref lhs, ref rhs) => {
@@ -2469,7 +2469,7 @@ fn contains_exterior_struct_lit(value: &hir::Expr<'_>) -> bool {
             contains_exterior_struct_lit(&x)
         }
 
-        hir::ExprKind::MethodCall(.., ref exprs) => {
+        hir::ExprKind::MethodCall(.., ref exprs, _) => {
             // `X { y: 1 }.bar(...)`
             contains_exterior_struct_lit(&exprs[0])
         }

src/librustc_infer/infer/error_reporting/need_type_info.rs

@@ -107,7 +107,7 @@ impl<'a, 'tcx> Visitor<'tcx> for FindHirNodeVisitor<'a, 'tcx> {
     }
 
     fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
-        if let ExprKind::MethodCall(_, call_span, exprs) = expr.kind {
+        if let ExprKind::MethodCall(_, call_span, exprs, _) = expr.kind {
             if call_span == self.target_span
                 && Some(self.target)
                     == self.infcx.in_progress_tables.and_then(|tables| {
@@ -294,7 +294,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
             // 3 |     let _ = x.sum() as f64;
             //   |               ^^^ cannot infer type for `S`
             span
-        } else if let Some(ExprKind::MethodCall(_, call_span, _)) =
+        } else if let Some(ExprKind::MethodCall(_, call_span, _, _)) =
             local_visitor.found_method_call.map(|e| &e.kind)
         {
             // Point at the call instead of the whole expression:

src/librustc_lint/array_into_iter.rs

@@ -24,7 +24,7 @@ declare_lint_pass!(
 impl<'a, 'tcx> LateLintPass<'a, 'tcx> for ArrayIntoIter {
     fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, expr: &'tcx hir::Expr<'tcx>) {
         // We only care about method call expressions.
-        if let hir::ExprKind::MethodCall(call, span, args) = &expr.kind {
+        if let hir::ExprKind::MethodCall(call, span, args, _) = &expr.kind {
             if call.ident.name != sym::into_iter {
                 return;
             }

src/librustc_lint/builtin.rs

@@ -1899,7 +1899,7 @@ impl<'a, 'tcx> LateLintPass<'a, 'tcx> for InvalidValue {
                         }
                     }
                 }
-            } else if let hir::ExprKind::MethodCall(_, _, ref args) = expr.kind {
+            } else if let hir::ExprKind::MethodCall(_, _, ref args, _) = expr.kind {
                 // Find problematic calls to `MaybeUninit::assume_init`.
                 let def_id = cx.tables.type_dependent_def_id(expr.hir_id)?;
                 if cx.tcx.is_diagnostic_item(sym::assume_init, def_id) {

src/librustc_lint/unused.rs

@@ -526,7 +526,7 @@ trait UnusedDelimLint {
                 let (args_to_check, ctx) = match *call_or_other {
                     Call(_, ref args) => (&args[..], UnusedDelimsCtx::FunctionArg),
                     // first "argument" is self (which sometimes needs delims)
-                    MethodCall(_, ref args) => (&args[1..], UnusedDelimsCtx::MethodArg),
+                    MethodCall(_, ref args, _) => (&args[1..], UnusedDelimsCtx::MethodArg),
                     // actual catch-all arm
                     _ => {
                         return;

src/librustc_middle/mir/mod.rs

@@ -1131,6 +1131,7 @@ pub enum TerminatorKind<'tcx> {
         /// `true` if this is from a call in HIR rather than from an overloaded
         /// operator. True for overloaded function call.
         from_hir_call: bool,
+        fn_span: Span,
     },
 
     /// Jump to the target if the condition has the expected value,

src/librustc_middle/mir/type_foldable.rs

@@ -42,7 +42,7 @@ impl<'tcx> TypeFoldable<'tcx> for Terminator<'tcx> {
                 resume_arg: resume_arg.fold_with(folder),
                 drop,
             },
-            Call { ref func, ref args, ref destination, cleanup, from_hir_call } => {
+            Call { ref func, ref args, ref destination, cleanup, from_hir_call, fn_span } => {
                 let dest =
                     destination.as_ref().map(|&(ref loc, dest)| (loc.fold_with(folder), dest));
 
@@ -52,6 +52,7 @@ impl<'tcx> TypeFoldable<'tcx> for Terminator<'tcx> {
                     destination: dest,
                     cleanup,
                     from_hir_call,
+                    fn_span,
                 }
             }
             Assert { ref cond, expected, ref msg, target, cleanup } => {

src/librustc_middle/mir/visit.rs

@@ -492,6 +492,7 @@ macro_rules! make_mir_visitor {
                         destination,
                         cleanup: _,
                         from_hir_call: _,
+                        fn_span: _
                     } => {
                         self.visit_operand(func, source_location);
                         for arg in args {

src/librustc_mir/borrow_check/invalidation.rs

@@ -142,6 +142,7 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                 destination,
                 cleanup: _,
                 from_hir_call: _,
+                fn_span: _,
             } => {
                 self.consume_operand(location, func);
                 for arg in args {

src/librustc_mir/borrow_check/mod.rs

@@ -699,6 +699,7 @@ impl<'cx, 'tcx> dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtxt<'cx, 'tc
                 ref destination,
                 cleanup: _,
                 from_hir_call: _,
+                fn_span: _,
             } => {
                 self.consume_operand(loc, (func, span), flow_state);
                 for arg in args {

src/librustc_mir/dataflow/framework/direction.rs

@@ -467,7 +467,7 @@ impl Direction for Forward {
                 propagate(target, exit_state);
             }
 
-            Call { cleanup, destination, ref func, ref args, from_hir_call: _ } => {
+            Call { cleanup, destination, ref func, ref args, from_hir_call: _, fn_span: _ } => {
                 if let Some(unwind) = cleanup {
                     if dead_unwinds.map_or(true, |dead| !dead.contains(bb)) {
                         propagate(unwind, exit_state);

src/librustc_mir/dataflow/move_paths/builder.rs

@@ -401,6 +401,7 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
                 ref destination,
                 cleanup: _,
                 from_hir_call: _,
+                fn_span: _,
             } => {
                 self.gather_operand(func);
                 for arg in args {

src/librustc_mir/interpret/intrinsics/caller_location.rs

@@ -1,6 +1,7 @@
 use std::convert::TryFrom;
 
 use rustc_hir::lang_items::PanicLocationLangItem;
+use rustc_middle::mir::TerminatorKind;
 use rustc_middle::ty::subst::Subst;
 use rustc_span::{Span, Symbol};
 use rustc_target::abi::LayoutOf;
@@ -14,19 +15,32 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     /// Walks up the callstack from the intrinsic's callsite, searching for the first callsite in a
     /// frame which is not `#[track_caller]`.
     crate fn find_closest_untracked_caller_location(&self) -> Span {
-        self.stack()
+        let frame = self
+            .stack()
             .iter()
             .rev()
             // Find first non-`#[track_caller]` frame.
-            .find(|frame| !frame.instance.def.requires_caller_location(*self.tcx))
+            .find(|frame| {
+                debug!(
+                    "find_closest_untracked_caller_location: checking frame {:?}",
+                    frame.instance
+                );
+                !frame.instance.def.requires_caller_location(*self.tcx)
+            })
             // Assert that there is always such a frame.
-            .unwrap()
-            .current_source_info()
-            // Assert that the frame we look at is actually executing code currently
-            // (`current_source_info` is None when we are unwinding and the frame does
-            // not require cleanup).
-            .unwrap()
-            .span
+            .unwrap();
+        let loc = frame.loc.unwrap();
+        let block = &frame.body.basic_blocks()[loc.block];
+        assert_eq!(block.statements.len(), loc.statement_index);
+        debug!(
+            "find_closest_untracked_caller_location:: got terminator {:?} ({:?})",
+            block.terminator(),
+            block.terminator().kind
+        );
+        if let TerminatorKind::Call { fn_span, .. } = block.terminator().kind {
+            return fn_span;
+        }
+        unreachable!();
     }
 
     /// Allocate a `const core::panic::Location` with the provided filename and line/column numbers.

src/librustc_mir/interpret/terminator.rs

@@ -56,6 +56,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 destination,
                 ref cleanup,
                 from_hir_call: _from_hir_call,
+                fn_span: _,
             } => {
                 let old_stack = self.frame_idx();
                 let old_loc = self.frame().loc;

src/librustc_mir/shim.rs

@@ -460,6 +460,7 @@ impl CloneShimBuilder<'tcx> {
                 destination: Some((dest, next)),
                 cleanup: Some(cleanup),
                 from_hir_call: true,
+                fn_span: self.span,
             },
             false,
         );
@@ -788,6 +789,7 @@ fn build_call_shim<'tcx>(
                 None
             },
             from_hir_call: true,
+            fn_span: span,
         },
         false,
     );

src/librustc_mir/transform/promote_consts.rs

@@ -909,7 +909,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
             };
 
             match terminator.kind {
-                TerminatorKind::Call { mut func, mut args, from_hir_call, .. } => {
+                TerminatorKind::Call { mut func, mut args, from_hir_call, fn_span, .. } => {
                     self.visit_operand(&mut func, loc);
                     for arg in &mut args {
                         self.visit_operand(arg, loc);
@@ -925,6 +925,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
                             cleanup: None,
                             destination: Some((Place::from(new_temp), new_target)),
                             from_hir_call,
+                            fn_span,
                         },
                         ..terminator
                     };

src/librustc_mir/transform/qualify_min_const_fn.rs

@@ -368,7 +368,14 @@ fn check_terminator(
             Err((span, "const fn generators are unstable".into()))
         }
 
-        TerminatorKind::Call { func, args, from_hir_call: _, destination: _, cleanup: _ } => {
+        TerminatorKind::Call {
+            func,
+            args,
+            from_hir_call: _,
+            destination: _,
+            cleanup: _,
+            fn_span: _,
+        } => {
             let fn_ty = func.ty(body, tcx);
             if let ty::FnDef(def_id, _) = fn_ty.kind {
                 if !crate::const_eval::is_min_const_fn(tcx, def_id) {