1 files changed, 33 insertions, 11 deletions

diff --git a/llvm/lib/ProfileData/SampleProfReader.cpp b/llvm/lib/ProfileData/SampleProfReader.cpp
index 38cbca8..089ad7b 100644
--- a/llvm/lib/ProfileData/SampleProfReader.cpp
+++ b/llvm/lib/ProfileData/SampleProfReader.cpp
@@ -584,6 +584,8 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
     bool UseMD5 = hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name);
     assert((!FixedLengthMD5 || UseMD5) &&
            "If FixedLengthMD5 is true, UseMD5 has to be true");
+    FunctionSamples::HasUniqSuffix =
+        hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix);
     if (std::error_code EC = readNameTableSec(UseMD5))
       return EC;
     break;
@@ -615,11 +617,13 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
   return sampleprof_error::success;
 }
 
-void SampleProfileReaderExtBinaryBase::collectFuncsFrom(const Module &M) {
-  UseAllFuncs = false;
+bool SampleProfileReaderExtBinaryBase::collectFuncsFromModule() {
+  if (!M)
+    return false;
   FuncsToUse.clear();
-  for (auto &F : M)
+  for (auto &F : *M)
     FuncsToUse.insert(FunctionSamples::getCanonicalFnName(F));
+  return true;
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
@@ -648,14 +652,24 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
+  // Collect functions used by current module if the Reader has been
+  // given a module.
+  // collectFuncsFromModule uses FunctionSamples::getCanonicalFnName
+  // which will query FunctionSamples::HasUniqSuffix, so it has to be
+  // called after FunctionSamples::HasUniqSuffix is set, i.e. after
+  // NameTable section is read.
+  bool LoadFuncsToBeUsed = collectFuncsFromModule();
+
+  // When LoadFuncsToBeUsed is false, load all the function profiles.
   const uint8_t *Start = Data;
-  if (UseAllFuncs) {
+  if (!LoadFuncsToBeUsed) {
     while (Data < End) {
       if (std::error_code EC = readFuncProfile(Data))
         return EC;
     }
     assert(Data == End && "More data is read than expected");
   } else {
+    // Load function profiles on demand.
     if (Remapper) {
       for (auto Name : FuncsToUse) {
         Remapper->insert(Name);
@@ -688,7 +702,6 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
     }
     Data = End;
   }
-
   assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
          "Cannot have both context-sensitive and regular profile");
   ProfileIsCS = (CSProfileCount > 0);
@@ -783,13 +796,18 @@ std::error_code SampleProfileReaderExtBinaryBase::readImpl() {
 }
 
 std::error_code SampleProfileReaderCompactBinary::readImpl() {
+  // Collect functions used by current module if the Reader has been
+  // given a module.
+  bool LoadFuncsToBeUsed = collectFuncsFromModule();
+
   std::vector<uint64_t> OffsetsToUse;
-  if (UseAllFuncs) {
+  if (!LoadFuncsToBeUsed) {
+    // load all the function profiles.
     for (auto FuncEntry : FuncOffsetTable) {
       OffsetsToUse.push_back(FuncEntry.second);
     }
-  }
-  else {
+  } else {
+    // load function profiles on demand.
     for (auto Name : FuncsToUse) {
       auto GUID = std::to_string(MD5Hash(Name));
       auto iter = FuncOffsetTable.find(StringRef(GUID));
@@ -1010,6 +1028,8 @@ static std::string getSecFlagsStr(const SecHdrTableEntry &Entry) {
       Flags.append("fixlenmd5,");
     else if (hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name))
       Flags.append("md5,");
+    if (hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix))
+      Flags.append("uniq,");
     break;
   case SecProfSummary:
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagPartial))
@@ -1117,11 +1137,13 @@ std::error_code SampleProfileReaderCompactBinary::readFuncOffsetTable() {
   return sampleprof_error::success;
 }
 
-void SampleProfileReaderCompactBinary::collectFuncsFrom(const Module &M) {
-  UseAllFuncs = false;
+bool SampleProfileReaderCompactBinary::collectFuncsFromModule() {
+  if (!M)
+    return false;
   FuncsToUse.clear();
-  for (auto &F : M)
+  for (auto &F : *M)
     FuncsToUse.insert(FunctionSamples::getCanonicalFnName(F));
+  return true;
 }
 
 std::error_code SampleProfileReaderBinary::readSummaryEntry(