1 files changed, 112 insertions, 6 deletions

diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index b07f4027771..c8d2c2781c7 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -15,30 +15,62 @@
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/InstrProf.h"
 
+#include "InstrProfIndexed.h"
+
 #include <cassert>
 
 using namespace llvm;
 
-error_code InstrProfReader::create(std::string Path,
-                                   std::unique_ptr<InstrProfReader> &Result) {
-  std::unique_ptr<MemoryBuffer> Buffer;
+static error_code setupMemoryBuffer(std::string Path,
+                                    std::unique_ptr<MemoryBuffer> &Buffer) {
   if (error_code EC = MemoryBuffer::getFileOrSTDIN(Path, Buffer))
     return EC;
 
   // Sanity check the file.
   if (Buffer->getBufferSize() > std::numeric_limits<unsigned>::max())
     return instrprof_error::too_large;
+  return instrprof_error::success;
+}
+
+static error_code initializeReader(InstrProfReader &Reader) {
+  return Reader.readHeader();
+}
+
+error_code InstrProfReader::create(std::string Path,
+                                   std::unique_ptr<InstrProfReader> &Result) {
+  // Set up the buffer to read.
+  std::unique_ptr<MemoryBuffer> Buffer;
+  if (error_code EC = setupMemoryBuffer(Path, Buffer))
+    return EC;
 
   // Create the reader.
-  if (RawInstrProfReader64::hasFormat(*Buffer))
+  if (IndexedInstrProfReader::hasFormat(*Buffer))
+    Result.reset(new IndexedInstrProfReader(std::move(Buffer)));
+  else if (RawInstrProfReader64::hasFormat(*Buffer))
     Result.reset(new RawInstrProfReader64(std::move(Buffer)));
   else if (RawInstrProfReader32::hasFormat(*Buffer))
     Result.reset(new RawInstrProfReader32(std::move(Buffer)));
   else
     Result.reset(new TextInstrProfReader(std::move(Buffer)));
 
-  // Read the header and return the result.
-  return Result->readHeader();
+  // Initialize the reader and return the result.
+  return initializeReader(*Result);
+}
+
+error_code IndexedInstrProfReader::create(
+    std::string Path, std::unique_ptr<IndexedInstrProfReader> &Result) {
+  // Set up the buffer to read.
+  std::unique_ptr<MemoryBuffer> Buffer;
+  if (error_code EC = setupMemoryBuffer(Path, Buffer))
+    return EC;
+
+  // Create the reader.
+  if (!IndexedInstrProfReader::hasFormat(*Buffer))
+    return instrprof_error::bad_magic;
+  Result.reset(new IndexedInstrProfReader(std::move(Buffer)));
+
+  // Initialize the reader and return the result.
+  return initializeReader(*Result);
 }
 
 void InstrProfIterator::Increment() {
@@ -210,3 +242,77 @@ namespace llvm {
 template class RawInstrProfReader<uint32_t>;
 template class RawInstrProfReader<uint64_t>;
 }
+
+bool IndexedInstrProfReader::hasFormat(const MemoryBuffer &DataBuffer) {
+  if (DataBuffer.getBufferSize() < 8)
+    return false;
+  using namespace support;
+  uint64_t Magic =
+      endian::read<uint64_t, little, aligned>(DataBuffer.getBufferStart());
+  return Magic == IndexedInstrProf::Magic;
+}
+
+error_code IndexedInstrProfReader::readHeader() {
+  const unsigned char *Start = (unsigned char *)DataBuffer->getBufferStart();
+  const unsigned char *Cur = Start;
+  if ((unsigned char *)DataBuffer->getBufferEnd() - Cur < 24)
+    return error(instrprof_error::truncated);
+
+  using namespace support;
+
+  // Check the magic number.
+  uint64_t Magic = endian::readNext<uint64_t, little, unaligned>(Cur);
+  if (Magic != IndexedInstrProf::Magic)
+    return error(instrprof_error::bad_magic);
+
+  // Read the version.
+  uint64_t Version = endian::readNext<uint64_t, little, unaligned>(Cur);
+  if (Version != IndexedInstrProf::Version)
+    return error(instrprof_error::unsupported_version);
+
+  // Read the maximal function count.
+  MaxFunctionCount = endian::readNext<uint64_t, little, unaligned>(Cur);
+
+  // Read the hash type and start offset.
+  IndexedInstrProf::HashT HashType = static_cast<IndexedInstrProf::HashT>(
+      endian::readNext<uint64_t, little, unaligned>(Cur));
+  if (HashType > IndexedInstrProf::HashT::Last)
+    return error(instrprof_error::unsupported_hash_type);
+  uint64_t HashOffset = endian::readNext<uint64_t, little, unaligned>(Cur);
+
+  // The rest of the file is an on disk hash table.
+  Index.reset(InstrProfReaderIndex::Create(Start + HashOffset, Cur, Start,
+                                           InstrProfLookupTrait(HashType)));
+  // Set up our iterator for readNextRecord.
+  RecordIterator = Index->data_begin();
+
+  return success();
+}
+
+error_code IndexedInstrProfReader::getFunctionCounts(
+    StringRef FuncName, uint64_t &FuncHash, std::vector<uint64_t> &Counts) {
+  const auto &Iter = Index->find(FuncName);
+  if (Iter == Index->end())
+    return error(instrprof_error::unknown_function);
+
+  // Found it. Make sure it's valid before giving back a result.
+  const InstrProfRecord &Record = *Iter;
+  if (Record.Name.empty())
+    return error(instrprof_error::malformed);
+  FuncHash = Record.Hash;
+  Counts = Record.Counts;
+  return success();
+}
+
+error_code IndexedInstrProfReader::readNextRecord(InstrProfRecord &Record) {
+  // Are we out of records?
+  if (RecordIterator == Index->data_end())
+    return error(instrprof_error::eof);
+
+  // Read the next one.
+  Record = *RecordIterator;
+  ++RecordIterator;
+  if (Record.Name.empty())
+    return error(instrprof_error::malformed);
+  return success();
+}