[XRay] Define the InstrumentationMap type

Summary:
This change implements the instrumentation map loading library which can
understand both YAML-defined instrumentation maps, and ELF 64-bit object
files that have the XRay instrumentation map section. We break it out
into a library on its own to allow for other applications to deal with
the XRay instrumentation map defined in XRay-instrumented binaries.

This type provides both raw access to the logical representation of the
instrumentation map entries as well as higher level functions for
converting a function ID into a function address.

At this point we only support ELF64 binaries and YAML-defined XRay
instrumentation maps. Future changes should extend this to support
32-bit ELF binaries, as well as other binary formats (like MachO).

As part of this change we also migrate all uses of the extraction logic
that used to be defined in tools/llvm-xray/ to use this new type and
interface for loading from files. We also remove the flag from the
`llvm-xray` tool that required users to specify the type of the
instrumentation map file being provided to instead make the library
auto-detect the file type.

Reviewers: dblaikie

Subscribers: mgorny, varno, llvm-commits

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

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

2 files changed, 197 insertions, 0 deletions

diff --git a/lib/XRay/CMakeLists.txt b/lib/XRay/CMakeLists.txt
index 6c1acba79bfa..8d558209d8ee 100644
--- a/lib/XRay/CMakeLists.txt
+++ b/lib/XRay/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_llvm_library(LLVMXRay
+  InstrumentationMap.cpp
   Trace.cpp
 
   ADDITIONAL_HEADER_DIRS
@@ -7,7 +8,9 @@ add_llvm_library(LLVMXRay
 
   DEPENDS
   LLVMSupport
+  LLVMObject
 
   LINK_LIBS
   LLVMSupport
+  LLVMObject
   )
diff --git a/lib/XRay/InstrumentationMap.cpp b/lib/XRay/InstrumentationMap.cpp
new file mode 100644
index 000000000000..f14ad3815780
--- /dev/null
+++ b/lib/XRay/InstrumentationMap.cpp
@@ -0,0 +1,194 @@
+//===- InstrumentationMap.cpp - XRay Instrumentation Map ------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of the InstrumentationMap type for XRay sleds.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef XRAY_INSTRUMENTATIONMAP_H
+#define XRAY_INSTRUMENTATIONMAP_H
+
+#include "llvm/XRay/InstrumentationMap.h"
+
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/DataExtractor.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/XRay/XRayRecord.h"
+#include <system_error>
+
+namespace llvm {
+namespace xray {
+
+Optional<int32_t> InstrumentationMap::getFunctionId(uint64_t Addr) const {
+  auto I = FunctionIds.find(Addr);
+  if (I != FunctionIds.end())
+    return I->second;
+  return None;
+}
+
+Optional<uint64_t> InstrumentationMap::getFunctionAddr(int32_t FuncId) const {
+  auto I = FunctionAddresses.find(FuncId);
+  if (I != FunctionAddresses.end())
+    return I->second;
+  return None;
+}
+
+namespace {
+Error loadELF64(StringRef Filename,
+                object::OwningBinary<object::ObjectFile> &ObjFile,
+                InstrumentationMap::SledContainer &Sleds,
+                InstrumentationMap::FunctionAddressMap &FunctionAddresses,
+                InstrumentationMap::FunctionAddressReverseMap &FunctionIds) {
+  InstrumentationMap Map;
+
+  // Find the section named "xray_instr_map".
+  if (!ObjFile.getBinary()->isELF() ||
+      ObjFile.getBinary()->getArch() != Triple::x86_64)
+    return make_error<StringError>(
+        "File format not supported (only does ELF little endian 64-bit).",
+        std::make_error_code(std::errc::not_supported));
+
+  StringRef Contents = "";
+  const auto &Sections = ObjFile.getBinary()->sections();
+  auto I = find_if(Sections, [&](object::SectionRef Section) {
+    StringRef Name = "";
+    if (Section.getName(Name))
+      return false;
+    return Name == "xray_instr_map";
+  });
+
+  if (I == Sections.end())
+    return make_error<StringError>(
+        "Failed to find XRay instrumentation map.",
+        std::make_error_code(std::errc::executable_format_error));
+
+  if (I->getContents(Contents))
+    return errorCodeToError(
+        std::make_error_code(std::errc::executable_format_error));
+
+  // Copy the instrumentation map data into the Sleds data structure.
+  auto C = Contents.bytes_begin();
+  static constexpr size_t ELF64SledEntrySize = 32;
+
+  if ((C - Contents.bytes_end()) % ELF64SledEntrySize != 0)
+    return make_error<StringError>(
+        Twine("Instrumentation map entries not evenly divisible by size of "
+              "an XRay sled entry in ELF64."),
+        std::make_error_code(std::errc::executable_format_error));
+
+  int32_t FuncId = 1;
+  uint64_t CurFn = 0;
+  for (; C != Contents.bytes_end(); C += ELF64SledEntrySize) {
+    DataExtractor Extractor(
+        StringRef(reinterpret_cast<const char *>(C), ELF64SledEntrySize), true,
+        8);
+    Sleds.push_back({});
+    auto &Entry = Sleds.back();
+    uint32_t OffsetPtr = 0;
+    Entry.Address = Extractor.getU64(&OffsetPtr);
+    Entry.Function = Extractor.getU64(&OffsetPtr);
+    auto Kind = Extractor.getU8(&OffsetPtr);
+    static constexpr SledEntry::FunctionKinds Kinds[] = {
+        SledEntry::FunctionKinds::ENTRY, SledEntry::FunctionKinds::EXIT,
+        SledEntry::FunctionKinds::TAIL,
+    };
+    if (Kind >= sizeof(Kinds))
+      return errorCodeToError(
+          std::make_error_code(std::errc::executable_format_error));
+    Entry.Kind = Kinds[Kind];
+    Entry.AlwaysInstrument = Extractor.getU8(&OffsetPtr) != 0;
+
+    // We do replicate the function id generation scheme implemented in the
+    // XRay runtime.
+    // FIXME: Figure out how to keep this consistent with the XRay runtime.
+    if (CurFn == 0) {
+      CurFn = Entry.Function;
+      FunctionAddresses[FuncId] = Entry.Function;
+      FunctionIds[Entry.Function] = FuncId;
+    }
+    if (Entry.Function != CurFn) {
+      ++FuncId;
+      CurFn = Entry.Function;
+      FunctionAddresses[FuncId] = Entry.Function;
+      FunctionIds[Entry.Function] = FuncId;
+    }
+  }
+  return Error::success();
+}
+
+Error loadYAML(int Fd, size_t FileSize, StringRef Filename,
+               InstrumentationMap::SledContainer &Sleds,
+               InstrumentationMap::FunctionAddressMap &FunctionAddresses,
+               InstrumentationMap::FunctionAddressReverseMap &FunctionIds) {
+  std::error_code EC;
+  sys::fs::mapped_file_region MappedFile(
+      Fd, sys::fs::mapped_file_region::mapmode::readonly, FileSize, 0, EC);
+  if (EC)
+    return make_error<StringError>(
+        Twine("Failed memory-mapping file '") + Filename + "'.", EC);
+
+  std::vector<YAMLXRaySledEntry> YAMLSleds;
+  yaml::Input In(StringRef(MappedFile.data(), MappedFile.size()));
+  In >> YAMLSleds;
+  if (In.error())
+    return make_error<StringError>(
+        Twine("Failed loading YAML document from '") + Filename + "'.",
+        In.error());
+
+  Sleds.reserve(YAMLSleds.size());
+  for (const auto &Y : YAMLSleds) {
+    FunctionAddresses[Y.FuncId] = Y.Function;
+    FunctionIds[Y.Function] = Y.FuncId;
+    Sleds.push_back(
+        SledEntry{Y.Address, Y.Function, Y.Kind, Y.AlwaysInstrument});
+  }
+  return Error::success();
+}
+} // namespace
+
+// FIXME: Create error types that encapsulate a bit more information than what
+// StringError instances contain.
+Expected<InstrumentationMap> loadInstrumentationMap(StringRef Filename) {
+  // At this point we assume the file is an object file -- and if that doesn't
+  // work, we treat it as YAML.
+  // FIXME: Extend to support non-ELF and non-x86_64 binaries.
+
+  InstrumentationMap Map;
+  auto ObjectFileOrError = object::ObjectFile::createObjectFile(Filename);
+  if (!ObjectFileOrError) {
+    auto E = ObjectFileOrError.takeError();
+    // We try to load it as YAML if the ELF load didn't work.
+    int Fd;
+    if (sys::fs::openFileForRead(Filename, Fd))
+      return std::move(E);
+
+    uint64_t FileSize;
+    if (sys::fs::file_size(Filename, FileSize))
+      return std::move(E);
+
+    // If the file is empty, we return the original error.
+    if (FileSize == 0)
+      return std::move(E);
+
+    // From this point on the errors will be only for the YAML parts, so we
+    // consume the errors at this point.
+    consumeError(std::move(E));
+    if (auto E = loadYAML(Fd, FileSize, Filename, Map.Sleds,
+                          Map.FunctionAddresses, Map.FunctionIds))
+      return std::move(E);
+  } else if (auto E = loadELF64(Filename, *ObjectFileOrError, Map.Sleds,
+                                Map.FunctionAddresses, Map.FunctionIds)) {
+    return std::move(E);
+  }
+  return Map;
+}
+}
+}
+
+#endif // XRAY_INSTRUMENTATIONMAP_H