[PGO] Indexed Prof Reader refactoring (NFC)

Change InstrProfReaderIndex from typedef into a wrapper
class with helper methods. This makes the index profile
reader code more readable. It also hides the implementation 
detail of the index format and make it more flexible to allow 
support different (or more than one) format in the future.

llvm-svn: 251491

2 files changed, 90 insertions, 49 deletions

diff --git a/llvm/include/llvm/ProfileData/InstrProfReader.h b/llvm/include/llvm/ProfileData/InstrProfReader.h
index cdead0f112d..e4dfc386000 100644
--- a/llvm/include/llvm/ProfileData/InstrProfReader.h
+++ b/llvm/include/llvm/ProfileData/InstrProfReader.h
@@ -64,10 +64,7 @@ public:
   InstrProfIterator begin() { return InstrProfIterator(this); }
   InstrProfIterator end() { return InstrProfIterator(); }
 
-protected:
-  /// String table for holding a unique copy of all the strings in the profile.
-  InstrProfStringTable StringTable;
-
+ protected:
   /// Set the current std::error_code and return same.
   std::error_code error(std::error_code EC) {
     LastError = EC;
@@ -221,8 +218,33 @@ public:
   data_type ReadData(StringRef K, const unsigned char *D, offset_type N);
 };
 
-typedef OnDiskIterableChainedHashTable<InstrProfLookupTrait>
-    InstrProfReaderIndex;
+class InstrProfReaderIndex {
+ private:
+  typedef OnDiskIterableChainedHashTable<InstrProfLookupTrait> IndexType;
+
+  std::unique_ptr<IndexType> Index;
+  IndexType::data_iterator RecordIterator;
+  uint64_t FormatVersion;
+
+  // String table for holding a unique copy of all the strings in the profile.
+  InstrProfStringTable StringTable;
+
+ public:
+  InstrProfReaderIndex() : Index(nullptr) {}
+  void Init(const unsigned char *Buckets, const unsigned char *const Payload,
+            const unsigned char *const Base, IndexedInstrProf::HashT HashType,
+            uint64_t Version);
+
+  // Read all the pofile records with the same key pointed to the current
+  // iterator.
+  std::error_code getRecords(ArrayRef<InstrProfRecord> &Data);
+  // Read all the profile records with the key equal to FuncName
+  std::error_code getRecords(StringRef FuncName,
+                             ArrayRef<InstrProfRecord> &Data);
+
+  void advanceToNextKey() { RecordIterator++; }
+  bool atEnd() const { return RecordIterator == Index->data_end(); }
+};
 
 /// Reader for the indexed binary instrprof format.
 class IndexedInstrProfReader : public InstrProfReader {
@@ -230,19 +252,16 @@ private:
   /// The profile data file contents.
   std::unique_ptr<MemoryBuffer> DataBuffer;
   /// The index into the profile data.
-  std::unique_ptr<InstrProfReaderIndex> Index;
-  /// Iterator over the profile data.
-  InstrProfReaderIndex::data_iterator RecordIterator;
-  /// The file format version of the profile data.
-  uint64_t FormatVersion;
+  InstrProfReaderIndex Index;
   /// The maximal execution count among all functions.
   uint64_t MaxFunctionCount;
 
   IndexedInstrProfReader(const IndexedInstrProfReader &) = delete;
   IndexedInstrProfReader &operator=(const IndexedInstrProfReader &) = delete;
-public:
+
+ public:
   IndexedInstrProfReader(std::unique_ptr<MemoryBuffer> DataBuffer)
-      : DataBuffer(std::move(DataBuffer)), Index(nullptr) {}
+      : DataBuffer(std::move(DataBuffer)), Index() {}
 
   /// Return true if the given buffer is in an indexed instrprof format.
   static bool hasFormat(const MemoryBuffer &DataBuffer);
diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index 61d1f11076c..838ec0f29ca 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -370,9 +370,51 @@ data_type InstrProfLookupTrait::ReadData(StringRef K, const unsigned char *D,
   return DataBuffer;
 }
 
+std::error_code InstrProfReaderIndex::getRecords(
+    StringRef FuncName, ArrayRef<InstrProfRecord> &Data) {
+  auto Iter = Index->find(FuncName);
+  if (Iter == Index->end()) return instrprof_error::unknown_function;
+
+  Data = (*Iter);
+  if (Data.empty()) return instrprof_error::malformed;
+
+  return instrprof_error::success;
+}
+
+std::error_code InstrProfReaderIndex::getRecords(
+    ArrayRef<InstrProfRecord> &Data) {
+  if (atEnd()) return instrprof_error::eof;
+
+  Data = *RecordIterator;
+
+  if (Data.empty()) return instrprof_error::malformed;
+
+  return instrprof_error::success;
+}
+
+void InstrProfReaderIndex::Init(const unsigned char *Buckets,
+                                const unsigned char *const Payload,
+                                const unsigned char *const Base,
+                                IndexedInstrProf::HashT HashType,
+                                uint64_t Version) {
+  FormatVersion = Version;
+  Index.reset(IndexType::Create(Buckets, Payload, Base,
+                                InstrProfLookupTrait(HashType, Version)));
+  // Form the map of hash values to const char* keys in profiling data.
+  std::vector<std::pair<uint64_t, const char *>> HashKeys;
+  for (auto Key : Index->keys()) {
+    const char *KeyTableRef = StringTable.insertString(Key);
+    HashKeys.push_back(std::make_pair(ComputeHash(HashType, Key), KeyTableRef));
+  }
+  std::sort(HashKeys.begin(), HashKeys.end(), less_first());
+  HashKeys.erase(std::unique(HashKeys.begin(), HashKeys.end()), HashKeys.end());
+  // Set the hash key map for the InstrLookupTrait
+  Index->getInfoObj().setHashKeys(std::move(HashKeys));
+  RecordIterator = Index->data_begin();
+}
+
 bool IndexedInstrProfReader::hasFormat(const MemoryBuffer &DataBuffer) {
-  if (DataBuffer.getBufferSize() < 8)
-    return false;
+  if (DataBuffer.getBufferSize() < 8) return false;
   using namespace support;
   uint64_t Magic =
       endian::read<uint64_t, little, aligned>(DataBuffer.getBufferStart());
@@ -398,7 +440,7 @@ std::error_code IndexedInstrProfReader::readHeader() {
     return error(instrprof_error::bad_magic);
 
   // Read the version.
-  FormatVersion = endian::byte_swap<uint64_t, little>(Header->Version);
+  uint64_t FormatVersion = endian::byte_swap<uint64_t, little>(Header->Version);
   if (FormatVersion > IndexedInstrProf::Version)
     return error(instrprof_error::unsupported_version);
 
@@ -415,37 +457,19 @@ std::error_code IndexedInstrProfReader::readHeader() {
   uint64_t HashOffset = endian::byte_swap<uint64_t, little>(Header->HashOffset);
 
   // The rest of the file is an on disk hash table.
-  Index.reset(InstrProfReaderIndex::Create(
-      Start + HashOffset, Cur, Start,
-      InstrProfLookupTrait(HashType, FormatVersion)));
-
-  // Form the map of hash values to const char* keys in profiling data.
-  std::vector<std::pair<uint64_t, const char *>> HashKeys;
-  for (auto Key : Index->keys()) {
-    const char *KeyTableRef = StringTable.insertString(Key);
-    HashKeys.push_back(std::make_pair(ComputeHash(HashType, Key), KeyTableRef));
-  }
-  std::sort(HashKeys.begin(), HashKeys.end(), less_first());
-  HashKeys.erase(std::unique(HashKeys.begin(), HashKeys.end()), HashKeys.end());
-  // Set the hash key map for the InstrLookupTrait
-  Index->getInfoObj().setHashKeys(std::move(HashKeys));
-  // Set up our iterator for readNextRecord.
-  RecordIterator = Index->data_begin();
+  Index.Init(Start + HashOffset, Cur, Start, HashType, FormatVersion);
 
   return success();
 }
 
 std::error_code IndexedInstrProfReader::getFunctionCounts(
     StringRef FuncName, uint64_t FuncHash, std::vector<uint64_t> &Counts) {
-  auto Iter = Index->find(FuncName);
-  if (Iter == Index->end())
-    return error(instrprof_error::unknown_function);
+  ArrayRef<InstrProfRecord> Data;
 
-  // Found it. Look for counters with the right hash.
-  ArrayRef<InstrProfRecord> Data = (*Iter);
-  if (Data.empty())
-    return error(instrprof_error::malformed);
+  std::error_code EC = Index.getRecords(FuncName, Data);
+  if (EC != instrprof_error::success) return EC;
 
+  // Found it. Look for counters with the right hash.
   for (unsigned I = 0, E = Data.size(); I < E; ++I) {
     // Check for a match and fill the vector if there is one.
     if (Data[I].Hash == FuncHash) {
@@ -456,20 +480,18 @@ std::error_code IndexedInstrProfReader::getFunctionCounts(
   return error(instrprof_error::hash_mismatch);
 }
 
-std::error_code
-IndexedInstrProfReader::readNextRecord(InstrProfRecord &Record) {
-  // Are we out of records?
-  if (RecordIterator == Index->data_end())
-    return error(instrprof_error::eof);
+std::error_code IndexedInstrProfReader::readNextRecord(
+    InstrProfRecord &Record) {
+  static unsigned RecordIndex = 0;
 
-  if ((*RecordIterator).empty())
-    return error(instrprof_error::malformed);
+  ArrayRef<InstrProfRecord> Data;
+
+  std::error_code EC = Index.getRecords(Data);
+  if (EC != instrprof_error::success) return error(EC);
 
-  static unsigned RecordIndex = 0;
-  ArrayRef<InstrProfRecord> Data = (*RecordIterator);
   Record = Data[RecordIndex++];
   if (RecordIndex >= Data.size()) {
-    ++RecordIterator;
+    Index.advanceToNextKey();
     RecordIndex = 0;
   }
   return success();