7 files changed, 102 insertions, 19 deletions
diff --git a/llvm/lib/ProfileData/CoverageMapping.cpp b/llvm/lib/ProfileData/CoverageMapping.cpp
index e130565500f..334de3c3b06 100644
--- a/llvm/lib/ProfileData/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/CoverageMapping.cpp
@@ -375,21 +375,7 @@ static SmallBitVector gatherFileIDs(StringRef SourceFile,
   return FilenameEquivalence;
 }
 
-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 &&
-        FilenameEquivalence[CR.FileID])
-      IsNotExpandedFile[CR.ExpandedFileID] = false;
-  IsNotExpandedFile &= FilenameEquivalence;
-  int I = IsNotExpandedFile.find_first();
-  if (I == -1)
-    return None;
-  return I;
-}
-
+/// Return the ID of the file where the definition of the function is located.
 static Optional<unsigned> findMainViewFileID(const FunctionRecord &Function) {
   SmallBitVector IsNotExpandedFile(Function.Filenames.size(), true);
   for (const auto &CR : Function.CountedRegions)
@@ -401,6 +387,16 @@ static Optional<unsigned> findMainViewFileID(const FunctionRecord &Function) {
   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;
+}
+
 /// Sort a nested sequence of regions from a single file.
 template <class It> static void sortNestedRegions(It First, It Last) {
   std::sort(First, Last,
@@ -422,13 +418,11 @@ 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 (isExpansion(CR, *MainFileID))
+        if (MainFileID && 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
new file mode 100644
index 00000000000..c255aa3272f
--- /dev/null
+++ b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.covmapping
diff --git a/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.cpp b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.cpp
new file mode 100644
index 00000000000..2a84f1dd58e
--- /dev/null
+++ b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.cpp
@@ -0,0 +1,15 @@
+#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
new file mode 100644
index 00000000000..ad5362e17e4
--- /dev/null
+++ b/llvm/test/tools/llvm-cov/Inputs/prevent_false_instantiations.proftext
@@ -0,0 +1,26 @@
+_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 05aef8645e4..8b6fe62772f 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']
+config.suffixes = ['.test', '.m', '.cpp', '.c', '.h']
diff --git a/llvm/test/tools/llvm-cov/prevent_false_instantiations.h b/llvm/test/tools/llvm-cov/prevent_false_instantiations.h
new file mode 100644
index 00000000000..667ebde6d16
--- /dev/null
+++ b/llvm/test/tools/llvm-cov/prevent_false_instantiations.h
@@ -0,0 +1,10 @@
+// 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 f199cc7a2b2..6535d68dadc 100644
--- a/llvm/unittests/ProfileData/CoverageMappingTest.cpp
+++ b/llvm/unittests/ProfileData/CoverageMappingTest.cpp
@@ -454,6 +454,44 @@ 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());