7 files changed, 19 insertions, 102 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);
       }
   }
diff --git a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.covmapping b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.covmapping
deleted file mode 100644
index c255aa3272f..00000000000
--- a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.covmapping
+++ /dev/null
diff --git a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.cpp b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.cpp
deleted file mode 100644
index 2a84f1dd58e..00000000000
--- a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-#include "prevent_false_instantiations.h"
-
-void func1() {
-  DO_SOMETHING();
-}
-
-void func2() {
-  DO_SOMETHING();
-}
-
-int main() {
-  func1();
-  func2();
-  return 0;
-}
diff --git a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.proftext b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.proftext
deleted file mode 100644
index ad5362e17e4..00000000000
--- a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.proftext
+++ /dev/null
@@ -1,26 +0,0 @@
-_Z5func1v
-# Func Hash:
-3
-# Num Counters:
-2
-# Counter Values:
-1
-0
-
-_Z5func2v
-# Func Hash:
-3
-# Num Counters:
-2
-# Counter Values:
-1
-0
-
-main
-# Func Hash:
-0
-# Num Counters:
-1
-# Counter Values:
-1
-
diff --git a/llvm/test/tools/llvm-cov/lit.local.cfg b/llvm/test/tools/llvm-cov/lit.local.cfg
index 8b6fe62772f..05aef8645e4 100644
--- a/llvm/test/tools/llvm-cov/lit.local.cfg
+++ b/llvm/test/tools/llvm-cov/lit.local.cfg
@@ -1 +1 @@
-config.suffixes = ['.test', '.m', '.cpp', '.c', '.h']
+config.suffixes = ['.test', '.m', '.cpp', '.c']
diff --git a/llvm/test/tools/llvm-cov/prevent_false_instantiations.h b/llvm/test/tools/llvm-cov/prevent_false_instantiations.h
deleted file mode 100644
index 667ebde6d16..00000000000
--- a/llvm/test/tools/llvm-cov/prevent_false_instantiations.h
+++ /dev/null
@@ -1,10 +0,0 @@
-// Checks that function instantiations don't go to a wrong file.
-
-// CHECK-NOT: {{_Z5func[1,2]v}}
-
-// RUN: llvm-profdata merge %S/Inputs/prevent_false_instantiations.proftext -o %t.profdata
-// RUN: llvm-cov show %S/Inputs/prevent_false_instantiations.covmapping -instr-profile %t.profdata -filename-equivalence %s | FileCheck %s
-
-#define DO_SOMETHING() \
-  do {                 \
-  } while (0)
diff --git a/llvm/unittests/ProfileData/CoverageMappingTest.cpp b/llvm/unittests/ProfileData/CoverageMappingTest.cpp
index 6535d68dadc..f199cc7a2b2 100644
--- a/llvm/unittests/ProfileData/CoverageMappingTest.cpp
+++ b/llvm/unittests/ProfileData/CoverageMappingTest.cpp
@@ -454,44 +454,6 @@ TEST_P(MaybeSparseCoverageMappingTest, strip_unknown_filename_prefix) {
   ASSERT_EQ("func", Names[0]);
 }
 
-TEST_P(MaybeSparseCoverageMappingTest, dont_detect_false_instantiations) {
-  InstrProfRecord Record1("foo", 0x1234, {10});
-  InstrProfRecord Record2("bar", 0x2345, {20});
-  ProfileWriter.addRecord(std::move(Record1));
-  ProfileWriter.addRecord(std::move(Record2));
-
-  startFunction("foo", 0x1234);
-  addCMR(Counter::getCounter(0), "expanded", 1, 1, 1, 10);
-  addExpansionCMR("main", "expanded", 4, 1, 4, 5);
-
-  startFunction("bar", 0x2345);
-  addCMR(Counter::getCounter(0), "expanded", 1, 1, 1, 10);
-  addExpansionCMR("main", "expanded", 9, 1, 9, 5);
-
-  loadCoverageMapping();
-
-  std::vector<const FunctionRecord *> Instantiations =
-      LoadedCoverage->getInstantiations("expanded");
-  ASSERT_TRUE(Instantiations.empty());
-}
-
-TEST_P(MaybeSparseCoverageMappingTest, load_coverage_for_expanded_file) {
-  InstrProfRecord Record("func", 0x1234, {10});
-  ProfileWriter.addRecord(std::move(Record));
-
-  startFunction("func", 0x1234);
-  addCMR(Counter::getCounter(0), "expanded", 1, 1, 1, 10);
-  addExpansionCMR("main", "expanded", 4, 1, 4, 5);
-
-  loadCoverageMapping();
-
-  CoverageData Data = LoadedCoverage->getCoverageForFile("expanded");
-  std::vector<CoverageSegment> Segments(Data.begin(), Data.end());
-  ASSERT_EQ(2U, Segments.size());
-  EXPECT_EQ(CoverageSegment(1, 1, 10, true), Segments[0]);
-  EXPECT_EQ(CoverageSegment(1, 10, false), Segments[1]);
-}
-
 INSTANTIATE_TEST_CASE_P(MaybeSparse, MaybeSparseCoverageMappingTest,
                         ::testing::Bool());