1 files changed, 18 insertions, 12 deletions
diff --git a/llvm/lib/ProfileData/CoverageMapping.cpp b/llvm/lib/ProfileData/CoverageMapping.cpp
index 334de3c3b06..e130565500f 100644
--- a/llvm/lib/ProfileData/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/CoverageMapping.cpp
@@ -375,26 +375,30 @@ static SmallBitVector gatherFileIDs(StringRef SourceFile,
   return FilenameEquivalence;
 }
 
-/// Return the ID of the file where the definition of the function is located.
-static Optional<unsigned> findMainViewFileID(const FunctionRecord &Function) {
+static Optional<unsigned> findMainViewFileID(StringRef SourceFile,
+                                             const FunctionRecord &Function) {
   SmallBitVector IsNotExpandedFile(Function.Filenames.size(), true);
+  SmallBitVector FilenameEquivalence = gatherFileIDs(SourceFile, Function);
   for (const auto &CR : Function.CountedRegions)
-    if (CR.Kind == CounterMappingRegion::ExpansionRegion)
+    if (CR.Kind == CounterMappingRegion::ExpansionRegion &&
+        FilenameEquivalence[CR.FileID])
       IsNotExpandedFile[CR.ExpandedFileID] = false;
+  IsNotExpandedFile &= FilenameEquivalence;
   int I = IsNotExpandedFile.find_first();
   if (I == -1)
     return None;
   return I;
 }
 
-/// Check if SourceFile is the file that contains the definition of
-/// the Function. Return the ID of the file in that case or None otherwise.
-static Optional<unsigned> findMainViewFileID(StringRef SourceFile,
-                                             const FunctionRecord &Function) {
-  Optional<unsigned> I = findMainViewFileID(Function);
-  if (I && SourceFile == Function.Filenames[*I])
-    return I;
-  return None;
+static Optional<unsigned> findMainViewFileID(const FunctionRecord &Function) {
+  SmallBitVector IsNotExpandedFile(Function.Filenames.size(), true);
+  for (const auto &CR : Function.CountedRegions)
+    if (CR.Kind == CounterMappingRegion::ExpansionRegion)
+      IsNotExpandedFile[CR.ExpandedFileID] = false;
+  int I = IsNotExpandedFile.find_first();
+  if (I == -1)
+    return None;
+  return I;
 }
 
 /// Sort a nested sequence of regions from a single file.
@@ -418,11 +422,13 @@ CoverageData CoverageMapping::getCoverageForFile(StringRef Filename) {
 
   for (const auto &Function : Functions) {
     auto MainFileID = findMainViewFileID(Filename, Function);
+    if (!MainFileID)
+      continue;
     auto FileIDs = gatherFileIDs(Filename, Function);
     for (const auto &CR : Function.CountedRegions)
       if (FileIDs.test(CR.FileID)) {
         Regions.push_back(CR);
-        if (MainFileID && isExpansion(CR, *MainFileID))
+        if (isExpansion(CR, *MainFileID))
           FileCoverage.Expansions.emplace_back(CR, Function);
       }
   }