1 files changed, 51 insertions, 21 deletions

diff --git a/lib/IR/RemarkStreamer.cpp b/lib/IR/RemarkStreamer.cpp
index 5b4c7e72b479..0fcc06b961f3 100644
--- a/lib/IR/RemarkStreamer.cpp
+++ b/lib/IR/RemarkStreamer.cpp
@@ -15,15 +15,17 @@
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
+#include "llvm/Remarks/BitstreamRemarkSerializer.h"
 #include "llvm/Remarks/RemarkFormat.h"
+#include "llvm/Remarks/RemarkSerializer.h"
 
 using namespace llvm;
 
-RemarkStreamer::RemarkStreamer(StringRef Filename,
-                               std::unique_ptr<remarks::Serializer> Serializer)
-    : Filename(Filename), PassFilter(), Serializer(std::move(Serializer)) {
-  assert(!Filename.empty() && "This needs to be a real filename.");
-}
+RemarkStreamer::RemarkStreamer(
+    std::unique_ptr<remarks::RemarkSerializer> RemarkSerializer,
+    Optional<StringRef> FilenameIn)
+    : PassFilter(), RemarkSerializer(std::move(RemarkSerializer)),
+      Filename(FilenameIn ? Optional<std::string>(FilenameIn->str()) : None) {}
 
 Error RemarkStreamer::setFilter(StringRef Filter) {
   Regex R = Regex(Filter);
@@ -99,24 +101,13 @@ void RemarkStreamer::emit(const DiagnosticInfoOptimizationBase &Diag) {
   // First, convert the diagnostic to a remark.
   remarks::Remark R = toRemark(Diag);
   // Then, emit the remark through the serializer.
-  Serializer->emit(R);
+  RemarkSerializer->emit(R);
 }
 
 char RemarkSetupFileError::ID = 0;
 char RemarkSetupPatternError::ID = 0;
 char RemarkSetupFormatError::ID = 0;
 
-static std::unique_ptr<remarks::Serializer>
-formatToSerializer(remarks::Format RemarksFormat, raw_ostream &OS) {
-  switch (RemarksFormat) {
-  default:
-    llvm_unreachable("Unknown remark serializer format.");
-    return nullptr;
-  case remarks::Format::YAML:
-    return llvm::make_unique<remarks::YAMLSerializer>(OS);
-  };
-}
-
 Expected<std::unique_ptr<ToolOutputFile>>
 llvm::setupOptimizationRemarks(LLVMContext &Context, StringRef RemarksFilename,
                                StringRef RemarksPasses, StringRef RemarksFormat,
@@ -131,24 +122,63 @@ llvm::setupOptimizationRemarks(LLVMContext &Context, StringRef RemarksFilename,
   if (RemarksFilename.empty())
     return nullptr;
 
+  Expected<remarks::Format> Format = remarks::parseFormat(RemarksFormat);
+  if (Error E = Format.takeError())
+    return make_error<RemarkSetupFormatError>(std::move(E));
+
   std::error_code EC;
+  auto Flags = *Format == remarks::Format::YAML ? sys::fs::OF_Text
+                                                : sys::fs::OF_None;
   auto RemarksFile =
-      llvm::make_unique<ToolOutputFile>(RemarksFilename, EC, sys::fs::F_None);
+      std::make_unique<ToolOutputFile>(RemarksFilename, EC, Flags);
   // We don't use llvm::FileError here because some diagnostics want the file
   // name separately.
   if (EC)
     return make_error<RemarkSetupFileError>(errorCodeToError(EC));
 
+  Expected<std::unique_ptr<remarks::RemarkSerializer>> RemarkSerializer =
+      remarks::createRemarkSerializer(
+          *Format, remarks::SerializerMode::Separate, RemarksFile->os());
+  if (Error E = RemarkSerializer.takeError())
+    return make_error<RemarkSetupFormatError>(std::move(E));
+
+  Context.setRemarkStreamer(std::make_unique<RemarkStreamer>(
+      std::move(*RemarkSerializer), RemarksFilename));
+
+  if (!RemarksPasses.empty())
+    if (Error E = Context.getRemarkStreamer()->setFilter(RemarksPasses))
+      return make_error<RemarkSetupPatternError>(std::move(E));
+
+  return std::move(RemarksFile);
+}
+
+Error llvm::setupOptimizationRemarks(LLVMContext &Context, raw_ostream &OS,
+                                     StringRef RemarksPasses,
+                                     StringRef RemarksFormat,
+                                     bool RemarksWithHotness,
+                                     unsigned RemarksHotnessThreshold) {
+  if (RemarksWithHotness)
+    Context.setDiagnosticsHotnessRequested(true);
+
+  if (RemarksHotnessThreshold)
+    Context.setDiagnosticsHotnessThreshold(RemarksHotnessThreshold);
+
   Expected<remarks::Format> Format = remarks::parseFormat(RemarksFormat);
   if (Error E = Format.takeError())
     return make_error<RemarkSetupFormatError>(std::move(E));
 
-  Context.setRemarkStreamer(llvm::make_unique<RemarkStreamer>(
-      RemarksFilename, formatToSerializer(*Format, RemarksFile->os())));
+  Expected<std::unique_ptr<remarks::RemarkSerializer>> RemarkSerializer =
+      remarks::createRemarkSerializer(*Format,
+                                      remarks::SerializerMode::Separate, OS);
+  if (Error E = RemarkSerializer.takeError())
+    return make_error<RemarkSetupFormatError>(std::move(E));
+
+  Context.setRemarkStreamer(
+      std::make_unique<RemarkStreamer>(std::move(*RemarkSerializer)));
 
   if (!RemarksPasses.empty())
     if (Error E = Context.getRemarkStreamer()->setFilter(RemarksPasses))
       return make_error<RemarkSetupPatternError>(std::move(E));
 
-  return std::move(RemarksFile);
+  return Error::success();
 }