1 files changed, 28 insertions, 8 deletions

diff --git a/llvm/unittests/ProfileData/SampleProfTest.cpp b/llvm/unittests/ProfileData/SampleProfTest.cpp
index dd5ded5bc9b..59ed19d292e 100644
--- a/llvm/unittests/ProfileData/SampleProfTest.cpp
+++ b/llvm/unittests/ProfileData/SampleProfTest.cpp
@@ -54,7 +54,7 @@ struct SampleProfTest : ::testing::Test {
     auto ReaderOrErr = SampleProfileReader::create(Profile, Context);
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
-    Reader->collectFuncsToUse(M);
+    Reader->collectFuncsFrom(M);
   }
 
   void testRoundTrip(SampleProfileFormat Format, bool Remap) {
@@ -86,6 +86,13 @@ struct SampleProfTest : ::testing::Test {
     BarSamples.addCalledTargetSamples(1, 0, MconstructName, 1000);
     BarSamples.addCalledTargetSamples(1, 0, StringviewName, 437);
 
+    StringRef BazName("_Z3bazi");
+    FunctionSamples BazSamples;
+    BazSamples.setName(BazName);
+    BazSamples.addTotalSamples(12557);
+    BazSamples.addHeadSamples(1257);
+    BazSamples.addBodySamples(1, 0, 12557);
+
     Module M("my_module", Context);
     FunctionType *fn_type =
         FunctionType::get(Type::getVoidTy(Context), {}, false);
@@ -95,6 +102,7 @@ struct SampleProfTest : ::testing::Test {
     StringMap<FunctionSamples> Profiles;
     Profiles[FooName] = std::move(FooSamples);
     Profiles[BarName] = std::move(BarSamples);
+    Profiles[BazName] = std::move(BazSamples);
 
     ProfileSymbolList List;
     if (Format == SampleProfileFormat::SPF_Ext_Binary) {
@@ -137,8 +145,6 @@ struct SampleProfTest : ::testing::Test {
       ASSERT_TRUE(NoError(EC));
     }
 
-    ASSERT_EQ(2u, Reader->getProfiles().size());
-
     FunctionSamples *ReadFooSamples = Reader->getSamplesFor(FooName);
     ASSERT_TRUE(ReadFooSamples != nullptr);
     if (Format != SampleProfileFormat::SPF_Compact_Binary) {
@@ -158,6 +164,20 @@ struct SampleProfTest : ::testing::Test {
         ReadBarSamples->findCallTargetMapAt(1, 0);
     ASSERT_FALSE(CTMap.getError());
 
+    // Because _Z3bazi is not defined in module M, expect _Z3bazi's profile
+    // is not loaded when the profile is ExtBinary or Compact format because
+    // these formats support loading function profiles on demand.
+    FunctionSamples *ReadBazSamples = Reader->getSamplesFor(BazName);
+    if (Format == SampleProfileFormat::SPF_Ext_Binary ||
+        Format == SampleProfileFormat::SPF_Compact_Binary) {
+      ASSERT_TRUE(ReadBazSamples == nullptr);
+      ASSERT_EQ(2u, Reader->getProfiles().size());
+    } else {
+      ASSERT_TRUE(ReadBazSamples != nullptr);
+      ASSERT_EQ(12557u, ReadBazSamples->getTotalSamples());
+      ASSERT_EQ(3u, Reader->getProfiles().size());
+    }
+
     std::string MconstructGUID;
     StringRef MconstructRep =
         getRepInFormat(MconstructName, Format, MconstructGUID);
@@ -169,9 +189,9 @@ struct SampleProfTest : ::testing::Test {
 
     auto VerifySummary = [](ProfileSummary &Summary) mutable {
       ASSERT_EQ(ProfileSummary::PSK_Sample, Summary.getKind());
-      ASSERT_EQ(123603u, Summary.getTotalCount());
-      ASSERT_EQ(6u, Summary.getNumCounts());
-      ASSERT_EQ(2u, Summary.getNumFunctions());
+      ASSERT_EQ(136160u, Summary.getTotalCount());
+      ASSERT_EQ(7u, Summary.getNumCounts());
+      ASSERT_EQ(3u, Summary.getNumFunctions());
       ASSERT_EQ(1437u, Summary.getMaxFunctionCount());
       ASSERT_EQ(60351u, Summary.getMaxCount());
 
@@ -188,8 +208,8 @@ struct SampleProfTest : ::testing::Test {
       Cutoff = 990000;
       auto NinetyNinePerc = find_if(Details, Predicate);
       ASSERT_EQ(60000u, EightyPerc->MinCount);
-      ASSERT_EQ(60000u, NinetyPerc->MinCount);
-      ASSERT_EQ(60000u, NinetyFivePerc->MinCount);
+      ASSERT_EQ(12557u, NinetyPerc->MinCount);
+      ASSERT_EQ(12557u, NinetyFivePerc->MinCount);
       ASSERT_EQ(610u, NinetyNinePerc->MinCount);
     };