Use diagnostic handler in the LLVMContext

[oota-llvm.git] / lib / Bitcode / Reader / BitReader.cpp

diff --git a/lib/Bitcode/Reader/BitReader.cpp b/lib/Bitcode/Reader/BitReader.cpp
index 4136cec7c3278b30cff6bcc8185acf8d241c2eba..beef56bec4263e5f8439d3483ff415d4de2cc999 100644 (file)
--- a/lib/Bitcode/Reader/BitReader.cpp
+++ b/lib/Bitcode/Reader/BitReader.cpp
@@ -9,9 +9,11 @@
 
 #include "llvm-c/BitReader.h"
 #include "llvm/Bitcode/ReaderWriter.h"
+#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/raw_ostream.h"
 #include <cstring>
 #include <string>
 
@@ -26,20 +28,38 @@ 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,
                                    char **OutMessage) {
+  MemoryBufferRef Buf = unwrap(MemBuf)->getMemBufferRef();
+  LLVMContext &Ctx = *unwrap(ContextRef);
+
+  LLVMContext::DiagnosticHandlerTy OldDiagnosticHandler =
+      Ctx.getDiagnosticHandler();
+  void *OldDiagnosticContext = Ctx.getDiagnosticContext();
   std::string Message;
+  Ctx.setDiagnosticHandler(diagnosticHandler, &Message, true);
 
-  *OutModule = wrap(ParseBitcodeFile(unwrap(MemBuf), *unwrap(ContextRef),
-                                     &Message));
-  if (!*OutModule) {
+  ErrorOr<std::unique_ptr<Module>> ModuleOrErr = parseBitcodeFile(Buf, Ctx);
+
+  Ctx.setDiagnosticHandler(OldDiagnosticHandler, OldDiagnosticContext, true);
+
+  if (ModuleOrErr.getError()) {
     if (OutMessage)
       *OutMessage = strdup(Message.c_str());
+    *OutModule = wrap((Module*)nullptr);
     return 1;
   }
 
+  *OutModule = wrap(ModuleOrErr.get().release());
   return 0;
 }
 
@@ -51,17 +71,20 @@ LLVMBool LLVMGetBitcodeModuleInContext(LLVMContextRef ContextRef,
                                        LLVMModuleRef *OutM,
                                        char **OutMessage) {
   std::string Message;
-  ErrorOr<Module *> ModuleOrErr =
-      getLazyBitcodeModule(unwrap(MemBuf), *unwrap(ContextRef));
+  std::unique_ptr<MemoryBuffer> Owner(unwrap(MemBuf));
+
+  ErrorOr<std::unique_ptr<Module>> ModuleOrErr =
+      getLazyBitcodeModule(std::move(Owner), *unwrap(ContextRef));
+  Owner.release();
 
-  if (error_code EC = ModuleOrErr.getError()) {
-    *OutM = wrap((Module *)NULL);
+  if (std::error_code EC = ModuleOrErr.getError()) {
+    *OutM = wrap((Module *)nullptr);
     if (OutMessage)
       *OutMessage = strdup(EC.message().c_str());
     return 1;
   }
 
-  *OutM = wrap(ModuleOrErr.get());
+  *OutM = wrap(ModuleOrErr.get().release());
 
   return 0;