[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

llvm-svn: 293721

1 files changed, 19 insertions, 225 deletions

diff --git a/llvm/tools/llvm-xray/xray-extract.cc b/llvm/tools/llvm-xray/xray-extract.cc
index 49ecd742113..bf1f696de55 100644
--- a/llvm/tools/llvm-xray/xray-extract.cc
+++ b/llvm/tools/llvm-xray/xray-extract.cc
@@ -16,10 +16,7 @@
 #include <type_traits>
 #include <utility>
 
-#include "xray-extract.h"
-
 #include "xray-registry.h"
-#include "xray-sleds.h"
 #include "llvm/Object/ELF.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/CommandLine.h"
@@ -28,8 +25,8 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Format.h"
-#include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/XRay/InstrumentationMap.h"
 
 using namespace llvm;
 using namespace llvm::xray;
@@ -49,243 +46,40 @@ static cl::alias ExtractOutput2("o", cl::aliasopt(ExtractOutput),
                                 cl::desc("Alias for -output"),
                                 cl::sub(Extract));
 
-struct YAMLXRaySledEntry {
-  int32_t FuncId;
-  Hex64 Address;
-  Hex64 Function;
-  SledEntry::FunctionKinds Kind;
-  bool AlwaysInstrument;
-};
-
-namespace llvm {
-namespace yaml {
-
-template <> struct ScalarEnumerationTraits<SledEntry::FunctionKinds> {
-  static void enumeration(IO &IO, SledEntry::FunctionKinds &Kind) {
-    IO.enumCase(Kind, "function-enter", SledEntry::FunctionKinds::ENTRY);
-    IO.enumCase(Kind, "function-exit", SledEntry::FunctionKinds::EXIT);
-    IO.enumCase(Kind, "tail-exit", SledEntry::FunctionKinds::TAIL);
-  }
-};
-
-template <> struct MappingTraits<YAMLXRaySledEntry> {
-  static void mapping(IO &IO, YAMLXRaySledEntry &Entry) {
-    IO.mapRequired("id", Entry.FuncId);
-    IO.mapRequired("address", Entry.Address);
-    IO.mapRequired("function", Entry.Function);
-    IO.mapRequired("kind", Entry.Kind);
-    IO.mapRequired("always-instrument", Entry.AlwaysInstrument);
-  }
-
-  static constexpr bool flow = true;
-};
-}
-}
-
-LLVM_YAML_IS_SEQUENCE_VECTOR(YAMLXRaySledEntry)
-
 namespace {
 
-llvm::Error LoadBinaryInstrELF(
-    StringRef Filename, std::deque<SledEntry> &OutputSleds,
-    InstrumentationMapExtractor::FunctionAddressMap &InstrMap,
-    InstrumentationMapExtractor::FunctionAddressReverseMap &FunctionIds) {
-  auto ObjectFile = object::ObjectFile::createObjectFile(Filename);
-
-  if (!ObjectFile)
-    return ObjectFile.takeError();
-
-  // FIXME: Maybe support other ELF formats. For now, 64-bit Little Endian only.
-  if (!ObjectFile->getBinary()->isELF())
-    return make_error<StringError>(
-        "File format not supported (only does ELF).",
-        std::make_error_code(std::errc::not_supported));
-  if (ObjectFile->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));
-
-  // Find the section named "xray_instr_map".
-  StringRef Contents = "";
-  const auto &Sections = ObjectFile->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::not_supported));
-  if (I->getContents(Contents))
-    return make_error<StringError>(
-        "Failed to get contents of 'xray_instr_map' section.",
-        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>(
-        "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;
-  std::deque<SledEntry> Sleds;
-  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);
-    switch (Kind) {
-    case 0: // ENTRY
-      Entry.Kind = SledEntry::FunctionKinds::ENTRY;
-      break;
-    case 1: // EXIT
-      Entry.Kind = SledEntry::FunctionKinds::EXIT;
-      break;
-    case 2: // TAIL
-      Entry.Kind = SledEntry::FunctionKinds::TAIL;
-      break;
-    default:
-      return make_error<StringError>(
-          Twine("Encountered unknown sled type ") + "'" + Twine(int32_t{Kind}) +
-              "'.",
-          std::make_error_code(std::errc::executable_format_error));
-    }
-    Entry.AlwaysInstrument = Extractor.getU8(&OffsetPtr) != 0;
-
-    // We replicate the function id generation scheme implemented in the runtime
-    // here. Ideally we should be able to break it out, or output this map from
-    // the runtime, but that's a design point we can discuss later on. For now,
-    // we replicate the logic and move on.
-    if (CurFn == 0) {
-      CurFn = Entry.Function;
-      InstrMap[FuncId] = Entry.Function;
-      FunctionIds[Entry.Function] = FuncId;
-    }
-    if (Entry.Function != CurFn) {
-      ++FuncId;
-      CurFn = Entry.Function;
-      InstrMap[FuncId] = Entry.Function;
-      FunctionIds[Entry.Function] = FuncId;
-    }
-  }
-  OutputSleds = std::move(Sleds);
-  return llvm::Error::success();
-}
-
-Error LoadYAMLInstrMap(
-    StringRef Filename, std::deque<SledEntry> &Sleds,
-    InstrumentationMapExtractor::FunctionAddressMap &InstrMap,
-    InstrumentationMapExtractor::FunctionAddressReverseMap &FunctionIds) {
-  int Fd;
-  if (auto EC = sys::fs::openFileForRead(Filename, Fd))
-    return make_error<StringError>(
-        Twine("Failed opening file '") + Filename + "' for reading.", EC);
-
-  uint64_t FileSize;
-  if (auto EC = sys::fs::file_size(Filename, FileSize))
-    return make_error<StringError>(
-        Twine("Failed getting size of file '") + Filename + "'.", EC);
-
-  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;
-  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());
-
-  for (const auto &Y : YAMLSleds) {
-    InstrMap[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
-
-InstrumentationMapExtractor::InstrumentationMapExtractor(std::string Filename,
-                                                         InputFormats Format,
-                                                         Error &EC) {
-  ErrorAsOutParameter ErrAsOutputParam(&EC);
-  if (Filename.empty()) {
-    EC = Error::success();
-    return;
-  }
-  switch (Format) {
-  case InputFormats::ELF: {
-    EC = handleErrors(
-        LoadBinaryInstrELF(Filename, Sleds, FunctionAddresses, FunctionIds),
-        [&](std::unique_ptr<ErrorInfoBase> E) {
-          return joinErrors(
-              make_error<StringError>(
-                  Twine("Cannot extract instrumentation map from '") +
-                      Filename + "'.",
-                  std::make_error_code(std::errc::executable_format_error)),
-              std::move(E));
-        });
-    break;
-  }
-  case InputFormats::YAML: {
-    EC = handleErrors(
-        LoadYAMLInstrMap(Filename, Sleds, FunctionAddresses, FunctionIds),
-        [&](std::unique_ptr<ErrorInfoBase> E) {
-          return joinErrors(
-              make_error<StringError>(
-                  Twine("Cannot load YAML instrumentation map from '") +
-                      Filename + "'.",
-                  std::make_error_code(std::errc::executable_format_error)),
-              std::move(E));
-        });
-    break;
-  }
-  }
-}
-
-void InstrumentationMapExtractor::exportAsYAML(raw_ostream &OS) {
+void exportAsYAML(const InstrumentationMap &Map, raw_ostream &OS) {
   // First we translate the sleds into the YAMLXRaySledEntry objects in a deque.
   std::vector<YAMLXRaySledEntry> YAMLSleds;
-  YAMLSleds.reserve(Sleds.size());
+  auto Sleds = Map.sleds();
+  YAMLSleds.reserve(std::distance(Sleds.begin(), Sleds.end()));
   for (const auto &Sled : Sleds) {
-    YAMLSleds.push_back({FunctionIds[Sled.Function], Sled.Address,
-                         Sled.Function, Sled.Kind, Sled.AlwaysInstrument});
+    auto FuncId = Map.getFunctionId(Sled.Function);
+    if (!FuncId)
+      return;
+    YAMLSleds.push_back({*FuncId, Sled.Address, Sled.Function, Sled.Kind,
+                         Sled.AlwaysInstrument});
   }
   Output Out(OS);
   Out << YAMLSleds;
 }
 
+} // namespace
+
 static CommandRegistration Unused(&Extract, []() -> Error {
-  Error Err = Error::success();
-  xray::InstrumentationMapExtractor Extractor(
-      ExtractInput, InstrumentationMapExtractor::InputFormats::ELF, Err);
-  if (Err)
-    return Err;
+  auto InstrumentationMapOrError = loadInstrumentationMap(ExtractInput);
+  if (!InstrumentationMapOrError)
+    return joinErrors(make_error<StringError>(
+                          Twine("Cannot extract instrumentation map from '") +
+                              ExtractInput + "'.",
+                          std::make_error_code(std::errc::invalid_argument)),
+                      InstrumentationMapOrError.takeError());
 
   std::error_code EC;
   raw_fd_ostream OS(ExtractOutput, EC, sys::fs::OpenFlags::F_Text);
   if (EC)
     return make_error<StringError>(
         Twine("Cannot open file '") + ExtractOutput + "' for writing.", EC);
-  Extractor.exportAsYAML(OS);
+  exportAsYAML(*InstrumentationMapOrError, OS);
   return Error::success();
 });