1 files changed, 16 insertions, 19 deletions
diff --git a/llvm/tools/llvm-dis/llvm-dis.cpp b/llvm/tools/llvm-dis/llvm-dis.cpp
index 072f6369e12..1349ecc85a4 100644
--- a/llvm/tools/llvm-dis/llvm-dis.cpp
+++ b/llvm/tools/llvm-dis/llvm-dis.cpp
@@ -20,6 +20,8 @@
 #include "llvm/Bitcode/ReaderWriter.h"
 #include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/IR/DebugInfo.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
@@ -112,6 +114,17 @@ public:
 
 } // end anon namespace
 
+static void diagnosticHandler(const DiagnosticInfo &DI, void *Context) {
+  assert(DI.getSeverity() == DS_Error && "Only expecting errors");
+
+  raw_ostream &OS = errs();
+  OS << (char *)Context << ": ";
+  DiagnosticPrinterRawOStream DP(OS);
+  DI.print(DP);
+  OS << '\n';
+  exit(1);
+}
+
 int main(int argc, char **argv) {
   // Print a stack trace if we signal out.
   sys::PrintStackTraceOnErrorSignal();
@@ -120,6 +133,7 @@ int main(int argc, char **argv) {
   LLVMContext &Context = getGlobalContext();
   llvm_shutdown_obj Y;  // Call llvm_shutdown() on exit.
 
+  Context.setDiagnosticHandler(diagnosticHandler, argv[0]);
 
   cl::ParseCommandLineOptions(argc, argv, "llvm .bc -> .ll disassembler\n");
 
@@ -136,25 +150,8 @@ int main(int argc, char **argv) {
       DisplayFilename = InputFilename;
     ErrorOr<std::unique_ptr<Module>> MOrErr =
         getStreamedBitcodeModule(DisplayFilename, Streamer, Context);
-    if (std::error_code EC = MOrErr.getError())
-      ErrorMessage = EC.message();
-    else
-      M = std::move(*MOrErr);
-    if(M.get()) {
-      if (std::error_code EC = M->materializeAllPermanently()) {
-        ErrorMessage = EC.message();
-        M.reset();
-      }
-    }
-  }
-
-  if (!M.get()) {
-    errs() << argv[0] << ": ";
-    if (ErrorMessage.size())
-      errs() << ErrorMessage << "\n";
-    else
-      errs() << "bitcode didn't read correctly.\n";
-    return 1;
+    M = std::move(*MOrErr);
+    M->materializeAllPermanently();
   }
 
   // Just use stdout.  We won't actually print anything on it.