1 files changed, 43 insertions, 20 deletions

diff --git a/llvm/lib/IR/LLVMRemarkStreamer.cpp b/llvm/lib/IR/LLVMRemarkStreamer.cpp
index 71f8d4a4b1c7..9e1e45998f2f 100644
--- a/llvm/lib/IR/LLVMRemarkStreamer.cpp
+++ b/llvm/lib/IR/LLVMRemarkStreamer.cpp
@@ -92,7 +92,7 @@ char LLVMRemarkSetupFileError::ID = 0;
 char LLVMRemarkSetupPatternError::ID = 0;
 char LLVMRemarkSetupFormatError::ID = 0;
 
-Expected<std::unique_ptr<ToolOutputFile>> llvm::setupLLVMOptimizationRemarks(
+Expected<LLVMRemarkFileHandle> llvm::setupLLVMOptimizationRemarks(
     LLVMContext &Context, StringRef RemarksFilename, StringRef RemarksPasses,
     StringRef RemarksFormat, bool RemarksWithHotness,
     std::optional<uint64_t> RemarksHotnessThreshold) {
@@ -102,7 +102,7 @@ Expected<std::unique_ptr<ToolOutputFile>> llvm::setupLLVMOptimizationRemarks(
   Context.setDiagnosticsHotnessThreshold(RemarksHotnessThreshold);
 
   if (RemarksFilename.empty())
-    return nullptr;
+    return LLVMRemarkFileHandle();
 
   Expected<remarks::Format> Format = remarks::parseFormat(RemarksFormat);
   if (Error E = Format.takeError())
@@ -119,24 +119,35 @@ Expected<std::unique_ptr<ToolOutputFile>> llvm::setupLLVMOptimizationRemarks(
     return make_error<LLVMRemarkSetupFileError>(errorCodeToError(EC));
 
   Expected<std::unique_ptr<remarks::RemarkSerializer>> RemarkSerializer =
-      remarks::createRemarkSerializer(
-          *Format, remarks::SerializerMode::Separate, RemarksFile->os());
+      remarks::createRemarkSerializer(*Format, RemarksFile->os());
   if (Error E = RemarkSerializer.takeError())
     return make_error<LLVMRemarkSetupFormatError>(std::move(E));
 
-  // Create the main remark streamer.
-  Context.setMainRemarkStreamer(std::make_unique<remarks::RemarkStreamer>(
-      std::move(*RemarkSerializer), RemarksFilename));
+  auto RS = std::make_unique<remarks::RemarkStreamer>(
+      std::move(*RemarkSerializer), RemarksFilename);
+
+  if (!RemarksPasses.empty())
+    if (Error E = RS->setFilter(RemarksPasses)) {
+      RS->releaseSerializer();
+      return make_error<LLVMRemarkSetupPatternError>(std::move(E));
+    }
+
+  // Install the main remark streamer. Only install this after setting the
+  // filter, because this might fail.
+  Context.setMainRemarkStreamer(std::move(RS));
 
   // Create LLVM's optimization remarks streamer.
   Context.setLLVMRemarkStreamer(
       std::make_unique<LLVMRemarkStreamer>(*Context.getMainRemarkStreamer()));
 
-  if (!RemarksPasses.empty())
-    if (Error E = Context.getMainRemarkStreamer()->setFilter(RemarksPasses))
-      return make_error<LLVMRemarkSetupPatternError>(std::move(E));
+  return LLVMRemarkFileHandle{std::move(RemarksFile), Context};
+}
 
-  return std::move(RemarksFile);
+void LLVMRemarkFileHandle::Finalizer::finalize() {
+  if (!Context)
+    return;
+  finalizeLLVMOptimizationRemarks(*Context);
+  Context = nullptr;
 }
 
 Error llvm::setupLLVMOptimizationRemarks(
@@ -153,22 +164,34 @@ Error llvm::setupLLVMOptimizationRemarks(
     return make_error<LLVMRemarkSetupFormatError>(std::move(E));
 
   Expected<std::unique_ptr<remarks::RemarkSerializer>> RemarkSerializer =
-      remarks::createRemarkSerializer(*Format,
-                                      remarks::SerializerMode::Separate, OS);
+      remarks::createRemarkSerializer(*Format, OS);
   if (Error E = RemarkSerializer.takeError())
     return make_error<LLVMRemarkSetupFormatError>(std::move(E));
 
-  // Create the main remark streamer.
-  Context.setMainRemarkStreamer(
-      std::make_unique<remarks::RemarkStreamer>(std::move(*RemarkSerializer)));
+  auto RS =
+      std::make_unique<remarks::RemarkStreamer>(std::move(*RemarkSerializer));
+
+  if (!RemarksPasses.empty())
+    if (Error E = RS->setFilter(RemarksPasses)) {
+      RS->releaseSerializer();
+      return make_error<LLVMRemarkSetupPatternError>(std::move(E));
+    }
+
+  // Install the main remark streamer. Only install this after setting the
+  // filter, because this might fail.
+  Context.setMainRemarkStreamer(std::move(RS));
 
   // Create LLVM's optimization remarks streamer.
   Context.setLLVMRemarkStreamer(
       std::make_unique<LLVMRemarkStreamer>(*Context.getMainRemarkStreamer()));
 
-  if (!RemarksPasses.empty())
-    if (Error E = Context.getMainRemarkStreamer()->setFilter(RemarksPasses))
-      return make_error<LLVMRemarkSetupPatternError>(std::move(E));
-
   return Error::success();
 }
+
+void llvm::finalizeLLVMOptimizationRemarks(LLVMContext &Context) {
+  Context.setLLVMRemarkStreamer(nullptr);
+  if (auto *RS = Context.getMainRemarkStreamer()) {
+    RS->releaseSerializer();
+    Context.setMainRemarkStreamer(nullptr);
+  }
+}