1 files changed, 20 insertions, 15 deletions

diff --git a/llvm/tools/lto/lto.cpp b/llvm/tools/lto/lto.cpp
index 1b218a64cbf..11885211797 100644
--- a/llvm/tools/lto/lto.cpp
+++ b/llvm/tools/lto/lto.cpp
@@ -75,20 +75,23 @@ static bool parsedOptions = false;
 
 static LLVMContext *LTOContext = nullptr;
 
-static void diagnosticHandler(const DiagnosticInfo &DI, void *Context) {
-  if (DI.getSeverity() != DS_Error) {
-    DiagnosticPrinterRawOStream DP(errs());
-    DI.print(DP);
-    errs() << '\n';
-    return;
-  }
-  sLastErrorString = "";
-  {
-    raw_string_ostream Stream(sLastErrorString);
-    DiagnosticPrinterRawOStream DP(Stream);
-    DI.print(DP);
+struct LTOToolDiagnosticHandler : public DiagnosticHandler {
+  bool handleDiagnostics(const DiagnosticInfo &DI) override {
+    if (DI.getSeverity() != DS_Error) {
+      DiagnosticPrinterRawOStream DP(errs());
+      DI.print(DP);
+      errs() << '\n';
+      return true;
+    }
+    sLastErrorString = "";
+    {
+      raw_string_ostream Stream(sLastErrorString);
+      DiagnosticPrinterRawOStream DP(Stream);
+      DI.print(DP);
+    }
+    return true;
   }
-}
+};
 
 // Initialize the configured targets if they have not been initialized.
 static void lto_initialize() {
@@ -108,7 +111,8 @@ static void lto_initialize() {
 
     static LLVMContext Context;
     LTOContext = &Context;
-    LTOContext->setDiagnosticHandler(diagnosticHandler, nullptr, true);
+    LTOContext->setDiagnosticHandler(
+        llvm::make_unique<LTOToolDiagnosticHandler>(), true);
     initialized = true;
   }
 }
@@ -274,7 +278,8 @@ lto_module_t lto_module_create_in_local_context(const void *mem, size_t length,
 
   // Create a local context. Ownership will be transferred to LTOModule.
   std::unique_ptr<LLVMContext> Context = llvm::make_unique<LLVMContext>();
-  Context->setDiagnosticHandler(diagnosticHandler, nullptr, true);
+  Context->setDiagnosticHandler(llvm::make_unique<LTOToolDiagnosticHandler>(),
+                                true);
 
   ErrorOr<std::unique_ptr<LTOModule>> M = LTOModule::createInLocalContext(
       std::move(Context), mem, length, Options, StringRef(path));