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;
+}