From 31a5066e0b5f0e7e79b6cc04ae09166de0352c63 Mon Sep 17 00:00:00 2001
From: Jun Wu <quark@fb.com>
Date: Sun, 30 Dec 2018 11:59:03 -0800
Subject: [PATCH] Add `-Z instrument-mcount`

This flag inserts `mcount` function call to the beginning of every function
after inline processing. So tracing tools like uftrace [1] (or ftrace for
Linux kernel modules) have a chance to examine function calls.

It is similar to the `-pg` flag provided by gcc or clang, but without
generating a `__gmon_start__` function for executables. If a program
runs without being traced, no `gmon.out` will be written to disk.

Under the hood, it simply adds `"instrument-function-entry-inlined"="mcount"`
attribute to every function. The `post-inline-ee-instrument` LLVM pass does
the actual job.

[1]: https://github.com/namhyung/uftrace
---
 src/librustc/session/config.rs          |  2 ++
 src/librustc/session/mod.rs             |  9 ++++++++-
 src/librustc_codegen_llvm/attributes.rs | 13 +++++++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/librustc/session/config.rs b/src/librustc/session/config.rs
index e58c9d75a9db2..33409f9b4a74f 100644
--- a/src/librustc/session/config.rs
+++ b/src/librustc/session/config.rs
@@ -1227,6 +1227,8 @@ options! {DebuggingOptions, DebuggingSetter, basic_debugging_options,
         "verify incr. comp. hashes of green query instances"),
     incremental_ignore_spans: bool = (false, parse_bool, [UNTRACKED],
         "ignore spans during ICH computation -- used for testing"),
+    instrument_mcount: bool = (false, parse_bool, [TRACKED],
+        "insert function instrument code for mcount-based tracing"),
     dump_dep_graph: bool = (false, parse_bool, [UNTRACKED],
         "dump the dependency graph to $RUST_DEP_GRAPH (default: /tmp/dep_graph.gv)"),
     query_dep_graph: bool = (false, parse_bool, [UNTRACKED],
diff --git a/src/librustc/session/mod.rs b/src/librustc/session/mod.rs
index 4fc9b87ceef20..5c977d5969e16 100644
--- a/src/librustc/session/mod.rs
+++ b/src/librustc/session/mod.rs
@@ -503,6 +503,9 @@ impl Session {
     pub fn profile_queries_and_keys(&self) -> bool {
         self.opts.debugging_opts.profile_queries_and_keys
     }
+    pub fn instrument_mcount(&self) -> bool {
+        self.opts.debugging_opts.instrument_mcount
+    }
     pub fn count_llvm_insns(&self) -> bool {
         self.opts.debugging_opts.count_llvm_insns
     }
@@ -667,7 +670,11 @@ impl Session {
     }
 
     pub fn must_not_eliminate_frame_pointers(&self) -> bool {
-        if let Some(x) = self.opts.cg.force_frame_pointers {
+        // "mcount" function relies on stack pointer.
+        // See https://sourceware.org/binutils/docs/gprof/Implementation.html
+        if self.instrument_mcount() {
+            true
+        } else if let Some(x) = self.opts.cg.force_frame_pointers {
             x
         } else {
             !self.target.target.options.eliminate_frame_pointer
diff --git a/src/librustc_codegen_llvm/attributes.rs b/src/librustc_codegen_llvm/attributes.rs
index acd1c4afd2e63..226b03c99c043 100644
--- a/src/librustc_codegen_llvm/attributes.rs
+++ b/src/librustc_codegen_llvm/attributes.rs
@@ -78,6 +78,18 @@ pub fn set_frame_pointer_elimination(cx: &CodegenCx<'ll, '_>, llfn: &'ll Value)
     }
 }
 
+/// Tell LLVM what instrument function to insert.
+#[inline]
+pub fn set_instrument_function(cx: &CodegenCx<'ll, '_>, llfn: &'ll Value) {
+    if cx.sess().instrument_mcount() {
+        // Similar to `clang -pg` behavior. Handled by the
+        // `post-inline-ee-instrument` LLVM pass.
+        llvm::AddFunctionAttrStringValue(
+            llfn, llvm::AttributePlace::Function,
+            const_cstr!("instrument-function-entry-inlined"), const_cstr!("mcount"));
+    }
+}
+
 pub fn set_probestack(cx: &CodegenCx<'ll, '_>, llfn: &'ll Value) {
     // Only use stack probes if the target specification indicates that we
     // should be using stack probes
@@ -174,6 +186,7 @@ pub fn from_fn_attrs(
     }
 
     set_frame_pointer_elimination(cx, llfn);
+    set_instrument_function(cx, llfn);
     set_probestack(cx, llfn);
 
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::COLD) {