1 //===- InstrProfReader.cpp - Instrumented profiling reader ----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file contains support for reading profiling data for clang's
10 // instrumentation based PGO and coverage.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "llvm/ProfileData/InstrProfReader.h"
15 #include "llvm/ADT/ArrayRef.h"
16 #include "llvm/ADT/DenseMap.h"
17 #include "llvm/ADT/STLExtras.h"
18 #include "llvm/ADT/StringRef.h"
19 #include "llvm/IR/ProfileSummary.h"
20 #include "llvm/ProfileData/InstrProf.h"
21 #include "llvm/ProfileData/ProfileCommon.h"
22 #include "llvm/Support/Endian.h"
23 #include "llvm/Support/Error.h"
24 #include "llvm/Support/ErrorOr.h"
25 #include "llvm/Support/MemoryBuffer.h"
26 #include "llvm/Support/SymbolRemappingReader.h"
27 #include "llvm/Support/SwapByteOrder.h"
28 #include <algorithm>
29 #include <cctype>
30 #include <cstddef>
31 #include <cstdint>
32 #include <limits>
33 #include <memory>
34 #include <system_error>
35 #include <utility>
36 #include <vector>
37 
38 using namespace llvm;
39 
40 static Expected<std::unique_ptr<MemoryBuffer>>
41 setupMemoryBuffer(const Twine &Path) {
42   ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
43       MemoryBuffer::getFileOrSTDIN(Path);
44   if (std::error_code EC = BufferOrErr.getError())
45     return errorCodeToError(EC);
46   return std::move(BufferOrErr.get());
47 }
48 
49 static Error initializeReader(InstrProfReader &Reader) {
50   return Reader.readHeader();
51 }
52 
53 Expected<std::unique_ptr<InstrProfReader>>
54 InstrProfReader::create(const Twine &Path) {
55   // Set up the buffer to read.
56   auto BufferOrError = setupMemoryBuffer(Path);
57   if (Error E = BufferOrError.takeError())
58     return std::move(E);
59   return InstrProfReader::create(std::move(BufferOrError.get()));
60 }
61 
62 Expected<std::unique_ptr<InstrProfReader>>
63 InstrProfReader::create(std::unique_ptr<MemoryBuffer> Buffer) {
64   // Sanity check the buffer.
65   if (uint64_t(Buffer->getBufferSize()) > std::numeric_limits<unsigned>::max())
66     return make_error<InstrProfError>(instrprof_error::too_large);
67 
68   if (Buffer->getBufferSize() == 0)
69     return make_error<InstrProfError>(instrprof_error::empty_raw_profile);
70 
71   std::unique_ptr<InstrProfReader> Result;
72   // Create the reader.
73   if (IndexedInstrProfReader::hasFormat(*Buffer))
74     Result.reset(new IndexedInstrProfReader(std::move(Buffer)));
75   else if (RawInstrProfReader64::hasFormat(*Buffer))
76     Result.reset(new RawInstrProfReader64(std::move(Buffer)));
77   else if (RawInstrProfReader32::hasFormat(*Buffer))
78     Result.reset(new RawInstrProfReader32(std::move(Buffer)));
79   else if (TextInstrProfReader::hasFormat(*Buffer))
80     Result.reset(new TextInstrProfReader(std::move(Buffer)));
81   else
82     return make_error<InstrProfError>(instrprof_error::unrecognized_format);
83 
84   // Initialize the reader and return the result.
85   if (Error E = initializeReader(*Result))
86     return std::move(E);
87 
88   return std::move(Result);
89 }
90 
91 Expected<std::unique_ptr<IndexedInstrProfReader>>
92 IndexedInstrProfReader::create(const Twine &Path, const Twine &RemappingPath) {
93   // Set up the buffer to read.
94   auto BufferOrError = setupMemoryBuffer(Path);
95   if (Error E = BufferOrError.takeError())
96     return std::move(E);
97 
98   // Set up the remapping buffer if requested.
99   std::unique_ptr<MemoryBuffer> RemappingBuffer;
100   std::string RemappingPathStr = RemappingPath.str();
101   if (!RemappingPathStr.empty()) {
102     auto RemappingBufferOrError = setupMemoryBuffer(RemappingPathStr);
103     if (Error E = RemappingBufferOrError.takeError())
104       return std::move(E);
105     RemappingBuffer = std::move(RemappingBufferOrError.get());
106   }
107 
108   return IndexedInstrProfReader::create(std::move(BufferOrError.get()),
109                                         std::move(RemappingBuffer));
110 }
111 
112 Expected<std::unique_ptr<IndexedInstrProfReader>>
113 IndexedInstrProfReader::create(std::unique_ptr<MemoryBuffer> Buffer,
114                                std::unique_ptr<MemoryBuffer> RemappingBuffer) {
115   // Sanity check the buffer.
116   if (uint64_t(Buffer->getBufferSize()) > std::numeric_limits<unsigned>::max())
117     return make_error<InstrProfError>(instrprof_error::too_large);
118 
119   // Create the reader.
120   if (!IndexedInstrProfReader::hasFormat(*Buffer))
121     return make_error<InstrProfError>(instrprof_error::bad_magic);
122   auto Result = llvm::make_unique<IndexedInstrProfReader>(
123       std::move(Buffer), std::move(RemappingBuffer));
124 
125   // Initialize the reader and return the result.
126   if (Error E = initializeReader(*Result))
127     return std::move(E);
128 
129   return std::move(Result);
130 }
131 
132 void InstrProfIterator::Increment() {
133   if (auto E = Reader->readNextRecord(Record)) {
134     // Handle errors in the reader.
135     InstrProfError::take(std::move(E));
136     *this = InstrProfIterator();
137   }
138 }
139 
140 bool TextInstrProfReader::hasFormat(const MemoryBuffer &Buffer) {
141   // Verify that this really looks like plain ASCII text by checking a
142   // 'reasonable' number of characters (up to profile magic size).
143   size_t count = std::min(Buffer.getBufferSize(), sizeof(uint64_t));
144   StringRef buffer = Buffer.getBufferStart();
145   return count == 0 ||
146          std::all_of(buffer.begin(), buffer.begin() + count,
147                      [](char c) { return isPrint(c) || ::isspace(c); });
148 }
149 
150 // Read the profile variant flag from the header: ":FE" means this is a FE
151 // generated profile. ":IR" means this is an IR level profile. Other strings
152 // with a leading ':' will be reported an error format.
153 Error TextInstrProfReader::readHeader() {
154   Symtab.reset(new InstrProfSymtab());
155   bool IsIRInstr = false;
156   if (!Line->startswith(":")) {
157     IsIRLevelProfile = false;
158     return success();
159   }
160   StringRef Str = (Line)->substr(1);
161   if (Str.equals_lower("ir"))
162     IsIRInstr = true;
163   else if (Str.equals_lower("fe"))
164     IsIRInstr = false;
165   else
166     return error(instrprof_error::bad_header);
167 
168   ++Line;
169   IsIRLevelProfile = IsIRInstr;
170   return success();
171 }
172 
173 Error
174 TextInstrProfReader::readValueProfileData(InstrProfRecord &Record) {
175 
176 #define CHECK_LINE_END(Line)                                                   \
177   if (Line.is_at_end())                                                        \
178     return error(instrprof_error::truncated);
179 #define READ_NUM(Str, Dst)                                                     \
180   if ((Str).getAsInteger(10, (Dst)))                                           \
181     return error(instrprof_error::malformed);
182 #define VP_READ_ADVANCE(Val)                                                   \
183   CHECK_LINE_END(Line);                                                        \
184   uint32_t Val;                                                                \
185   READ_NUM((*Line), (Val));                                                    \
186   Line++;
187 
188   if (Line.is_at_end())
189     return success();
190 
191   uint32_t NumValueKinds;
192   if (Line->getAsInteger(10, NumValueKinds)) {
193     // No value profile data
194     return success();
195   }
196   if (NumValueKinds == 0 || NumValueKinds > IPVK_Last + 1)
197     return error(instrprof_error::malformed);
198   Line++;
199 
200   for (uint32_t VK = 0; VK < NumValueKinds; VK++) {
201     VP_READ_ADVANCE(ValueKind);
202     if (ValueKind > IPVK_Last)
203       return error(instrprof_error::malformed);
204     VP_READ_ADVANCE(NumValueSites);
205     if (!NumValueSites)
206       continue;
207 
208     Record.reserveSites(VK, NumValueSites);
209     for (uint32_t S = 0; S < NumValueSites; S++) {
210       VP_READ_ADVANCE(NumValueData);
211 
212       std::vector<InstrProfValueData> CurrentValues;
213       for (uint32_t V = 0; V < NumValueData; V++) {
214         CHECK_LINE_END(Line);
215         std::pair<StringRef, StringRef> VD = Line->rsplit(':');
216         uint64_t TakenCount, Value;
217         if (ValueKind == IPVK_IndirectCallTarget) {
218           if (InstrProfSymtab::isExternalSymbol(VD.first)) {
219             Value = 0;
220           } else {
221             if (Error E = Symtab->addFuncName(VD.first))
222               return E;
223             Value = IndexedInstrProf::ComputeHash(VD.first);
224           }
225         } else {
226           READ_NUM(VD.first, Value);
227         }
228         READ_NUM(VD.second, TakenCount);
229         CurrentValues.push_back({Value, TakenCount});
230         Line++;
231       }
232       Record.addValueData(ValueKind, S, CurrentValues.data(), NumValueData,
233                           nullptr);
234     }
235   }
236   return success();
237 
238 #undef CHECK_LINE_END
239 #undef READ_NUM
240 #undef VP_READ_ADVANCE
241 }
242 
243 Error TextInstrProfReader::readNextRecord(NamedInstrProfRecord &Record) {
244   // Skip empty lines and comments.
245   while (!Line.is_at_end() && (Line->empty() || Line->startswith("#")))
246     ++Line;
247   // If we hit EOF while looking for a name, we're done.
248   if (Line.is_at_end()) {
249     return error(instrprof_error::eof);
250   }
251 
252   // Read the function name.
253   Record.Name = *Line++;
254   if (Error E = Symtab->addFuncName(Record.Name))
255     return error(std::move(E));
256 
257   // Read the function hash.
258   if (Line.is_at_end())
259     return error(instrprof_error::truncated);
260   if ((Line++)->getAsInteger(0, Record.Hash))
261     return error(instrprof_error::malformed);
262 
263   // Read the number of counters.
264   uint64_t NumCounters;
265   if (Line.is_at_end())
266     return error(instrprof_error::truncated);
267   if ((Line++)->getAsInteger(10, NumCounters))
268     return error(instrprof_error::malformed);
269   if (NumCounters == 0)
270     return error(instrprof_error::malformed);
271 
272   // Read each counter and fill our internal storage with the values.
273   Record.Clear();
274   Record.Counts.reserve(NumCounters);
275   for (uint64_t I = 0; I < NumCounters; ++I) {
276     if (Line.is_at_end())
277       return error(instrprof_error::truncated);
278     uint64_t Count;
279     if ((Line++)->getAsInteger(10, Count))
280       return error(instrprof_error::malformed);
281     Record.Counts.push_back(Count);
282   }
283 
284   // Check if value profile data exists and read it if so.
285   if (Error E = readValueProfileData(Record))
286     return error(std::move(E));
287 
288   return success();
289 }
290 
291 template <class IntPtrT>
292 bool RawInstrProfReader<IntPtrT>::hasFormat(const MemoryBuffer &DataBuffer) {
293   if (DataBuffer.getBufferSize() < sizeof(uint64_t))
294     return false;
295   uint64_t Magic =
296     *reinterpret_cast<const uint64_t *>(DataBuffer.getBufferStart());
297   return RawInstrProf::getMagic<IntPtrT>() == Magic ||
298          sys::getSwappedBytes(RawInstrProf::getMagic<IntPtrT>()) == Magic;
299 }
300 
301 template <class IntPtrT>
302 Error RawInstrProfReader<IntPtrT>::readHeader() {
303   if (!hasFormat(*DataBuffer))
304     return error(instrprof_error::bad_magic);
305   if (DataBuffer->getBufferSize() < sizeof(RawInstrProf::Header))
306     return error(instrprof_error::bad_header);
307   auto *Header = reinterpret_cast<const RawInstrProf::Header *>(
308       DataBuffer->getBufferStart());
309   ShouldSwapBytes = Header->Magic != RawInstrProf::getMagic<IntPtrT>();
310   return readHeader(*Header);
311 }
312 
313 template <class IntPtrT>
314 Error RawInstrProfReader<IntPtrT>::readNextHeader(const char *CurrentPos) {
315   const char *End = DataBuffer->getBufferEnd();
316   // Skip zero padding between profiles.
317   while (CurrentPos != End && *CurrentPos == 0)
318     ++CurrentPos;
319   // If there's nothing left, we're done.
320   if (CurrentPos == End)
321     return make_error<InstrProfError>(instrprof_error::eof);
322   // If there isn't enough space for another header, this is probably just
323   // garbage at the end of the file.
324   if (CurrentPos + sizeof(RawInstrProf::Header) > End)
325     return make_error<InstrProfError>(instrprof_error::malformed);
326   // The writer ensures each profile is padded to start at an aligned address.
327   if (reinterpret_cast<size_t>(CurrentPos) % alignof(uint64_t))
328     return make_error<InstrProfError>(instrprof_error::malformed);
329   // The magic should have the same byte order as in the previous header.
330   uint64_t Magic = *reinterpret_cast<const uint64_t *>(CurrentPos);
331   if (Magic != swap(RawInstrProf::getMagic<IntPtrT>()))
332     return make_error<InstrProfError>(instrprof_error::bad_magic);
333 
334   // There's another profile to read, so we need to process the header.
335   auto *Header = reinterpret_cast<const RawInstrProf::Header *>(CurrentPos);
336   return readHeader(*Header);
337 }
338 
339 template <class IntPtrT>
340 Error RawInstrProfReader<IntPtrT>::createSymtab(InstrProfSymtab &Symtab) {
341   if (Error E = Symtab.create(StringRef(NamesStart, NamesSize)))
342     return error(std::move(E));
343   for (const RawInstrProf::ProfileData<IntPtrT> *I = Data; I != DataEnd; ++I) {
344     const IntPtrT FPtr = swap(I->FunctionPointer);
345     if (!FPtr)
346       continue;
347     Symtab.mapAddress(FPtr, I->NameRef);
348   }
349   return success();
350 }
351 
352 template <class IntPtrT>
353 Error RawInstrProfReader<IntPtrT>::readHeader(
354     const RawInstrProf::Header &Header) {
355   Version = swap(Header.Version);
356   if (GET_VERSION(Version) != RawInstrProf::Version)
357     return error(instrprof_error::unsupported_version);
358 
359   CountersDelta = swap(Header.CountersDelta);
360   NamesDelta = swap(Header.NamesDelta);
361   auto DataSize = swap(Header.DataSize);
362   auto CountersSize = swap(Header.CountersSize);
363   NamesSize = swap(Header.NamesSize);
364   ValueKindLast = swap(Header.ValueKindLast);
365 
366   auto DataSizeInBytes = DataSize * sizeof(RawInstrProf::ProfileData<IntPtrT>);
367   auto PaddingSize = getNumPaddingBytes(NamesSize);
368 
369   ptrdiff_t DataOffset = sizeof(RawInstrProf::Header);
370   ptrdiff_t CountersOffset = DataOffset + DataSizeInBytes;
371   ptrdiff_t NamesOffset = CountersOffset + sizeof(uint64_t) * CountersSize;
372   ptrdiff_t ValueDataOffset = NamesOffset + NamesSize + PaddingSize;
373 
374   auto *Start = reinterpret_cast<const char *>(&Header);
375   if (Start + ValueDataOffset > DataBuffer->getBufferEnd())
376     return error(instrprof_error::bad_header);
377 
378   Data = reinterpret_cast<const RawInstrProf::ProfileData<IntPtrT> *>(
379       Start + DataOffset);
380   DataEnd = Data + DataSize;
381   CountersStart = reinterpret_cast<const uint64_t *>(Start + CountersOffset);
382   NamesStart = Start + NamesOffset;
383   ValueDataStart = reinterpret_cast<const uint8_t *>(Start + ValueDataOffset);
384 
385   std::unique_ptr<InstrProfSymtab> NewSymtab = make_unique<InstrProfSymtab>();
386   if (Error E = createSymtab(*NewSymtab.get()))
387     return E;
388 
389   Symtab = std::move(NewSymtab);
390   return success();
391 }
392 
393 template <class IntPtrT>
394 Error RawInstrProfReader<IntPtrT>::readName(NamedInstrProfRecord &Record) {
395   Record.Name = getName(Data->NameRef);
396   return success();
397 }
398 
399 template <class IntPtrT>
400 Error RawInstrProfReader<IntPtrT>::readFuncHash(NamedInstrProfRecord &Record) {
401   Record.Hash = swap(Data->FuncHash);
402   return success();
403 }
404 
405 template <class IntPtrT>
406 Error RawInstrProfReader<IntPtrT>::readRawCounts(
407     InstrProfRecord &Record) {
408   uint32_t NumCounters = swap(Data->NumCounters);
409   IntPtrT CounterPtr = Data->CounterPtr;
410   if (NumCounters == 0)
411     return error(instrprof_error::malformed);
412 
413   auto RawCounts = makeArrayRef(getCounter(CounterPtr), NumCounters);
414   auto *NamesStartAsCounter = reinterpret_cast<const uint64_t *>(NamesStart);
415 
416   // Check bounds.
417   if (RawCounts.data() < CountersStart ||
418       RawCounts.data() + RawCounts.size() > NamesStartAsCounter)
419     return error(instrprof_error::malformed);
420 
421   if (ShouldSwapBytes) {
422     Record.Counts.clear();
423     Record.Counts.reserve(RawCounts.size());
424     for (uint64_t Count : RawCounts)
425       Record.Counts.push_back(swap(Count));
426   } else
427     Record.Counts = RawCounts;
428 
429   return success();
430 }
431 
432 template <class IntPtrT>
433 Error RawInstrProfReader<IntPtrT>::readValueProfilingData(
434     InstrProfRecord &Record) {
435   Record.clearValueData();
436   CurValueDataSize = 0;
437   // Need to match the logic in value profile dumper code in compiler-rt:
438   uint32_t NumValueKinds = 0;
439   for (uint32_t I = 0; I < IPVK_Last + 1; I++)
440     NumValueKinds += (Data->NumValueSites[I] != 0);
441 
442   if (!NumValueKinds)
443     return success();
444 
445   Expected<std::unique_ptr<ValueProfData>> VDataPtrOrErr =
446       ValueProfData::getValueProfData(
447           ValueDataStart, (const unsigned char *)DataBuffer->getBufferEnd(),
448           getDataEndianness());
449 
450   if (Error E = VDataPtrOrErr.takeError())
451     return E;
452 
453   // Note that besides deserialization, this also performs the conversion for
454   // indirect call targets.  The function pointers from the raw profile are
455   // remapped into function name hashes.
456   VDataPtrOrErr.get()->deserializeTo(Record, Symtab.get());
457   CurValueDataSize = VDataPtrOrErr.get()->getSize();
458   return success();
459 }
460 
461 template <class IntPtrT>
462 Error RawInstrProfReader<IntPtrT>::readNextRecord(NamedInstrProfRecord &Record) {
463   if (atEnd())
464     // At this point, ValueDataStart field points to the next header.
465     if (Error E = readNextHeader(getNextHeaderPos()))
466       return error(std::move(E));
467 
468   // Read name ad set it in Record.
469   if (Error E = readName(Record))
470     return error(std::move(E));
471 
472   // Read FuncHash and set it in Record.
473   if (Error E = readFuncHash(Record))
474     return error(std::move(E));
475 
476   // Read raw counts and set Record.
477   if (Error E = readRawCounts(Record))
478     return error(std::move(E));
479 
480   // Read value data and set Record.
481   if (Error E = readValueProfilingData(Record))
482     return error(std::move(E));
483 
484   // Iterate.
485   advanceData();
486   return success();
487 }
488 
489 namespace llvm {
490 
491 template class RawInstrProfReader<uint32_t>;
492 template class RawInstrProfReader<uint64_t>;
493 
494 } // end namespace llvm
495 
496 InstrProfLookupTrait::hash_value_type
497 InstrProfLookupTrait::ComputeHash(StringRef K) {
498   return IndexedInstrProf::ComputeHash(HashType, K);
499 }
500 
501 using data_type = InstrProfLookupTrait::data_type;
502 using offset_type = InstrProfLookupTrait::offset_type;
503 
504 bool InstrProfLookupTrait::readValueProfilingData(
505     const unsigned char *&D, const unsigned char *const End) {
506   Expected<std::unique_ptr<ValueProfData>> VDataPtrOrErr =
507       ValueProfData::getValueProfData(D, End, ValueProfDataEndianness);
508 
509   if (VDataPtrOrErr.takeError())
510     return false;
511 
512   VDataPtrOrErr.get()->deserializeTo(DataBuffer.back(), nullptr);
513   D += VDataPtrOrErr.get()->TotalSize;
514 
515   return true;
516 }
517 
518 data_type InstrProfLookupTrait::ReadData(StringRef K, const unsigned char *D,
519                                          offset_type N) {
520   using namespace support;
521 
522   // Check if the data is corrupt. If so, don't try to read it.
523   if (N % sizeof(uint64_t))
524     return data_type();
525 
526   DataBuffer.clear();
527   std::vector<uint64_t> CounterBuffer;
528 
529   const unsigned char *End = D + N;
530   while (D < End) {
531     // Read hash.
532     if (D + sizeof(uint64_t) >= End)
533       return data_type();
534     uint64_t Hash = endian::readNext<uint64_t, little, unaligned>(D);
535 
536     // Initialize number of counters for GET_VERSION(FormatVersion) == 1.
537     uint64_t CountsSize = N / sizeof(uint64_t) - 1;
538     // If format version is different then read the number of counters.
539     if (GET_VERSION(FormatVersion) != IndexedInstrProf::ProfVersion::Version1) {
540       if (D + sizeof(uint64_t) > End)
541         return data_type();
542       CountsSize = endian::readNext<uint64_t, little, unaligned>(D);
543     }
544     // Read counter values.
545     if (D + CountsSize * sizeof(uint64_t) > End)
546       return data_type();
547 
548     CounterBuffer.clear();
549     CounterBuffer.reserve(CountsSize);
550     for (uint64_t J = 0; J < CountsSize; ++J)
551       CounterBuffer.push_back(endian::readNext<uint64_t, little, unaligned>(D));
552 
553     DataBuffer.emplace_back(K, Hash, std::move(CounterBuffer));
554 
555     // Read value profiling data.
556     if (GET_VERSION(FormatVersion) > IndexedInstrProf::ProfVersion::Version2 &&
557         !readValueProfilingData(D, End)) {
558       DataBuffer.clear();
559       return data_type();
560     }
561   }
562   return DataBuffer;
563 }
564 
565 template <typename HashTableImpl>
566 Error InstrProfReaderIndex<HashTableImpl>::getRecords(
567     StringRef FuncName, ArrayRef<NamedInstrProfRecord> &Data) {
568   auto Iter = HashTable->find(FuncName);
569   if (Iter == HashTable->end())
570     return make_error<InstrProfError>(instrprof_error::unknown_function);
571 
572   Data = (*Iter);
573   if (Data.empty())
574     return make_error<InstrProfError>(instrprof_error::malformed);
575 
576   return Error::success();
577 }
578 
579 template <typename HashTableImpl>
580 Error InstrProfReaderIndex<HashTableImpl>::getRecords(
581     ArrayRef<NamedInstrProfRecord> &Data) {
582   if (atEnd())
583     return make_error<InstrProfError>(instrprof_error::eof);
584 
585   Data = *RecordIterator;
586 
587   if (Data.empty())
588     return make_error<InstrProfError>(instrprof_error::malformed);
589 
590   return Error::success();
591 }
592 
593 template <typename HashTableImpl>
594 InstrProfReaderIndex<HashTableImpl>::InstrProfReaderIndex(
595     const unsigned char *Buckets, const unsigned char *const Payload,
596     const unsigned char *const Base, IndexedInstrProf::HashT HashType,
597     uint64_t Version) {
598   FormatVersion = Version;
599   HashTable.reset(HashTableImpl::Create(
600       Buckets, Payload, Base,
601       typename HashTableImpl::InfoType(HashType, Version)));
602   RecordIterator = HashTable->data_begin();
603 }
604 
605 namespace {
606 /// A remapper that does not apply any remappings.
607 class InstrProfReaderNullRemapper : public InstrProfReaderRemapper {
608   InstrProfReaderIndexBase &Underlying;
609 
610 public:
611   InstrProfReaderNullRemapper(InstrProfReaderIndexBase &Underlying)
612       : Underlying(Underlying) {}
613 
614   Error getRecords(StringRef FuncName,
615                    ArrayRef<NamedInstrProfRecord> &Data) override {
616     return Underlying.getRecords(FuncName, Data);
617   }
618 };
619 }
620 
621 /// A remapper that applies remappings based on a symbol remapping file.
622 template <typename HashTableImpl>
623 class llvm::InstrProfReaderItaniumRemapper
624     : public InstrProfReaderRemapper {
625 public:
626   InstrProfReaderItaniumRemapper(
627       std::unique_ptr<MemoryBuffer> RemapBuffer,
628       InstrProfReaderIndex<HashTableImpl> &Underlying)
629       : RemapBuffer(std::move(RemapBuffer)), Underlying(Underlying) {
630   }
631 
632   /// Extract the original function name from a PGO function name.
633   static StringRef extractName(StringRef Name) {
634     // We can have multiple :-separated pieces; there can be pieces both
635     // before and after the mangled name. Find the first part that starts
636     // with '_Z'; we'll assume that's the mangled name we want.
637     std::pair<StringRef, StringRef> Parts = {StringRef(), Name};
638     while (true) {
639       Parts = Parts.second.split(':');
640       if (Parts.first.startswith("_Z"))
641         return Parts.first;
642       if (Parts.second.empty())
643         return Name;
644     }
645   }
646 
647   /// Given a mangled name extracted from a PGO function name, and a new
648   /// form for that mangled name, reconstitute the name.
649   static void reconstituteName(StringRef OrigName, StringRef ExtractedName,
650                                StringRef Replacement,
651                                SmallVectorImpl<char> &Out) {
652     Out.reserve(OrigName.size() + Replacement.size() - ExtractedName.size());
653     Out.insert(Out.end(), OrigName.begin(), ExtractedName.begin());
654     Out.insert(Out.end(), Replacement.begin(), Replacement.end());
655     Out.insert(Out.end(), ExtractedName.end(), OrigName.end());
656   }
657 
658   Error populateRemappings() override {
659     if (Error E = Remappings.read(*RemapBuffer))
660       return E;
661     for (StringRef Name : Underlying.HashTable->keys()) {
662       StringRef RealName = extractName(Name);
663       if (auto Key = Remappings.insert(RealName)) {
664         // FIXME: We could theoretically map the same equivalence class to
665         // multiple names in the profile data. If that happens, we should
666         // return NamedInstrProfRecords from all of them.
667         MappedNames.insert({Key, RealName});
668       }
669     }
670     return Error::success();
671   }
672 
673   Error getRecords(StringRef FuncName,
674                    ArrayRef<NamedInstrProfRecord> &Data) override {
675     StringRef RealName = extractName(FuncName);
676     if (auto Key = Remappings.lookup(RealName)) {
677       StringRef Remapped = MappedNames.lookup(Key);
678       if (!Remapped.empty()) {
679         if (RealName.begin() == FuncName.begin() &&
680             RealName.end() == FuncName.end())
681           FuncName = Remapped;
682         else {
683           // Try rebuilding the name from the given remapping.
684           SmallString<256> Reconstituted;
685           reconstituteName(FuncName, RealName, Remapped, Reconstituted);
686           Error E = Underlying.getRecords(Reconstituted, Data);
687           if (!E)
688             return E;
689 
690           // If we failed because the name doesn't exist, fall back to asking
691           // about the original name.
692           if (Error Unhandled = handleErrors(
693                   std::move(E), [](std::unique_ptr<InstrProfError> Err) {
694                     return Err->get() == instrprof_error::unknown_function
695                                ? Error::success()
696                                : Error(std::move(Err));
697                   }))
698             return Unhandled;
699         }
700       }
701     }
702     return Underlying.getRecords(FuncName, Data);
703   }
704 
705 private:
706   /// The memory buffer containing the remapping configuration. Remappings
707   /// holds pointers into this buffer.
708   std::unique_ptr<MemoryBuffer> RemapBuffer;
709 
710   /// The mangling remapper.
711   SymbolRemappingReader Remappings;
712 
713   /// Mapping from mangled name keys to the name used for the key in the
714   /// profile data.
715   /// FIXME: Can we store a location within the on-disk hash table instead of
716   /// redoing lookup?
717   DenseMap<SymbolRemappingReader::Key, StringRef> MappedNames;
718 
719   /// The real profile data reader.
720   InstrProfReaderIndex<HashTableImpl> &Underlying;
721 };
722 
723 bool IndexedInstrProfReader::hasFormat(const MemoryBuffer &DataBuffer) {
724   using namespace support;
725 
726   if (DataBuffer.getBufferSize() < 8)
727     return false;
728   uint64_t Magic =
729       endian::read<uint64_t, little, aligned>(DataBuffer.getBufferStart());
730   // Verify that it's magical.
731   return Magic == IndexedInstrProf::Magic;
732 }
733 
734 const unsigned char *
735 IndexedInstrProfReader::readSummary(IndexedInstrProf::ProfVersion Version,
736                                     const unsigned char *Cur) {
737   using namespace IndexedInstrProf;
738   using namespace support;
739 
740   if (Version >= IndexedInstrProf::Version4) {
741     const IndexedInstrProf::Summary *SummaryInLE =
742         reinterpret_cast<const IndexedInstrProf::Summary *>(Cur);
743     uint64_t NFields =
744         endian::byte_swap<uint64_t, little>(SummaryInLE->NumSummaryFields);
745     uint64_t NEntries =
746         endian::byte_swap<uint64_t, little>(SummaryInLE->NumCutoffEntries);
747     uint32_t SummarySize =
748         IndexedInstrProf::Summary::getSize(NFields, NEntries);
749     std::unique_ptr<IndexedInstrProf::Summary> SummaryData =
750         IndexedInstrProf::allocSummary(SummarySize);
751 
752     const uint64_t *Src = reinterpret_cast<const uint64_t *>(SummaryInLE);
753     uint64_t *Dst = reinterpret_cast<uint64_t *>(SummaryData.get());
754     for (unsigned I = 0; I < SummarySize / sizeof(uint64_t); I++)
755       Dst[I] = endian::byte_swap<uint64_t, little>(Src[I]);
756 
757     SummaryEntryVector DetailedSummary;
758     for (unsigned I = 0; I < SummaryData->NumCutoffEntries; I++) {
759       const IndexedInstrProf::Summary::Entry &Ent = SummaryData->getEntry(I);
760       DetailedSummary.emplace_back((uint32_t)Ent.Cutoff, Ent.MinBlockCount,
761                                    Ent.NumBlocks);
762     }
763     // initialize InstrProfSummary using the SummaryData from disk.
764     this->Summary = llvm::make_unique<ProfileSummary>(
765         ProfileSummary::PSK_Instr, DetailedSummary,
766         SummaryData->get(Summary::TotalBlockCount),
767         SummaryData->get(Summary::MaxBlockCount),
768         SummaryData->get(Summary::MaxInternalBlockCount),
769         SummaryData->get(Summary::MaxFunctionCount),
770         SummaryData->get(Summary::TotalNumBlocks),
771         SummaryData->get(Summary::TotalNumFunctions));
772     return Cur + SummarySize;
773   } else {
774     // For older version of profile data, we need to compute on the fly:
775     using namespace IndexedInstrProf;
776 
777     InstrProfSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
778     // FIXME: This only computes an empty summary. Need to call addRecord for
779     // all NamedInstrProfRecords to get the correct summary.
780     this->Summary = Builder.getSummary();
781     return Cur;
782   }
783 }
784 
785 Error IndexedInstrProfReader::readHeader() {
786   using namespace support;
787 
788   const unsigned char *Start =
789       (const unsigned char *)DataBuffer->getBufferStart();
790   const unsigned char *Cur = Start;
791   if ((const unsigned char *)DataBuffer->getBufferEnd() - Cur < 24)
792     return error(instrprof_error::truncated);
793 
794   auto *Header = reinterpret_cast<const IndexedInstrProf::Header *>(Cur);
795   Cur += sizeof(IndexedInstrProf::Header);
796 
797   // Check the magic number.
798   uint64_t Magic = endian::byte_swap<uint64_t, little>(Header->Magic);
799   if (Magic != IndexedInstrProf::Magic)
800     return error(instrprof_error::bad_magic);
801 
802   // Read the version.
803   uint64_t FormatVersion = endian::byte_swap<uint64_t, little>(Header->Version);
804   if (GET_VERSION(FormatVersion) >
805       IndexedInstrProf::ProfVersion::CurrentVersion)
806     return error(instrprof_error::unsupported_version);
807 
808   Cur = readSummary((IndexedInstrProf::ProfVersion)FormatVersion, Cur);
809 
810   // Read the hash type and start offset.
811   IndexedInstrProf::HashT HashType = static_cast<IndexedInstrProf::HashT>(
812       endian::byte_swap<uint64_t, little>(Header->HashType));
813   if (HashType > IndexedInstrProf::HashT::Last)
814     return error(instrprof_error::unsupported_hash_type);
815 
816   uint64_t HashOffset = endian::byte_swap<uint64_t, little>(Header->HashOffset);
817 
818   // The rest of the file is an on disk hash table.
819   auto IndexPtr =
820       llvm::make_unique<InstrProfReaderIndex<OnDiskHashTableImplV3>>(
821           Start + HashOffset, Cur, Start, HashType, FormatVersion);
822 
823   // Load the remapping table now if requested.
824   if (RemappingBuffer) {
825     Remapper = llvm::make_unique<
826         InstrProfReaderItaniumRemapper<OnDiskHashTableImplV3>>(
827         std::move(RemappingBuffer), *IndexPtr);
828     if (Error E = Remapper->populateRemappings())
829       return E;
830   } else {
831     Remapper = llvm::make_unique<InstrProfReaderNullRemapper>(*IndexPtr);
832   }
833   Index = std::move(IndexPtr);
834 
835   return success();
836 }
837 
838 InstrProfSymtab &IndexedInstrProfReader::getSymtab() {
839   if (Symtab.get())
840     return *Symtab.get();
841 
842   std::unique_ptr<InstrProfSymtab> NewSymtab = make_unique<InstrProfSymtab>();
843   if (Error E = Index->populateSymtab(*NewSymtab.get())) {
844     consumeError(error(InstrProfError::take(std::move(E))));
845   }
846 
847   Symtab = std::move(NewSymtab);
848   return *Symtab.get();
849 }
850 
851 Expected<InstrProfRecord>
852 IndexedInstrProfReader::getInstrProfRecord(StringRef FuncName,
853                                            uint64_t FuncHash) {
854   ArrayRef<NamedInstrProfRecord> Data;
855   Error Err = Remapper->getRecords(FuncName, Data);
856   if (Err)
857     return std::move(Err);
858   // Found it. Look for counters with the right hash.
859   for (unsigned I = 0, E = Data.size(); I < E; ++I) {
860     // Check for a match and fill the vector if there is one.
861     if (Data[I].Hash == FuncHash) {
862       return std::move(Data[I]);
863     }
864   }
865   return error(instrprof_error::hash_mismatch);
866 }
867 
868 Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
869                                                 uint64_t FuncHash,
870                                                 std::vector<uint64_t> &Counts) {
871   Expected<InstrProfRecord> Record = getInstrProfRecord(FuncName, FuncHash);
872   if (Error E = Record.takeError())
873     return error(std::move(E));
874 
875   Counts = Record.get().Counts;
876   return success();
877 }
878 
879 Error IndexedInstrProfReader::readNextRecord(NamedInstrProfRecord &Record) {
880   ArrayRef<NamedInstrProfRecord> Data;
881 
882   Error E = Index->getRecords(Data);
883   if (E)
884     return error(std::move(E));
885 
886   Record = Data[RecordIndex++];
887   if (RecordIndex >= Data.size()) {
888     Index->advanceToNextKey();
889     RecordIndex = 0;
890   }
891   return success();
892 }
893