1 files changed, 53 insertions, 30 deletions

diff --git a/llvm/unittests/ProfileData/SampleProfTest.cpp b/llvm/unittests/ProfileData/SampleProfTest.cpp
index 59ed19d292e..816f403f570 100644
--- a/llvm/unittests/ProfileData/SampleProfTest.cpp
+++ b/llvm/unittests/ProfileData/SampleProfTest.cpp
@@ -50,13 +50,31 @@ struct SampleProfTest : ::testing::Test {
     Writer = std::move(WriterOrErr.get());
   }
 
-  void readProfile(const Module &M, StringRef Profile) {
-    auto ReaderOrErr = SampleProfileReader::create(Profile, Context);
+  void readProfile(const Module &M, StringRef Profile,
+                   StringRef RemapFile = "") {
+    auto ReaderOrErr = SampleProfileReader::create(Profile, Context, RemapFile);
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
     Reader->collectFuncsFrom(M);
   }
 
+  void createRemapFile(SmallVectorImpl<char> &RemapPath, StringRef &RemapFile) {
+    std::error_code EC =
+        llvm::sys::fs::createTemporaryFile("remapfile", "", RemapPath);
+    ASSERT_TRUE(NoError(EC));
+    RemapFile = StringRef(RemapPath.data(), RemapPath.size());
+
+    std::unique_ptr<raw_fd_ostream> OS(
+        new raw_fd_ostream(RemapFile, EC, sys::fs::OF_None));
+    *OS << R"(
+      # Types 'int' and 'long' are equivalent
+      type i l
+      # Function names 'foo' and 'faux' are equivalent
+      name 3foo 4faux
+    )";
+    OS->close();
+  }
+
   void testRoundTrip(SampleProfileFormat Format, bool Remap) {
     SmallVector<char, 128> ProfilePath;
     ASSERT_TRUE(NoError(llvm::sys::fs::createTemporaryFile("profile", "", ProfilePath)));
@@ -93,16 +111,34 @@ struct SampleProfTest : ::testing::Test {
     BazSamples.addHeadSamples(1257);
     BazSamples.addBodySamples(1, 0, 12557);
 
-    Module M("my_module", Context);
-    FunctionType *fn_type =
-        FunctionType::get(Type::getVoidTy(Context), {}, false);
-    M.getOrInsertFunction(FooName, fn_type);
-    M.getOrInsertFunction(BarName, fn_type);
+    StringRef BooName("_Z3booi");
+    FunctionSamples BooSamples;
+    BooSamples.setName(BooName);
+    BooSamples.addTotalSamples(1232);
+    BooSamples.addHeadSamples(1);
+    BooSamples.addBodySamples(1, 0, 1232);
 
     StringMap<FunctionSamples> Profiles;
     Profiles[FooName] = std::move(FooSamples);
     Profiles[BarName] = std::move(BarSamples);
     Profiles[BazName] = std::move(BazSamples);
+    Profiles[BooName] = std::move(BooSamples);
+
+    Module M("my_module", Context);
+    FunctionType *fn_type =
+        FunctionType::get(Type::getVoidTy(Context), {}, false);
+
+    SmallVector<char, 128> RemapPath;
+    StringRef RemapFile;
+    if (Remap) {
+      createRemapFile(RemapPath, RemapFile);
+      FooName = "_Z4fauxi";
+      BarName = "_Z3barl";
+    }
+
+    M.getOrInsertFunction(FooName, fn_type);
+    M.getOrInsertFunction(BarName, fn_type);
+    M.getOrInsertFunction(BooName, fn_type);
 
     ProfileSymbolList List;
     if (Format == SampleProfileFormat::SPF_Ext_Binary) {
@@ -117,8 +153,7 @@ struct SampleProfTest : ::testing::Test {
 
     Writer->getOutputStream().flush();
 
-    readProfile(M, Profile);
-
+    readProfile(M, Profile, RemapFile);
     EC = Reader->read();
     ASSERT_TRUE(NoError(EC));
 
@@ -129,22 +164,6 @@ struct SampleProfTest : ::testing::Test {
       ReaderList->contains("moo");
     }
 
-    if (Remap) {
-      auto MemBuffer = llvm::MemoryBuffer::getMemBuffer(R"(
-        # Types 'int' and 'long' are equivalent
-        type i l
-        # Function names 'foo' and 'faux' are equivalent
-        name 3foo 4faux
-      )");
-      Reader.reset(new SampleProfileReaderItaniumRemapper(
-          std::move(MemBuffer), Context, std::move(Reader)));
-      FooName = "_Z4fauxi";
-      BarName = "_Z3barl";
-
-      EC = Reader->read();
-      ASSERT_TRUE(NoError(EC));
-    }
-
     FunctionSamples *ReadFooSamples = Reader->getSamplesFor(FooName);
     ASSERT_TRUE(ReadFooSamples != nullptr);
     if (Format != SampleProfileFormat::SPF_Compact_Binary) {
@@ -171,13 +190,17 @@ struct SampleProfTest : ::testing::Test {
     if (Format == SampleProfileFormat::SPF_Ext_Binary ||
         Format == SampleProfileFormat::SPF_Compact_Binary) {
       ASSERT_TRUE(ReadBazSamples == nullptr);
-      ASSERT_EQ(2u, Reader->getProfiles().size());
+      ASSERT_EQ(3u, Reader->getProfiles().size());
     } else {
       ASSERT_TRUE(ReadBazSamples != nullptr);
       ASSERT_EQ(12557u, ReadBazSamples->getTotalSamples());
-      ASSERT_EQ(3u, Reader->getProfiles().size());
+      ASSERT_EQ(4u, Reader->getProfiles().size());
     }
 
+    FunctionSamples *ReadBooSamples = Reader->getSamplesFor(BooName);
+    ASSERT_TRUE(ReadBooSamples != nullptr);
+    ASSERT_EQ(1232u, ReadBooSamples->getTotalSamples());
+
     std::string MconstructGUID;
     StringRef MconstructRep =
         getRepInFormat(MconstructName, Format, MconstructGUID);
@@ -189,9 +212,9 @@ struct SampleProfTest : ::testing::Test {
 
     auto VerifySummary = [](ProfileSummary &Summary) mutable {
       ASSERT_EQ(ProfileSummary::PSK_Sample, Summary.getKind());
-      ASSERT_EQ(136160u, Summary.getTotalCount());
-      ASSERT_EQ(7u, Summary.getNumCounts());
-      ASSERT_EQ(3u, Summary.getNumFunctions());
+      ASSERT_EQ(137392u, Summary.getTotalCount());
+      ASSERT_EQ(8u, Summary.getNumCounts());
+      ASSERT_EQ(4u, Summary.getNumFunctions());
       ASSERT_EQ(1437u, Summary.getMaxFunctionCount());
       ASSERT_EQ(60351u, Summary.getMaxCount());