1 files changed, 37 insertions, 24 deletions
diff --git a/llvm/unittests/ProfileData/SampleProfTest.cpp b/llvm/unittests/ProfileData/SampleProfTest.cpp
index 711310785c7..8d88ff3ccb7 100644
--- a/llvm/unittests/ProfileData/SampleProfTest.cpp
+++ b/llvm/unittests/ProfileData/SampleProfTest.cpp
@@ -8,6 +8,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/IR/Metadata.h"
 #include "llvm/ProfileData/SampleProfReader.h"
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include "gtest/gtest.h"
@@ -93,31 +94,43 @@ struct SampleProfTest : ::testing::Test {
     ASSERT_EQ(20301u, ReadBarSamples.getTotalSamples());
     ASSERT_EQ(1437u, ReadBarSamples.getHeadSamples());
 
-    SampleProfileSummary &Summary = Reader->getSummary();
-    ASSERT_EQ(123603u, Summary.getTotalSamples());
-    ASSERT_EQ(6u, Summary.getNumLinesWithSamples());
-    ASSERT_EQ(2u, Summary.getNumFunctions());
-    ASSERT_EQ(1437u, Summary.getMaxHeadSamples());
-    ASSERT_EQ(60351u, Summary.getMaxSamplesPerLine());
-
-    std::vector<ProfileSummaryEntry> &Details = Summary.getDetailedSummary();
-    uint32_t Cutoff = 800000;
-    auto Predicate = [&Cutoff](const ProfileSummaryEntry &PE) {
-      return PE.Cutoff == Cutoff;
+    auto VerifySummary = [](SampleProfileSummary &Summary) mutable {
+      ASSERT_EQ(123603u, Summary.getTotalSamples());
+      ASSERT_EQ(6u, Summary.getNumLinesWithSamples());
+      ASSERT_EQ(2u, Summary.getNumFunctions());
+      ASSERT_EQ(1437u, Summary.getMaxHeadSamples());
+      ASSERT_EQ(60351u, Summary.getMaxSamplesPerLine());
+
+      uint32_t Cutoff = 800000;
+      auto Predicate = [&Cutoff](const ProfileSummaryEntry &PE) {
+        return PE.Cutoff == Cutoff;
+      };
+      std::vector<ProfileSummaryEntry> &Details = Summary.getDetailedSummary();
+      auto EightyPerc = std::find_if(Details.begin(), Details.end(), Predicate);
+      Cutoff = 900000;
+      auto NinetyPerc = std::find_if(Details.begin(), Details.end(), Predicate);
+      Cutoff = 950000;
+      auto NinetyFivePerc =
+          std::find_if(Details.begin(), Details.end(), Predicate);
+      Cutoff = 990000;
+      auto NinetyNinePerc =
+          std::find_if(Details.begin(), Details.end(), Predicate);
+      ASSERT_EQ(60000u, EightyPerc->MinCount);
+      ASSERT_EQ(60000u, NinetyPerc->MinCount);
+      ASSERT_EQ(60000u, NinetyFivePerc->MinCount);
+      ASSERT_EQ(610u, NinetyNinePerc->MinCount);
     };
-    auto EightyPerc = std::find_if(Details.begin(), Details.end(), Predicate);
-    Cutoff = 900000;
-    auto NinetyPerc = std::find_if(Details.begin(), Details.end(), Predicate);
-    Cutoff = 950000;
-    auto NinetyFivePerc =
-        std::find_if(Details.begin(), Details.end(), Predicate);
-    Cutoff = 990000;
-    auto NinetyNinePerc =
-        std::find_if(Details.begin(), Details.end(), Predicate);
-    ASSERT_EQ(60000u, EightyPerc->MinCount);
-    ASSERT_EQ(60000u, NinetyPerc->MinCount);
-    ASSERT_EQ(60000u, NinetyFivePerc->MinCount);
-    ASSERT_EQ(610u, NinetyNinePerc->MinCount);
+
+    SampleProfileSummary &Summary = Reader->getSummary();
+    VerifySummary(Summary);
+
+    Metadata *MD = Summary.getMD(getGlobalContext());
+    ASSERT_TRUE(MD);
+    ProfileSummary *PS = ProfileSummary::getFromMD(MD);
+    ASSERT_TRUE(PS);
+    ASSERT_TRUE(isa<SampleProfileSummary>(PS));
+    SampleProfileSummary *SPS = cast<SampleProfileSummary>(PS);
+    VerifySummary(*SPS);
   }
 };