1 files changed, 39 insertions, 26 deletions

diff --git a/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp b/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
index 56a7ab2..ac8e6b5 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
@@ -493,9 +493,13 @@ Error InstrProfSymtab::create(SectionRef &Section) {
 
   // If this is a linked PE/COFF file, then we have to skip over the null byte
   // that is allocated in the .lprfn$A section in the LLVM profiling runtime.
+  // If the name section is .lprfcovnames, it doesn't have the null byte at the
+  // beginning.
   const ObjectFile *Obj = Section.getObject();
   if (isa<COFFObjectFile>(Obj) && !Obj->isRelocatableObject())
-    Data = Data.drop_front(1);
+    if (Expected<StringRef> NameOrErr = Section.getName())
+      if (*NameOrErr != getInstrProfSectionName(IPSK_covname, Triple::COFF))
+        Data = Data.drop_front(1);
 
   return Error::success();
 }
@@ -1024,10 +1028,13 @@ loadTestingFormat(StringRef Data, StringRef CompilationDir) {
       BytesInAddress, Endian, CompilationDir);
 }
 
-/// Find all sections that match \p Name. There may be more than one if comdats
-/// are in use, e.g. for the __llvm_covfun section on ELF.
-static Expected<std::vector<SectionRef>> lookupSections(ObjectFile &OF,
-                                                        StringRef Name) {
+/// Find all sections that match \p IPSK name. There may be more than one if
+/// comdats are in use, e.g. for the __llvm_covfun section on ELF.
+static Expected<std::vector<SectionRef>>
+lookupSections(ObjectFile &OF, InstrProfSectKind IPSK) {
+  auto ObjFormat = OF.getTripleObjectFormat();
+  auto Name =
+      getInstrProfSectionName(IPSK, ObjFormat, /*AddSegmentInfo=*/false);
   // On COFF, the object file section name may end in "$M". This tells the
   // linker to sort these sections between "$A" and "$Z". The linker removes the
   // dollar and everything after it in the final binary. Do the same to match.
@@ -1042,8 +1049,13 @@ static Expected<std::vector<SectionRef>> lookupSections(ObjectFile &OF,
     Expected<StringRef> NameOrErr = Section.getName();
     if (!NameOrErr)
       return NameOrErr.takeError();
-    if (stripSuffix(*NameOrErr) == Name)
+    if (stripSuffix(*NameOrErr) == Name) {
+      // COFF profile name section contains two null bytes indicating the
+      // start/end of the section. If its size is 2 bytes, it's empty.
+      if (IsCOFF && IPSK == IPSK_name && Section.getSize() == 2)
+        continue;
       Sections.push_back(Section);
+    }
   }
   if (Sections.empty())
     return make_error<CoverageMapError>(coveragemap_error::no_data_found);
@@ -1079,15 +1091,27 @@ loadBinaryFormat(std::unique_ptr<Binary> Bin, StringRef Arch,
       OF->isLittleEndian() ? llvm::endianness::little : llvm::endianness::big;
 
   // Look for the sections that we are interested in.
-  auto ObjFormat = OF->getTripleObjectFormat();
-  auto NamesSection =
-      lookupSections(*OF, getInstrProfSectionName(IPSK_name, ObjFormat,
-                                                 /*AddSegmentInfo=*/false));
-  if (auto E = NamesSection.takeError())
+  InstrProfSymtab ProfileNames;
+  std::vector<SectionRef> NamesSectionRefs;
+  // If IPSK_name is not found, fallback to search for IPK_covname, which is
+  // used when binary correlation is enabled.
+  auto NamesSection = lookupSections(*OF, IPSK_name);
+  if (auto E = NamesSection.takeError()) {
+    consumeError(std::move(E));
+    NamesSection = lookupSections(*OF, IPSK_covname);
+    if (auto E = NamesSection.takeError())
+      return std::move(E);
+  }
+  NamesSectionRefs = *NamesSection;
+
+  if (NamesSectionRefs.size() != 1)
+    return make_error<CoverageMapError>(
+        coveragemap_error::malformed,
+        "the size of coverage mapping section is not one");
+  if (Error E = ProfileNames.create(NamesSectionRefs.back()))
     return std::move(E);
-  auto CoverageSection =
-      lookupSections(*OF, getInstrProfSectionName(IPSK_covmap, ObjFormat,
-                                                  /*AddSegmentInfo=*/false));
+
+  auto CoverageSection = lookupSections(*OF, IPSK_covmap);
   if (auto E = CoverageSection.takeError())
     return std::move(E);
   std::vector<SectionRef> CoverageSectionRefs = *CoverageSection;
@@ -1099,19 +1123,8 @@ loadBinaryFormat(std::unique_ptr<Binary> Bin, StringRef Arch,
     return CoverageMappingOrErr.takeError();
   StringRef CoverageMapping = CoverageMappingOrErr.get();
 
-  InstrProfSymtab ProfileNames;
-  std::vector<SectionRef> NamesSectionRefs = *NamesSection;
-  if (NamesSectionRefs.size() != 1)
-    return make_error<CoverageMapError>(
-        coveragemap_error::malformed,
-        "the size of coverage mapping section is not one");
-  if (Error E = ProfileNames.create(NamesSectionRefs.back()))
-    return std::move(E);
-
   // Look for the coverage records section (Version4 only).
-  auto CoverageRecordsSections =
-      lookupSections(*OF, getInstrProfSectionName(IPSK_covfun, ObjFormat,
-                                                  /*AddSegmentInfo=*/false));
+  auto CoverageRecordsSections = lookupSections(*OF, IPSK_covfun);
 
   BinaryCoverageReader::FuncRecordsStorage FuncRecords;
   if (auto E = CoverageRecordsSections.takeError()) {