[XRay][compiler-rt] Add function id utilities for XRay

Summary:
This change allows us to provide users and implementers of XRay handlers
a means of converting XRay function id's to addresses. This, in
combination with the facilities provided in D32695, allows users to find
out:

  - How many function id's there are defined in the current binary.
  - Get the address of the function associated with this function id.
  - Patch only specific functions according to their requirements.

While we don't directly provide symbolization support in XRay, having
the function's address lets users determine this information easily
either during runtime, or offline with tools like 'addr2line'.

Reviewers: dblaikie, echristo, pelikan

Subscribers: kpw, llvm-commits

Differential Revision: https://reviews.llvm.org/D32846

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@302210 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 68 insertions, 1 deletions

diff --git a/include/xray/xray_interface.h b/include/xray/xray_interface.h
index 469bfffe5..c90025e38 100644
--- a/include/xray/xray_interface.h
+++ b/include/xray/xray_interface.h
@@ -15,6 +15,7 @@
 #define XRAY_XRAY_INTERFACE_H
 
 #include <cstdint>
+#include <stddef.h>
 
 extern "C" {
 
@@ -86,6 +87,14 @@ extern XRayPatchingStatus __xray_patch_function(int32_t FuncId);
 /// result values.
 extern XRayPatchingStatus __xray_unpatch_function(int32_t FuncId);
 
+/// This function returns the address of the function provided a valid function
+/// id. We return 0 if we encounter any error, even if 0 may be a valid function
+/// address.
+extern uintptr_t __xray_function_address(int32_t FuncId);
+
+/// This function returns the maximum valid function id. Returns 0 if we
+/// encounter errors (when there are no instrumented functions, etc.).
+extern size_t __xray_max_function_id();
 
 }
 
diff --git a/lib/xray/xray_interface.cc b/lib/xray/xray_interface.cc
index 73cca80bc..26f0ab122 100644
--- a/lib/xray/xray_interface.cc
+++ b/lib/xray/xray_interface.cc
@@ -255,7 +255,7 @@ XRayPatchingStatus patchFunction(int32_t FuncId,
 
   // FuncId must be a positive number, less than the number of functions
   // instrumented.
-  if (FuncId <= 0 || static_cast<size_t>(FuncId) >= InstrMap.Functions) {
+  if (FuncId <= 0 || static_cast<size_t>(FuncId) > InstrMap.Functions) {
     Report("Invalid function id provided: %d\n", FuncId);
     return XRayPatchingStatus::FAILED;
   }
@@ -302,3 +302,15 @@ int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t)) {
   return 1;
 }
 int __xray_remove_handler_arg1() { return __xray_set_handler_arg1(nullptr); }
+
+uintptr_t __xray_function_address(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
+  __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+  if (FuncId <= 0 || static_cast<size_t>(FuncId) > XRayInstrMap.Functions)
+    return 0;
+  return XRayInstrMap.SledsIndex[FuncId - 1].Begin->Address;
+}
+
+size_t __xray_max_function_id() XRAY_NEVER_INSTRUMENT {
+  __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+  return XRayInstrMap.Functions;
+}
diff --git a/test/xray/TestCases/Linux/func-id-utils.cc b/test/xray/TestCases/Linux/func-id-utils.cc
new file mode 100644
index 000000000..834e7b499
--- /dev/null
+++ b/test/xray/TestCases/Linux/func-id-utils.cc
@@ -0,0 +1,46 @@
+// Check that we can turn a function id to a function address, and also get the
+// maximum function id for the current binary.
+//
+// RUN: %clangxx_xray -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=false xray_naive_log=false" %run %t | FileCheck %s
+
+#include "xray/xray_interface.h"
+#include <algorithm>
+#include <cstdio>
+#include <set>
+#include <iterator>
+
+[[clang::xray_always_instrument]] void bar(){
+    // do nothing!
+}
+
+    [[clang::xray_always_instrument]] void foo() {
+  bar();
+}
+
+[[clang::xray_always_instrument]] int main(int argc, char *argv[]) {
+  printf("max function id: %zu\n", __xray_max_function_id());
+  // CHECK: max function id: [[MAX:.*]]
+
+  std::set<void *> must_be_instrumented;
+  must_be_instrumented.insert(reinterpret_cast<void*>(&foo));
+  must_be_instrumented.insert(reinterpret_cast<void*>(&bar));
+  printf("addresses:\n");
+  std::set<void *> all_instrumented;
+  for (auto i = __xray_max_function_id(); i != 0; --i) {
+    auto addr = __xray_function_address(i);
+    printf("#%lu -> @%04lx\n", i, addr);
+    all_instrumented.insert(reinterpret_cast<void *>(addr));
+  }
+
+  // CHECK-LABEL: addresses:
+  // CHECK: #[[MAX]] -> @[[ADDR:.*]]
+  // CHECK-NOT: #0 -> @{{.*}}
+  std::set<void *> common;
+
+  std::set_intersection(all_instrumented.begin(), all_instrumented.end(),
+                        must_be_instrumented.begin(),
+                        must_be_instrumented.end(),
+                        std::inserter(common, common.begin()));
+  return common == must_be_instrumented ? 0 : 1;
+}