1 files changed, 16 insertions, 7 deletions
diff --git a/llvm/lib/Bitcode/Reader/BitReader.cpp b/llvm/lib/Bitcode/Reader/BitReader.cpp
index 289c76e85b4..beef56bec42 100644
--- a/llvm/lib/Bitcode/Reader/BitReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitReader.cpp
@@ -28,6 +28,13 @@ LLVMBool LLVMParseBitcode(LLVMMemoryBufferRef MemBuf,
                                    OutMessage);
 }
 
+static void diagnosticHandler(const DiagnosticInfo &DI, void *C) {
+  auto *Message = reinterpret_cast<std::string *>(C);
+  raw_string_ostream Stream(*Message);
+  DiagnosticPrinterRawOStream DP(Stream);
+  DI.print(DP);
+}
+
 LLVMBool LLVMParseBitcodeInContext(LLVMContextRef ContextRef,
                                    LLVMMemoryBufferRef MemBuf,
                                    LLVMModuleRef *OutModule,
@@ -35,17 +42,19 @@ LLVMBool LLVMParseBitcodeInContext(LLVMContextRef ContextRef,
   MemoryBufferRef Buf = unwrap(MemBuf)->getMemBufferRef();
   LLVMContext &Ctx = *unwrap(ContextRef);
 
+  LLVMContext::DiagnosticHandlerTy OldDiagnosticHandler =
+      Ctx.getDiagnosticHandler();
+  void *OldDiagnosticContext = Ctx.getDiagnosticContext();
   std::string Message;
-  raw_string_ostream Stream(Message);
-  DiagnosticPrinterRawOStream DP(Stream);
+  Ctx.setDiagnosticHandler(diagnosticHandler, &Message, true);
+
+  ErrorOr<std::unique_ptr<Module>> ModuleOrErr = parseBitcodeFile(Buf, Ctx);
+
+  Ctx.setDiagnosticHandler(OldDiagnosticHandler, OldDiagnosticContext, true);
 
-  ErrorOr<std::unique_ptr<Module>> ModuleOrErr = parseBitcodeFile(
-      Buf, Ctx, [&](const DiagnosticInfo &DI) { DI.print(DP); });
   if (ModuleOrErr.getError()) {
-    if (OutMessage) {
-      Stream.flush();
+    if (OutMessage)
       *OutMessage = strdup(Message.c_str());
-    }
     *OutModule = wrap((Module*)nullptr);
     return 1;
   }