From 207f5208273667c5c72397fdf99390e594a3cdf1 Mon Sep 17 00:00:00 2001
From: Adam Perry <adam.n.perry@gmail.com>
Date: Thu, 7 Nov 2019 06:04:14 -0800
Subject: [PATCH] Pass a location to #[track_caller] functions in
 codegen_call_terminator.

---
 src/librustc_codegen_ssa/mir/block.rs | 12 ++++++++++++
 src/librustc_codegen_ssa/mir/mod.rs   | 21 +++++++++++++++++----
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/librustc_codegen_ssa/mir/block.rs b/src/librustc_codegen_ssa/mir/block.rs
index 4f47d84047d0..b33cca2ee8da 100644
--- a/src/librustc_codegen_ssa/mir/block.rs
+++ b/src/librustc_codegen_ssa/mir/block.rs
@@ -770,6 +770,18 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 &fn_abi.args[first_args.len()..])
         }
 
+        let needs_location =
+            instance.map(|i| i.def.requires_caller_location(self.cx.tcx())).unwrap_or_default();
+        if needs_location {
+            assert_eq!(
+                fn_abi.args.len(), 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 last_arg = &fn_abi.args.last().unwrap();
+            self.codegen_argument(&mut bx, location, &mut llargs, last_arg);
+        }
+
         let fn_ptr = match (llfn, instance) {
             (Some(llfn), _) => llfn,
             (None, Some(instance)) => bx.get_fn_addr(instance),
diff --git a/src/librustc_codegen_ssa/mir/mod.rs b/src/librustc_codegen_ssa/mir/mod.rs
index a6333ef23f5e..6be6e13cc1d3 100644
--- a/src/librustc_codegen_ssa/mir/mod.rs
+++ b/src/librustc_codegen_ssa/mir/mod.rs
@@ -182,7 +182,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     // Allocate variable and temp allocas
     fx.locals = {
-        let args = arg_local_refs(&mut bx, &fx, &memory_locals);
+        let args = arg_local_refs(&mut bx, &mut fx, &memory_locals);
 
         let mut allocate_local = |local| {
             let decl = &mir_body.local_decls[local];
@@ -324,14 +324,14 @@ fn create_funclets<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 /// indirect.
 fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     bx: &mut Bx,
-    fx: &FunctionCx<'a, 'tcx, Bx>,
+    fx: &mut FunctionCx<'a, 'tcx, Bx>,
     memory_locals: &BitSet<mir::Local>,
 ) -> Vec<LocalRef<'tcx, Bx::Value>> {
     let mir = fx.mir;
     let mut idx = 0;
     let mut llarg_idx = fx.fn_abi.ret.is_indirect() as usize;
 
-    mir.args_iter().enumerate().map(|(arg_index, local)| {
+    let args = mir.args_iter().enumerate().map(|(arg_index, local)| {
         let arg_decl = &mir.local_decls[local];
 
         if Some(local) == mir.spread_arg {
@@ -427,7 +427,20 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
             bx.store_fn_arg(arg, &mut llarg_idx, tmp);
             LocalRef::Place(tmp)
         }
-    }).collect()
+    }).collect::<Vec<_>>();
+
+    if fx.instance.def.requires_caller_location(bx.tcx()) {
+        assert_eq!(
+            fx.fn_abi.args.len(), args.len() + 1,
+            "#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
+        );
+        let arg = &fx.fn_abi.args.last().unwrap();
+        let place = PlaceRef::alloca(bx, arg.layout);
+        bx.store_fn_arg(arg, &mut llarg_idx, place);
+        fx.caller_location = Some(place);
+    }
+
+    args
 }
 
 mod analyze;