1 files changed, 106 insertions, 4 deletions

diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index 8b3600090b8..c563355599f 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -14,12 +14,23 @@
 
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/InstrProf.h"
-#include "llvm/Support/Endian.h"
 
 #include <cassert>
 
 using namespace llvm;
 
+static uint64_t getRawMagic() {
+  return
+    uint64_t('l') << 56 |
+    uint64_t('p') << 48 |
+    uint64_t('r') << 40 |
+    uint64_t('o') << 32 |
+    uint64_t('f') << 24 |
+    uint64_t('r') << 16 |
+    uint64_t('a') <<  8 |
+    uint64_t('w');
+}
+
 error_code InstrProfReader::create(std::string Path,
                                    std::unique_ptr<InstrProfReader> &Result) {
   std::unique_ptr<MemoryBuffer> Buffer;
@@ -30,10 +41,19 @@ error_code InstrProfReader::create(std::string Path,
   if (Buffer->getBufferSize() > std::numeric_limits<unsigned>::max())
     return instrprof_error::too_large;
 
-  // FIXME: This needs to determine which format the file is and construct the
-  // correct subclass.
-  Result.reset(new TextInstrProfReader(Buffer));
+  if (Buffer->getBufferSize() < sizeof(uint64_t)) {
+    Result.reset(new TextInstrProfReader(Buffer));
+    Result->readHeader();
+    return instrprof_error::success;
+  }
 
+  uint64_t Magic = *(uint64_t *)Buffer->getBufferStart();
+  uint64_t SwappedMagic = sys::SwapByteOrder(Magic);
+  if (Magic == getRawMagic() || SwappedMagic == getRawMagic())
+    Result.reset(new RawInstrProfReader(Buffer));
+  else
+    Result.reset(new TextInstrProfReader(Buffer));
+  Result->readHeader();
   return instrprof_error::success;
 }
 
@@ -82,3 +102,85 @@ error_code TextInstrProfReader::readNextRecord(InstrProfRecord &Record) {
 
   return success();
 }
+
+static uint64_t getRawVersion() {
+  return 1;
+}
+namespace {
+}
+RawInstrProfReader::RawInstrProfReader(std::unique_ptr<MemoryBuffer> &DataBuffer)
+    : DataBuffer(DataBuffer.release()) { }
+
+error_code RawInstrProfReader::readHeader() {
+  if (DataBuffer->getBufferSize() < sizeof(RawHeader))
+    return error(instrprof_error::malformed);
+  const RawHeader *Header = (RawHeader *)DataBuffer->getBufferStart();
+  if (Header->Magic == getRawMagic())
+    ShouldSwapBytes = false;
+  else {
+    if (sys::SwapByteOrder(Header->Magic) != getRawMagic())
+      return error(instrprof_error::malformed);
+
+    ShouldSwapBytes = true;
+  }
+  return readHeader(*Header);
+}
+
+error_code RawInstrProfReader::readHeader(const RawHeader &Header) {
+  if (swap(Header.Version) != getRawVersion())
+    return error(instrprof_error::unsupported_version);
+
+  CountersDelta = swap(Header.CountersDelta);
+  NamesDelta = swap(Header.NamesDelta);
+  auto DataSize = swap(Header.DataSize);
+  auto CountersSize = swap(Header.CountersSize);
+  auto NamesSize = swap(Header.NamesSize);
+
+  ptrdiff_t DataOffset = sizeof(RawHeader);
+  ptrdiff_t CountersOffset = DataOffset + sizeof(ProfileData) * DataSize;
+  ptrdiff_t NamesOffset = CountersOffset + sizeof(uint64_t) * CountersSize;
+  size_t FileSize = NamesOffset + sizeof(char) * NamesSize;
+
+  if (FileSize != DataBuffer->getBufferSize())
+    return error(instrprof_error::malformed);
+
+  Data = (ProfileData *)(DataBuffer->getBufferStart() + DataOffset);
+  DataEnd = Data + DataSize;
+  CountersStart = (uint64_t *)(DataBuffer->getBufferStart() + CountersOffset);
+  NamesStart = DataBuffer->getBufferStart() + NamesOffset;
+
+  return success();
+}
+
+error_code RawInstrProfReader::readNextRecord(InstrProfRecord &Record) {
+  if (Data == DataEnd)
+    return error(instrprof_error::eof);
+
+  // Get the raw data.
+  StringRef RawName(getName(Data->NamePtr), swap(Data->NameSize));
+  auto RawCounts = makeArrayRef(getCounter(Data->CounterPtr),
+                                swap(Data->NumCounters));
+
+  // Check bounds.
+  if (RawName.data() < NamesStart ||
+      RawName.data() + RawName.size() > DataBuffer->getBufferEnd() ||
+      RawCounts.data() < CountersStart ||
+      RawCounts.data() + RawCounts.size() > (uint64_t *)NamesStart)
+    return error(instrprof_error::malformed);
+
+  // Store the data in Record, byte-swapping as necessary.
+  Record.Hash = swap(Data->FuncHash);
+  Record.Name = RawName;
+  if (ShouldSwapBytes) {
+    Counts.clear();
+    Counts.reserve(RawCounts.size());
+    for (uint64_t Count : RawCounts)
+      Counts.push_back(swap(Count));
+    Record.Counts = Counts;
+  } else
+    Record.Counts = RawCounts;
+
+  // Iterate.
+  ++Data;
+  return success();
+}