1 files changed, 74 insertions, 51 deletions

diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index 885c1fe49240..37cdf4dd1fe2 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -52,16 +52,19 @@ static Error initializeReader(InstrProfReader &Reader) {
 }
 
 Expected<std::unique_ptr<InstrProfReader>>
-InstrProfReader::create(const Twine &Path) {
+InstrProfReader::create(const Twine &Path,
+                        const InstrProfCorrelator *Correlator) {
   // Set up the buffer to read.
   auto BufferOrError = setupMemoryBuffer(Path);
   if (Error E = BufferOrError.takeError())
     return std::move(E);
-  return InstrProfReader::create(std::move(BufferOrError.get()));
+  return InstrProfReader::create(std::move(BufferOrError.get()), Correlator);
 }
 
 Expected<std::unique_ptr<InstrProfReader>>
-InstrProfReader::create(std::unique_ptr<MemoryBuffer> Buffer) {
+InstrProfReader::create(std::unique_ptr<MemoryBuffer> Buffer,
+                        const InstrProfCorrelator *Correlator) {
+  // Sanity check the buffer.
   if (uint64_t(Buffer->getBufferSize()) > std::numeric_limits<uint64_t>::max())
     return make_error<InstrProfError>(instrprof_error::too_large);
 
@@ -73,9 +76,9 @@ InstrProfReader::create(std::unique_ptr<MemoryBuffer> Buffer) {
   if (IndexedInstrProfReader::hasFormat(*Buffer))
     Result.reset(new IndexedInstrProfReader(std::move(Buffer)));
   else if (RawInstrProfReader64::hasFormat(*Buffer))
-    Result.reset(new RawInstrProfReader64(std::move(Buffer)));
+    Result.reset(new RawInstrProfReader64(std::move(Buffer), Correlator));
   else if (RawInstrProfReader32::hasFormat(*Buffer))
-    Result.reset(new RawInstrProfReader32(std::move(Buffer)));
+    Result.reset(new RawInstrProfReader32(std::move(Buffer), Correlator));
   else if (TextInstrProfReader::hasFormat(*Buffer))
     Result.reset(new TextInstrProfReader(std::move(Buffer)));
   else
@@ -352,7 +355,7 @@ Error RawInstrProfReader<IntPtrT>::readNextHeader(const char *CurrentPos) {
 
 template <class IntPtrT>
 Error RawInstrProfReader<IntPtrT>::createSymtab(InstrProfSymtab &Symtab) {
-  if (Error E = Symtab.create(StringRef(NamesStart, NamesSize)))
+  if (Error E = Symtab.create(StringRef(NamesStart, NamesEnd - NamesStart)))
     return error(std::move(E));
   for (const RawInstrProf::ProfileData<IntPtrT> *I = Data; I != DataEnd; ++I) {
     const IntPtrT FPtr = swap(I->FunctionPointer);
@@ -369,6 +372,10 @@ Error RawInstrProfReader<IntPtrT>::readHeader(
   Version = swap(Header.Version);
   if (GET_VERSION(Version) != RawInstrProf::Version)
     return error(instrprof_error::unsupported_version);
+  if (useDebugInfoCorrelate() && !Correlator)
+    return error(instrprof_error::missing_debug_info_for_correlation);
+  if (!useDebugInfoCorrelate() && Correlator)
+    return error(instrprof_error::unexpected_debug_info_for_correlation);
 
   BinaryIdsSize = swap(Header.BinaryIdsSize);
   if (BinaryIdsSize % sizeof(uint64_t))
@@ -380,7 +387,7 @@ Error RawInstrProfReader<IntPtrT>::readHeader(
   auto PaddingBytesBeforeCounters = swap(Header.PaddingBytesBeforeCounters);
   auto CountersSize = swap(Header.CountersSize);
   auto PaddingBytesAfterCounters = swap(Header.PaddingBytesAfterCounters);
-  NamesSize = swap(Header.NamesSize);
+  auto NamesSize = swap(Header.NamesSize);
   ValueKindLast = swap(Header.ValueKindLast);
 
   auto DataSizeInBytes = DataSize * sizeof(RawInstrProf::ProfileData<IntPtrT>);
@@ -398,15 +405,27 @@ Error RawInstrProfReader<IntPtrT>::readHeader(
   if (Start + ValueDataOffset > DataBuffer->getBufferEnd())
     return error(instrprof_error::bad_header);
 
-  Data = reinterpret_cast<const RawInstrProf::ProfileData<IntPtrT> *>(
-      Start + DataOffset);
-  DataEnd = Data + DataSize;
+  if (Correlator) {
+    // These sizes in the raw file are zero because we constructed them in the
+    // Correlator.
+    assert(DataSize == 0 && NamesSize == 0);
+    assert(CountersDelta == 0 && NamesDelta == 0);
+    Data = Correlator->getDataPointer();
+    DataEnd = Data + Correlator->getDataSize();
+    NamesStart = Correlator->getCompressedNamesPointer();
+    NamesEnd = NamesStart + Correlator->getCompressedNamesSize();
+  } else {
+    Data = reinterpret_cast<const RawInstrProf::ProfileData<IntPtrT> *>(
+        Start + DataOffset);
+    DataEnd = Data + DataSize;
+    NamesStart = Start + NamesOffset;
+    NamesEnd = NamesStart + NamesSize;
+  }
 
   // Binary ids start just after the header.
   BinaryIdsStart =
       reinterpret_cast<const uint8_t *>(&Header) + sizeof(RawInstrProf::Header);
   CountersStart = reinterpret_cast<const uint64_t *>(Start + CountersOffset);
-  NamesStart = Start + NamesOffset;
   ValueDataStart = reinterpret_cast<const uint8_t *>(Start + ValueDataOffset);
 
   const uint8_t *BufferEnd = (const uint8_t *)DataBuffer->getBufferEnd();
@@ -440,45 +459,50 @@ Error RawInstrProfReader<IntPtrT>::readRawCounts(
   if (NumCounters == 0)
     return error(instrprof_error::malformed, "number of counters is zero");
 
-  IntPtrT CounterPtr = Data->CounterPtr;
-  auto *NamesStartAsCounter = reinterpret_cast<const uint64_t *>(NamesStart);
-  ptrdiff_t MaxNumCounters = NamesStartAsCounter - CountersStart;
-
-  // Check bounds. Note that the counter pointer embedded in the data record
-  // may itself be corrupt.
-  if (MaxNumCounters < 0 || NumCounters > (uint32_t)MaxNumCounters)
-    return error(instrprof_error::malformed,
-                 "counter pointer is out of bounds");
-
-  // We need to compute the in-buffer counter offset from the in-memory address
-  // distance. The initial CountersDelta is the in-memory address difference
-  // start(__llvm_prf_cnts)-start(__llvm_prf_data), so SrcData->CounterPtr -
-  // CountersDelta computes the offset into the in-buffer counter section.
-  //
-  // CountersDelta decreases as we advance to the next data record.
-  ptrdiff_t CounterOffset = getCounterOffset(CounterPtr);
-  CountersDelta -= sizeof(*Data);
-  if (CounterOffset < 0)
-    return error(
-        instrprof_error::malformed,
-        ("counter offset " + Twine(CounterOffset) + " is negative").str());
+  ArrayRef<uint64_t> RawCounts;
+  if (Correlator) {
+    uint64_t CounterOffset = swap<IntPtrT>(Data->CounterPtr) / sizeof(uint64_t);
+    RawCounts =
+        makeArrayRef<uint64_t>(CountersStart + CounterOffset, NumCounters);
+  } else {
+    IntPtrT CounterPtr = Data->CounterPtr;
+    ptrdiff_t CounterOffset = getCounterOffset(CounterPtr);
+    if (CounterOffset < 0)
+      return error(
+          instrprof_error::malformed,
+          ("counter offset " + Twine(CounterOffset) + " is negative").str());
 
-  if (CounterOffset > MaxNumCounters)
-    return error(instrprof_error::malformed,
-                 ("counter offset " + Twine(CounterOffset) +
-                  " is greater than the maximum number of counters " +
-                  Twine((uint32_t)MaxNumCounters))
-                     .str());
+    // Check bounds. Note that the counter pointer embedded in the data record
+    // may itself be corrupt.
+    auto *NamesStartAsCounter = reinterpret_cast<const uint64_t *>(NamesStart);
+    ptrdiff_t MaxNumCounters = NamesStartAsCounter - CountersStart;
+    if (MaxNumCounters < 0 || NumCounters > (uint32_t)MaxNumCounters)
+      return error(instrprof_error::malformed,
+                   "counter pointer is out of bounds");
+    // We need to compute the in-buffer counter offset from the in-memory
+    // address distance. The initial CountersDelta is the in-memory address
+    // difference start(__llvm_prf_cnts)-start(__llvm_prf_data), so
+    // SrcData->CounterPtr - CountersDelta computes the offset into the
+    // in-buffer counter section.
+    if (CounterOffset > MaxNumCounters)
+      return error(instrprof_error::malformed,
+                   ("counter offset " + Twine(CounterOffset) +
+                    " is greater than the maximum number of counters " +
+                    Twine((uint32_t)MaxNumCounters))
+                       .str());
 
-  if (((uint32_t)CounterOffset + NumCounters) > (uint32_t)MaxNumCounters)
-    return error(instrprof_error::malformed,
-                 ("number of counters " +
-                  Twine(((uint32_t)CounterOffset + NumCounters)) +
-                  " is greater than the maximum number of counters " +
-                  Twine((uint32_t)MaxNumCounters))
-                     .str());
+    if (((uint32_t)CounterOffset + NumCounters) > (uint32_t)MaxNumCounters)
+      return error(instrprof_error::malformed,
+                   ("number of counters " +
+                    Twine(((uint32_t)CounterOffset + NumCounters)) +
+                    " is greater than the maximum number of counters " +
+                    Twine((uint32_t)MaxNumCounters))
+                       .str());
+    // CountersDelta decreases as we advance to the next data record.
+    CountersDelta -= sizeof(*Data);
 
-  auto RawCounts = makeArrayRef(getCounter(CounterOffset), NumCounters);
+    RawCounts = makeArrayRef(getCounter(CounterOffset), NumCounters);
+  }
 
   if (ShouldSwapBytes) {
     Record.Counts.clear();
@@ -977,11 +1001,10 @@ IndexedInstrProfReader::getInstrProfRecord(StringRef FuncName,
   if (Err)
     return std::move(Err);
   // Found it. Look for counters with the right hash.
-  for (unsigned I = 0, E = Data.size(); I < E; ++I) {
+  for (const NamedInstrProfRecord &I : Data) {
     // Check for a match and fill the vector if there is one.
-    if (Data[I].Hash == FuncHash) {
-      return std::move(Data[I]);
-    }
+    if (I.Hash == FuncHash)
+      return std::move(I);
   }
   return error(instrprof_error::hash_mismatch);
 }