1 files changed, 83 insertions, 37 deletions

diff --git a/llvm/lib/ProfileData/SampleProfReader.cpp b/llvm/lib/ProfileData/SampleProfReader.cpp
index cf3e56728e2..001aafce7bf 100644
--- a/llvm/lib/ProfileData/SampleProfReader.cpp
+++ b/llvm/lib/ProfileData/SampleProfReader.cpp
@@ -191,7 +191,7 @@ static bool ParseLine(const StringRef &Input, bool &IsCallsite, uint32_t &Depth,
 /// the expected format.
 ///
 /// \returns true if the file was loaded successfully, false otherwise.
-std::error_code SampleProfileReaderText::read() {
+std::error_code SampleProfileReaderText::readImpl() {
   line_iterator LineIt(*Buffer, /*SkipBlanks=*/true, '#');
   sampleprof_error Result = sampleprof_error::success;
 
@@ -461,7 +461,7 @@ SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderBinary::read() {
+std::error_code SampleProfileReaderBinary::readImpl() {
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))
       return EC;
@@ -540,15 +540,23 @@ std::error_code SampleProfileReaderExtBinary::readFuncProfiles() {
     return sampleprof_error::success;
   }
 
-  for (auto Name : FuncsToUse) {
-    auto iter = FuncOffsetTable.find(Name);
-    if (iter == FuncOffsetTable.end())
+  if (Remapper) {
+    for (auto Name : FuncsToUse) {
+      Remapper->insert(Name);
+    }
+  }
+
+  for (auto NameOffset : FuncOffsetTable) {
+    auto FuncName = NameOffset.first;
+    if (!FuncsToUse.count(FuncName) &&
+        (!Remapper || !Remapper->exist(FuncName)))
       continue;
-    const uint8_t *FuncProfileAddr = Start + iter->second;
+    const uint8_t *FuncProfileAddr = Start + NameOffset.second;
     assert(FuncProfileAddr < End && "out of LBRProfile section");
     if (std::error_code EC = readFuncProfile(FuncProfileAddr))
       return EC;
   }
+
   Data = End;
   return sampleprof_error::success;
 }
@@ -593,7 +601,7 @@ std::error_code SampleProfileReaderExtBinaryBase::decompressSection(
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinaryBase::read() {
+std::error_code SampleProfileReaderExtBinaryBase::readImpl() {
   const uint8_t *BufStart =
       reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
 
@@ -635,7 +643,7 @@ std::error_code SampleProfileReaderExtBinaryBase::read() {
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderCompactBinary::read() {
+std::error_code SampleProfileReaderCompactBinary::readImpl() {
   std::vector<uint64_t> OffsetsToUse;
   if (UseAllFuncs) {
     for (auto FuncEntry : FuncOffsetTable) {
@@ -1184,7 +1192,7 @@ std::error_code SampleProfileReaderGCC::readOneFunctionProfile(
 ///
 /// This format is generated by the Linux Perf conversion tool at
 /// https://github.com/google/autofdo.
-std::error_code SampleProfileReaderGCC::read() {
+std::error_code SampleProfileReaderGCC::readImpl() {
   // Read the string table.
   if (std::error_code EC = readNameTable())
     return EC;
@@ -1201,38 +1209,31 @@ bool SampleProfileReaderGCC::hasFormat(const MemoryBuffer &Buffer) {
   return Magic == "adcg*704";
 }
 
-std::error_code SampleProfileReaderItaniumRemapper::read() {
-  // If the underlying data is in compact format, we can't remap it because
+void SampleProfileReaderItaniumRemapper::applyRemapping(LLVMContext &Ctx) {
+  // If the reader is in compact format, we can't remap it because
   // we don't know what the original function names were.
-  if (getFormat() == SPF_Compact_Binary) {
+  if (Reader.getFormat() == SPF_Compact_Binary) {
     Ctx.diagnose(DiagnosticInfoSampleProfile(
-        Buffer->getBufferIdentifier(),
+        Reader.getBuffer()->getBufferIdentifier(),
         "Profile data remapping cannot be applied to profile data "
         "in compact format (original mangled names are not available).",
         DS_Warning));
-    return sampleprof_error::success;
-  }
-
-  if (Error E = Remappings.read(*Buffer)) {
-    handleAllErrors(
-        std::move(E), [&](const SymbolRemappingParseError &ParseError) {
-          reportError(ParseError.getLineNum(), ParseError.getMessage());
-        });
-    return sampleprof_error::malformed;
+    return;
   }
 
-  for (auto &Sample : getProfiles())
-    if (auto Key = Remappings.insert(Sample.first()))
+  assert(Remappings && "should be initialized while creating remapper");
+  for (auto &Sample : Reader.getProfiles())
+    if (auto Key = Remappings->insert(Sample.first()))
       SampleMap.insert({Key, &Sample.second});
 
-  return sampleprof_error::success;
+  RemappingApplied = true;
 }
 
 FunctionSamples *
 SampleProfileReaderItaniumRemapper::getSamplesFor(StringRef Fname) {
-  if (auto Key = Remappings.lookup(Fname))
+  if (auto Key = Remappings->lookup(Fname))
     return SampleMap.lookup(Key);
-  return SampleProfileReader::getSamplesFor(Fname);
+  return nullptr;
 }
 
 /// Prepare a memory buffer for the contents of \p Filename.
@@ -1258,13 +1259,16 @@ setupMemoryBuffer(const Twine &Filename) {
 ///
 /// \param C The LLVM context to use to emit diagnostics.
 ///
+/// \param RemapFilename The file used for profile remapping.
+///
 /// \returns an error code indicating the status of the created reader.
 ErrorOr<std::unique_ptr<SampleProfileReader>>
-SampleProfileReader::create(const Twine &Filename, LLVMContext &C) {
+SampleProfileReader::create(const std::string Filename, LLVMContext &C,
+                            const std::string RemapFilename) {
   auto BufferOrError = setupMemoryBuffer(Filename);
   if (std::error_code EC = BufferOrError.getError())
     return EC;
-  return create(BufferOrError.get(), C);
+  return create(BufferOrError.get(), C, RemapFilename);
 }
 
 /// Create a sample profile remapper from the given input, to remap the
@@ -1272,20 +1276,48 @@ SampleProfileReader::create(const Twine &Filename, LLVMContext &C) {
 ///
 /// \param Filename The file to open.
 ///
-/// \param C The LLVM context to use to emit diagnostics.
+/// \param Reader The profile reader the remapper is going to be applied to.
 ///
-/// \param Underlying The underlying profile data reader to remap.
+/// \param C The LLVM context to use to emit diagnostics.
 ///
 /// \returns an error code indicating the status of the created reader.
-ErrorOr<std::unique_ptr<SampleProfileReader>>
-SampleProfileReaderItaniumRemapper::create(
-    const Twine &Filename, LLVMContext &C,
-    std::unique_ptr<SampleProfileReader> Underlying) {
+ErrorOr<std::unique_ptr<SampleProfileReaderItaniumRemapper>>
+SampleProfileReaderItaniumRemapper::create(const std::string Filename,
+                                           SampleProfileReader &Reader,
+                                           LLVMContext &C) {
   auto BufferOrError = setupMemoryBuffer(Filename);
   if (std::error_code EC = BufferOrError.getError())
     return EC;
+  return create(BufferOrError.get(), Reader, C);
+}
+
+/// Create a sample profile remapper from the given input, to remap the
+/// function names in the given profile data.
+///
+/// \param B The memory buffer to create the reader from (assumes ownership).
+///
+/// \param C The LLVM context to use to emit diagnostics.
+///
+/// \param Reader The profile reader the remapper is going to be applied to.
+///
+/// \returns an error code indicating the status of the created reader.
+ErrorOr<std::unique_ptr<SampleProfileReaderItaniumRemapper>>
+SampleProfileReaderItaniumRemapper::create(std::unique_ptr<MemoryBuffer> &B,
+                                           SampleProfileReader &Reader,
+                                           LLVMContext &C) {
+  auto Remappings = std::make_unique<SymbolRemappingReader>();
+  if (Error E = Remappings->read(*B.get())) {
+    handleAllErrors(
+        std::move(E), [&](const SymbolRemappingParseError &ParseError) {
+          C.diagnose(DiagnosticInfoSampleProfile(B->getBufferIdentifier(),
+                                                 ParseError.getLineNum(),
+                                                 ParseError.getMessage()));
+        });
+    return sampleprof_error::malformed;
+  }
+
   return std::make_unique<SampleProfileReaderItaniumRemapper>(
-      std::move(BufferOrError.get()), C, std::move(Underlying));
+      std::move(B), std::move(Remappings), Reader);
 }
 
 /// Create a sample profile reader based on the format of the input data.
@@ -1294,9 +1326,12 @@ SampleProfileReaderItaniumRemapper::create(
 ///
 /// \param C The LLVM context to use to emit diagnostics.
 ///
+/// \param RemapFilename The file used for profile remapping.
+///
 /// \returns an error code indicating the status of the created reader.
 ErrorOr<std::unique_ptr<SampleProfileReader>>
-SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C) {
+SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C,
+                            const std::string RemapFilename) {
   std::unique_ptr<SampleProfileReader> Reader;
   if (SampleProfileReaderRawBinary::hasFormat(*B))
     Reader.reset(new SampleProfileReaderRawBinary(std::move(B), C));
@@ -1311,6 +1346,17 @@ SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C) {
   else
     return sampleprof_error::unrecognized_format;
 
+  if (!RemapFilename.empty()) {
+    auto ReaderOrErr =
+        SampleProfileReaderItaniumRemapper::create(RemapFilename, *Reader, C);
+    if (std::error_code EC = ReaderOrErr.getError()) {
+      std::string Msg = "Could not create remapper: " + EC.message();
+      C.diagnose(DiagnosticInfoSampleProfile(RemapFilename, Msg));
+      return EC;
+    }
+    Reader->Remapper = std::move(ReaderOrErr.get());
+  }
+
   FunctionSamples::Format = Reader->getFormat();
   if (std::error_code EC = Reader->readHeader()) {
     return EC;