From: Rafael Espindola <rafael.espindola@gmail.com>
Date: Mon, 14 Dec 2015 23:17:03 +0000 (+0000)
Subject: Use diagnostic handler in the LLVMContext
X-Git-Url: http://plrg.eecs.uci.edu/git/?p=oota-llvm.git;a=commitdiff_plain;h=7a1fc2d33e22333b6ebe8f1ef2f6c14ad82bd114

Use diagnostic handler in the LLVMContext

This patch converts code that has access to a LLVMContext to not take a
diagnostic handler.

This has a few advantages

* It is easier to use a consistent diagnostic handler in a single program.
* Less clutter since we are not passing a handler around.

It does make it a bit awkward to implement some C APIs that return a
diagnostic string. I will propose new versions of these APIs and
deprecate the current ones.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@255571 91177308-0d34-0410-b5e6-96231b3b80d8
---

diff --git a/include/llvm/Bitcode/ReaderWriter.h b/include/llvm/Bitcode/ReaderWriter.h
index 3e127290f37..60d865fd235 100644
--- a/include/llvm/Bitcode/ReaderWriter.h
+++ b/include/llvm/Bitcode/ReaderWriter.h
@@ -37,34 +37,30 @@ namespace llvm {
   ErrorOr<std::unique_ptr<Module>>
   getLazyBitcodeModule(std::unique_ptr<MemoryBuffer> &&Buffer,
                        LLVMContext &Context,
-                       DiagnosticHandlerFunction DiagnosticHandler = nullptr,
                        bool ShouldLazyLoadMetadata = false);
 
   /// Read the header of the specified stream and prepare for lazy
   /// deserialization and streaming of function bodies.
-  ErrorOr<std::unique_ptr<Module>> getStreamedBitcodeModule(
-      StringRef Name, std::unique_ptr<DataStreamer> Streamer,
-      LLVMContext &Context,
-      DiagnosticHandlerFunction DiagnosticHandler = nullptr);
+  ErrorOr<std::unique_ptr<Module>>
+  getStreamedBitcodeModule(StringRef Name,
+                           std::unique_ptr<DataStreamer> Streamer,
+                           LLVMContext &Context);
 
   /// Read the header of the specified bitcode buffer and extract just the
   /// triple information. If successful, this returns a string. On error, this
   /// returns "".
-  std::string
-  getBitcodeTargetTriple(MemoryBufferRef Buffer, LLVMContext &Context,
-                         DiagnosticHandlerFunction DiagnosticHandler = nullptr);
+  std::string getBitcodeTargetTriple(MemoryBufferRef Buffer,
+                                     LLVMContext &Context);
 
   /// Read the header of the specified bitcode buffer and extract just the
   /// producer string information. If successful, this returns a string. On
   /// error, this returns "".
-  std::string getBitcodeProducerString(
-      MemoryBufferRef Buffer, LLVMContext &Context,
-      DiagnosticHandlerFunction DiagnosticHandler = nullptr);
+  std::string getBitcodeProducerString(MemoryBufferRef Buffer,
+                                       LLVMContext &Context);
 
   /// Read the specified bitcode file, returning the module.
-  ErrorOr<std::unique_ptr<Module>>
-  parseBitcodeFile(MemoryBufferRef Buffer, LLVMContext &Context,
-                   DiagnosticHandlerFunction DiagnosticHandler = nullptr);
+  ErrorOr<std::unique_ptr<Module>> parseBitcodeFile(MemoryBufferRef Buffer,
+                                                    LLVMContext &Context);
 
   /// Check if the given bitcode buffer contains a function summary block.
   bool hasFunctionSummary(MemoryBufferRef Buffer,
@@ -75,9 +71,10 @@ namespace llvm {
   /// the index. Otherwise skip the function summary section, and only create
   /// an index object with a map from function name to function summary offset.
   /// The index is used to perform lazy function summary reading later.
-  ErrorOr<std::unique_ptr<FunctionInfoIndex>> getFunctionInfoIndex(
-      MemoryBufferRef Buffer, DiagnosticHandlerFunction DiagnosticHandler,
-      bool IsLazy = false);
+  ErrorOr<std::unique_ptr<FunctionInfoIndex>>
+  getFunctionInfoIndex(MemoryBufferRef Buffer,
+                       DiagnosticHandlerFunction DiagnosticHandler,
+                       bool IsLazy = false);
 
   /// This method supports lazy reading of function summary data from the
   /// combined index during function importing. When reading the combined index
diff --git a/include/llvm/Linker/IRMover.h b/include/llvm/Linker/IRMover.h
index 89d02f129c8..e21899cb61a 100644
--- a/include/llvm/Linker/IRMover.h
+++ b/include/llvm/Linker/IRMover.h
@@ -12,7 +12,6 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
-#include "llvm/IR/DiagnosticInfo.h"
 
 namespace llvm {
 class GlobalValue;
@@ -54,7 +53,7 @@ public:
     bool hasType(StructType *Ty);
   };
 
-  IRMover(Module &M, DiagnosticHandlerFunction DiagnosticHandler);
+  IRMover(Module &M);
 
   typedef std::function<void(GlobalValue &)> ValueAdder;
   /// Move in the provide values. The source is destroyed.
@@ -63,14 +62,9 @@ public:
             std::function<void(GlobalValue &GV, ValueAdder Add)> AddLazyFor);
   Module &getModule() { return Composite; }
 
-  DiagnosticHandlerFunction getDiagnosticHandler() const {
-    return DiagnosticHandler;
-  }
-
 private:
   Module &Composite;
   IdentifiedStructTypeSet IdentifiedStructTypes;
-  DiagnosticHandlerFunction DiagnosticHandler;
 };
 
 } // End llvm namespace
diff --git a/include/llvm/Linker/Linker.h b/include/llvm/Linker/Linker.h
index 9ff61bc518b..26c27c3aae8 100644
--- a/include/llvm/Linker/Linker.h
+++ b/include/llvm/Linker/Linker.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_LINKER_LINKER_H
 #define LLVM_LINKER_LINKER_H
 
-#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/FunctionInfo.h"
 #include "llvm/Linker/IRMover.h"
 
@@ -34,7 +33,7 @@ public:
     InternalizeLinkedSymbols = (1 << 2)
   };
 
-  Linker(Module &M, DiagnosticHandlerFunction DiagnosticHandler);
+  Linker(Module &M);
 
   /// \brief Link \p Src into the composite. The source is destroyed.
   ///
@@ -50,20 +49,14 @@ public:
                     DenseSet<const GlobalValue *> *FunctionsToImport = nullptr);
 
   static bool linkModules(Module &Dest, Module &Src,
-                          DiagnosticHandlerFunction DiagnosticHandler,
                           unsigned Flags = Flags::None);
 
-  DiagnosticHandlerFunction getDiagnosticHandler() const {
-    return Mover.getDiagnosticHandler();
-  }
 };
 
 /// Create a new module with exported local functions renamed and promoted
 /// for ThinLTO.
-std::unique_ptr<Module>
-renameModuleForThinLTO(std::unique_ptr<Module> &M,
-                       const FunctionInfoIndex *Index,
-                       DiagnosticHandlerFunction DiagnosticHandler);
+std::unique_ptr<Module> renameModuleForThinLTO(std::unique_ptr<Module> &M,
+                                               const FunctionInfoIndex *Index);
 
 } // End llvm namespace
 
diff --git a/include/llvm/Transforms/IPO/FunctionImport.h b/include/llvm/Transforms/IPO/FunctionImport.h
index e3da8a79c23..0f857d70073 100644
--- a/include/llvm/Transforms/IPO/FunctionImport.h
+++ b/include/llvm/Transforms/IPO/FunctionImport.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_FUNCTIONIMPORT_H
 #define LLVM_FUNCTIONIMPORT_H
 
-#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace llvm {
@@ -25,9 +24,6 @@ class FunctionImporter {
   /// The summaries index used to trigger importing.
   const FunctionInfoIndex &Index;
 
-  /// Diagnostic will be sent to this handler.
-  DiagnosticHandlerFunction DiagnosticHandler;
-
   /// Factory function to load a Module for a given identifier
   std::function<std::unique_ptr<Module>(StringRef Identifier)> ModuleLoader;
 
@@ -35,10 +31,8 @@ public:
   /// Create a Function Importer.
   FunctionImporter(
       const FunctionInfoIndex &Index,
-      DiagnosticHandlerFunction DiagnosticHandler,
       std::function<std::unique_ptr<Module>(StringRef Identifier)> ModuleLoader)
-      : Index(Index), DiagnosticHandler(DiagnosticHandler),
-        ModuleLoader(ModuleLoader) {}
+      : Index(Index), ModuleLoader(ModuleLoader) {}
 
   /// Import functions in Module \p M based on the summary informations.
   bool importFunctions(Module &M);
diff --git a/lib/Bitcode/Reader/BitReader.cpp b/lib/Bitcode/Reader/BitReader.cpp
index 289c76e85b4..beef56bec42 100644
--- a/lib/Bitcode/Reader/BitReader.cpp
+++ b/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;
   }
diff --git a/lib/Bitcode/Reader/BitcodeReader.cpp b/lib/Bitcode/Reader/BitcodeReader.cpp
index d59f12898e9..d80e70d3515 100644
--- a/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -132,7 +132,6 @@ public:
 
 class BitcodeReader : public GVMaterializer {
   LLVMContext &Context;
-  DiagnosticHandlerFunction DiagnosticHandler;
   Module *TheModule = nullptr;
   std::unique_ptr<MemoryBuffer> Buffer;
   std::unique_ptr<BitstreamReader> StreamFile;
@@ -239,10 +238,8 @@ public:
   std::error_code error(BitcodeError E);
   std::error_code error(const Twine &Message);
 
-  BitcodeReader(MemoryBuffer *Buffer, LLVMContext &Context,
-                DiagnosticHandlerFunction DiagnosticHandler);
-  BitcodeReader(LLVMContext &Context,
-                DiagnosticHandlerFunction DiagnosticHandler);
+  BitcodeReader(MemoryBuffer *Buffer, LLVMContext &Context);
+  BitcodeReader(LLVMContext &Context);
   ~BitcodeReader() override { freeState(); }
 
   std::error_code materializeForwardReferencedFunctions();
@@ -518,54 +515,51 @@ static std::error_code error(DiagnosticHandlerFunction DiagnosticHandler,
   return error(DiagnosticHandler, EC, EC.message());
 }
 
-static std::error_code error(DiagnosticHandlerFunction DiagnosticHandler,
+static std::error_code error(LLVMContext &Context, std::error_code EC,
                              const Twine &Message) {
-  return error(DiagnosticHandler,
-               make_error_code(BitcodeError::CorruptedBitcode), Message);
+  return error([&](const DiagnosticInfo &DI) { Context.diagnose(DI); }, EC,
+               Message);
+}
+
+static std::error_code error(LLVMContext &Context, std::error_code EC) {
+  return error(Context, EC, EC.message());
+}
+
+static std::error_code error(LLVMContext &Context, const Twine &Message) {
+  return error(Context, make_error_code(BitcodeError::CorruptedBitcode),
+               Message);
 }
 
 std::error_code BitcodeReader::error(BitcodeError E, const Twine &Message) {
   if (!ProducerIdentification.empty()) {
-    return ::error(DiagnosticHandler, make_error_code(E),
+    return ::error(Context, make_error_code(E),
                    Message + " (Producer: '" + ProducerIdentification +
                        "' Reader: 'LLVM " + LLVM_VERSION_STRING "')");
   }
-  return ::error(DiagnosticHandler, make_error_code(E), Message);
+  return ::error(Context, make_error_code(E), Message);
 }
 
 std::error_code BitcodeReader::error(const Twine &Message) {
   if (!ProducerIdentification.empty()) {
-    return ::error(DiagnosticHandler,
-                   make_error_code(BitcodeError::CorruptedBitcode),
+    return ::error(Context, make_error_code(BitcodeError::CorruptedBitcode),
                    Message + " (Producer: '" + ProducerIdentification +
                        "' Reader: 'LLVM " + LLVM_VERSION_STRING "')");
   }
-  return ::error(DiagnosticHandler,
-                 make_error_code(BitcodeError::CorruptedBitcode), Message);
+  return ::error(Context, make_error_code(BitcodeError::CorruptedBitcode),
+                 Message);
 }
 
 std::error_code BitcodeReader::error(BitcodeError E) {
-  return ::error(DiagnosticHandler, make_error_code(E));
+  return ::error(Context, make_error_code(E));
 }
 
-static DiagnosticHandlerFunction getDiagHandler(DiagnosticHandlerFunction F,
-                                                LLVMContext &C) {
-  if (F)
-    return F;
-  return [&C](const DiagnosticInfo &DI) { C.diagnose(DI); };
-}
-
-BitcodeReader::BitcodeReader(MemoryBuffer *Buffer, LLVMContext &Context,
-                             DiagnosticHandlerFunction DiagnosticHandler)
-    : Context(Context),
-      DiagnosticHandler(getDiagHandler(DiagnosticHandler, Context)),
-      Buffer(Buffer), ValueList(Context), MDValueList(Context) {}
+BitcodeReader::BitcodeReader(MemoryBuffer *Buffer, LLVMContext &Context)
+    : Context(Context), Buffer(Buffer), ValueList(Context),
+      MDValueList(Context) {}
 
-BitcodeReader::BitcodeReader(LLVMContext &Context,
-                             DiagnosticHandlerFunction DiagnosticHandler)
-    : Context(Context),
-      DiagnosticHandler(getDiagHandler(DiagnosticHandler, Context)),
-      Buffer(nullptr), ValueList(Context), MDValueList(Context) {}
+BitcodeReader::BitcodeReader(LLVMContext &Context)
+    : Context(Context), Buffer(nullptr), ValueList(Context),
+      MDValueList(Context) {}
 
 std::error_code BitcodeReader::materializeForwardReferencedFunctions() {
   if (WillMaterializeAllForwardRefs)
@@ -3898,17 +3892,17 @@ std::error_code BitcodeReader::parseMetadataAttachment(Function &F) {
   }
 }
 
-static std::error_code typeCheckLoadStoreInst(DiagnosticHandlerFunction DH,
-                                              Type *ValType, Type *PtrType) {
+static std::error_code typeCheckLoadStoreInst(Type *ValType, Type *PtrType) {
+  LLVMContext &Context = PtrType->getContext();
   if (!isa<PointerType>(PtrType))
-    return error(DH, "Load/Store operand is not a pointer type");
+    return error(Context, "Load/Store operand is not a pointer type");
   Type *ElemType = cast<PointerType>(PtrType)->getElementType();
 
   if (ValType && ValType != ElemType)
-    return error(DH, "Explicit load/store type does not match pointee type of "
-                     "pointer operand");
+    return error(Context, "Explicit load/store type does not match pointee "
+                          "type of pointer operand");
   if (!PointerType::isLoadableOrStorableType(ElemType))
-    return error(DH, "Cannot load/store from pointer");
+    return error(Context, "Cannot load/store from pointer");
   return std::error_code();
 }
 
@@ -4822,8 +4816,7 @@ std::error_code BitcodeReader::parseFunctionBody(Function *F) {
       Type *Ty = nullptr;
       if (OpNum + 3 == Record.size())
         Ty = getTypeByID(Record[OpNum++]);
-      if (std::error_code EC =
-              typeCheckLoadStoreInst(DiagnosticHandler, Ty, Op->getType()))
+      if (std::error_code EC = typeCheckLoadStoreInst(Ty, Op->getType()))
         return EC;
       if (!Ty)
         Ty = cast<PointerType>(Op->getType())->getElementType();
@@ -4847,8 +4840,7 @@ std::error_code BitcodeReader::parseFunctionBody(Function *F) {
       Type *Ty = nullptr;
       if (OpNum + 5 == Record.size())
         Ty = getTypeByID(Record[OpNum++]);
-      if (std::error_code EC =
-              typeCheckLoadStoreInst(DiagnosticHandler, Ty, Op->getType()))
+      if (std::error_code EC = typeCheckLoadStoreInst(Ty, Op->getType()))
         return EC;
       if (!Ty)
         Ty = cast<PointerType>(Op->getType())->getElementType();
@@ -4882,8 +4874,8 @@ std::error_code BitcodeReader::parseFunctionBody(Function *F) {
           OpNum + 2 != Record.size())
         return error("Invalid record");
 
-      if (std::error_code EC = typeCheckLoadStoreInst(
-              DiagnosticHandler, Val->getType(), Ptr->getType()))
+      if (std::error_code EC =
+              typeCheckLoadStoreInst(Val->getType(), Ptr->getType()))
         return EC;
       unsigned Align;
       if (std::error_code EC = parseAlignmentValue(Record[OpNum], Align))
@@ -4906,8 +4898,8 @@ std::error_code BitcodeReader::parseFunctionBody(Function *F) {
           OpNum + 4 != Record.size())
         return error("Invalid record");
 
-      if (std::error_code EC = typeCheckLoadStoreInst(
-              DiagnosticHandler, Val->getType(), Ptr->getType()))
+      if (std::error_code EC =
+              typeCheckLoadStoreInst(Val->getType(), Ptr->getType()))
         return EC;
       AtomicOrdering Ordering = getDecodedOrdering(Record[OpNum + 2]);
       if (Ordering == NotAtomic || Ordering == Acquire ||
@@ -4944,8 +4936,8 @@ std::error_code BitcodeReader::parseFunctionBody(Function *F) {
         return error("Invalid record");
       SynchronizationScope SynchScope = getDecodedSynchScope(Record[OpNum + 2]);
 
-      if (std::error_code EC = typeCheckLoadStoreInst(
-              DiagnosticHandler, Cmp->getType(), Ptr->getType()))
+      if (std::error_code EC =
+              typeCheckLoadStoreInst(Cmp->getType(), Ptr->getType()))
         return EC;
       AtomicOrdering FailureOrdering;
       if (Record.size() < 7)
@@ -5863,10 +5855,8 @@ getBitcodeModuleImpl(std::unique_ptr<DataStreamer> Streamer, StringRef Name,
 static ErrorOr<std::unique_ptr<Module>>
 getLazyBitcodeModuleImpl(std::unique_ptr<MemoryBuffer> &&Buffer,
                          LLVMContext &Context, bool MaterializeAll,
-                         DiagnosticHandlerFunction DiagnosticHandler,
                          bool ShouldLazyLoadMetadata = false) {
-  BitcodeReader *R =
-      new BitcodeReader(Buffer.get(), Context, DiagnosticHandler);
+  BitcodeReader *R = new BitcodeReader(Buffer.get(), Context);
 
   ErrorOr<std::unique_ptr<Module>> Ret =
       getBitcodeModuleImpl(nullptr, Buffer->getBufferIdentifier(), R, Context,
@@ -5878,50 +5868,46 @@ getLazyBitcodeModuleImpl(std::unique_ptr<MemoryBuffer> &&Buffer,
   return Ret;
 }
 
-ErrorOr<std::unique_ptr<Module>> llvm::getLazyBitcodeModule(
-    std::unique_ptr<MemoryBuffer> &&Buffer, LLVMContext &Context,
-    DiagnosticHandlerFunction DiagnosticHandler, bool ShouldLazyLoadMetadata) {
+ErrorOr<std::unique_ptr<Module>>
+llvm::getLazyBitcodeModule(std::unique_ptr<MemoryBuffer> &&Buffer,
+                           LLVMContext &Context, bool ShouldLazyLoadMetadata) {
   return getLazyBitcodeModuleImpl(std::move(Buffer), Context, false,
-                                  DiagnosticHandler, ShouldLazyLoadMetadata);
+                                  ShouldLazyLoadMetadata);
 }
 
-ErrorOr<std::unique_ptr<Module>> llvm::getStreamedBitcodeModule(
-    StringRef Name, std::unique_ptr<DataStreamer> Streamer,
-    LLVMContext &Context, DiagnosticHandlerFunction DiagnosticHandler) {
+ErrorOr<std::unique_ptr<Module>>
+llvm::getStreamedBitcodeModule(StringRef Name,
+                               std::unique_ptr<DataStreamer> Streamer,
+                               LLVMContext &Context) {
   std::unique_ptr<Module> M = make_unique<Module>(Name, Context);
-  BitcodeReader *R = new BitcodeReader(Context, DiagnosticHandler);
+  BitcodeReader *R = new BitcodeReader(Context);
 
   return getBitcodeModuleImpl(std::move(Streamer), Name, R, Context, false,
                               false);
 }
 
-ErrorOr<std::unique_ptr<Module>>
-llvm::parseBitcodeFile(MemoryBufferRef Buffer, LLVMContext &Context,
-                       DiagnosticHandlerFunction DiagnosticHandler) {
+ErrorOr<std::unique_ptr<Module>> llvm::parseBitcodeFile(MemoryBufferRef Buffer,
+                                                        LLVMContext &Context) {
   std::unique_ptr<MemoryBuffer> Buf = MemoryBuffer::getMemBuffer(Buffer, false);
-  return getLazyBitcodeModuleImpl(std::move(Buf), Context, true,
-                                  DiagnosticHandler);
+  return getLazyBitcodeModuleImpl(std::move(Buf), Context, true);
   // TODO: Restore the use-lists to the in-memory state when the bitcode was
   // written.  We must defer until the Module has been fully materialized.
 }
 
-std::string
-llvm::getBitcodeTargetTriple(MemoryBufferRef Buffer, LLVMContext &Context,
-                             DiagnosticHandlerFunction DiagnosticHandler) {
+std::string llvm::getBitcodeTargetTriple(MemoryBufferRef Buffer,
+                                         LLVMContext &Context) {
   std::unique_ptr<MemoryBuffer> Buf = MemoryBuffer::getMemBuffer(Buffer, false);
-  auto R = llvm::make_unique<BitcodeReader>(Buf.release(), Context,
-                                            DiagnosticHandler);
+  auto R = llvm::make_unique<BitcodeReader>(Buf.release(), Context);
   ErrorOr<std::string> Triple = R->parseTriple();
   if (Triple.getError())
     return "";
   return Triple.get();
 }
 
-std::string
-llvm::getBitcodeProducerString(MemoryBufferRef Buffer, LLVMContext &Context,
-                               DiagnosticHandlerFunction DiagnosticHandler) {
+std::string llvm::getBitcodeProducerString(MemoryBufferRef Buffer,
+                                           LLVMContext &Context) {
   std::unique_ptr<MemoryBuffer> Buf = MemoryBuffer::getMemBuffer(Buffer, false);
-  BitcodeReader R(Buf.release(), Context, DiagnosticHandler);
+  BitcodeReader R(Buf.release(), Context);
   ErrorOr<std::string> ProducerString = R.parseIdentificationBlock();
   if (ProducerString.getError())
     return "";
diff --git a/lib/LTO/LTOCodeGenerator.cpp b/lib/LTO/LTOCodeGenerator.cpp
index bf3cde59443..525ca37c2f1 100644
--- a/lib/LTO/LTOCodeGenerator.cpp
+++ b/lib/LTO/LTOCodeGenerator.cpp
@@ -66,9 +66,7 @@ const char* LTOCodeGenerator::getVersionString() {
 
 LTOCodeGenerator::LTOCodeGenerator(LLVMContext &Context)
     : Context(Context), MergedModule(new Module("ld-temp.o", Context)),
-      IRLinker(new Linker(*MergedModule, [this](const DiagnosticInfo &DI) {
-        MergedModule->getContext().diagnose(DI);
-      })) {
+      IRLinker(new Linker(*MergedModule)) {
   initializeLTOPasses();
 }
 
@@ -124,8 +122,7 @@ void LTOCodeGenerator::setModule(std::unique_ptr<LTOModule> Mod) {
   AsmUndefinedRefs.clear();
 
   MergedModule = Mod->takeModule();
-  IRLinker = llvm::make_unique<Linker>(*MergedModule,
-                                       IRLinker->getDiagnosticHandler());
+  IRLinker = make_unique<Linker>(*MergedModule);
 
   const std::vector<const char*> &Undefs = Mod->getAsmUndefinedRefs();
   for (int I = 0, E = Undefs.size(); I != E; ++I)
diff --git a/lib/LTO/LTOModule.cpp b/lib/LTO/LTOModule.cpp
index a6a3002e457..409b9490233 100644
--- a/lib/LTO/LTOModule.cpp
+++ b/lib/LTO/LTOModule.cpp
@@ -172,9 +172,8 @@ parseBitcodeFileImpl(MemoryBufferRef Buffer, LLVMContext &Context,
   // Parse lazily.
   std::unique_ptr<MemoryBuffer> LightweightBuf =
       MemoryBuffer::getMemBuffer(*MBOrErr, false);
-  ErrorOr<std::unique_ptr<Module>> M =
-      getLazyBitcodeModule(std::move(LightweightBuf), Context, nullptr,
-                           true /*ShouldLazyLoadMetadata*/);
+  ErrorOr<std::unique_ptr<Module>> M = getLazyBitcodeModule(
+      std::move(LightweightBuf), Context, true /*ShouldLazyLoadMetadata*/);
   if (std::error_code EC = M.getError())
     return EC;
   return std::move(*M);
diff --git a/lib/Linker/IRMover.cpp b/lib/Linker/IRMover.cpp
index 251cfb71894..8a11a0099b5 100644
--- a/lib/Linker/IRMover.cpp
+++ b/lib/Linker/IRMover.cpp
@@ -387,8 +387,6 @@ class IRLinker {
       Worklist.push_back(GV);
   }
 
-  DiagnosticHandlerFunction DiagnosticHandler;
-
   /// Set to true when all global value body linking is complete (including
   /// lazy linking). Used to prevent metadata linking from creating new
   /// references.
@@ -402,13 +400,13 @@ class IRLinker {
 
   /// Helper method for setting a message and returning an error code.
   bool emitError(const Twine &Message) {
-    DiagnosticHandler(LinkDiagnosticInfo(DS_Error, Message));
+    SrcM.getContext().diagnose(LinkDiagnosticInfo(DS_Error, Message));
     HasError = true;
     return true;
   }
 
   void emitWarning(const Twine &Message) {
-    DiagnosticHandler(LinkDiagnosticInfo(DS_Warning, Message));
+    SrcM.getContext().diagnose(LinkDiagnosticInfo(DS_Warning, Message));
   }
 
   /// Given a global in the source module, return the global in the
@@ -458,12 +456,10 @@ class IRLinker {
 
 public:
   IRLinker(Module &DstM, IRMover::IdentifiedStructTypeSet &Set, Module &SrcM,
-           DiagnosticHandlerFunction DiagnosticHandler,
            ArrayRef<GlobalValue *> ValuesToLink,
            std::function<void(GlobalValue &, IRMover::ValueAdder)> AddLazyFor)
       : DstM(DstM), SrcM(SrcM), AddLazyFor(AddLazyFor), TypeMap(Set),
-        GValMaterializer(this), LValMaterializer(this),
-        DiagnosticHandler(DiagnosticHandler) {
+        GValMaterializer(this), LValMaterializer(this) {
     for (GlobalValue *GV : ValuesToLink)
       maybeAdd(GV);
   }
@@ -1375,8 +1371,7 @@ bool IRMover::IdentifiedStructTypeSet::hasType(StructType *Ty) {
   return *I == Ty;
 }
 
-IRMover::IRMover(Module &M, DiagnosticHandlerFunction DiagnosticHandler)
-    : Composite(M), DiagnosticHandler(DiagnosticHandler) {
+IRMover::IRMover(Module &M) : Composite(M) {
   TypeFinder StructTypes;
   StructTypes.run(M, true);
   for (StructType *Ty : StructTypes) {
@@ -1390,8 +1385,8 @@ IRMover::IRMover(Module &M, DiagnosticHandlerFunction DiagnosticHandler)
 bool IRMover::move(
     Module &Src, ArrayRef<GlobalValue *> ValuesToLink,
     std::function<void(GlobalValue &, ValueAdder Add)> AddLazyFor) {
-  IRLinker TheLinker(Composite, IdentifiedStructTypes, Src, DiagnosticHandler,
-                     ValuesToLink, AddLazyFor);
+  IRLinker TheLinker(Composite, IdentifiedStructTypes, Src, ValuesToLink,
+                     AddLazyFor);
   bool RetCode = TheLinker.run();
   Composite.dropTriviallyDeadConstantArrays();
   return RetCode;
diff --git a/lib/Linker/LinkModules.cpp b/lib/Linker/LinkModules.cpp
index a596697e8f5..44b93696be1 100644
--- a/lib/Linker/LinkModules.cpp
+++ b/lib/Linker/LinkModules.cpp
@@ -17,6 +17,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/IR/DiagnosticPrinter.h"
+#include "llvm/IR/LLVMContext.h"
 using namespace llvm;
 
 namespace {
@@ -67,7 +68,7 @@ class ModuleLinker {
 
   /// Should we have mover and linker error diag info?
   bool emitError(const Twine &Message) {
-    Mover.getDiagnosticHandler()(LinkDiagnosticInfo(DS_Error, Message));
+    SrcM.getContext().diagnose(LinkDiagnosticInfo(DS_Error, Message));
     return true;
   }
 
@@ -786,8 +787,7 @@ bool ModuleLinker::run() {
   return false;
 }
 
-Linker::Linker(Module &M, DiagnosticHandlerFunction DiagnosticHandler)
-    : Mover(M, DiagnosticHandler) {}
+Linker::Linker(Module &M) : Mover(M) {}
 
 bool Linker::linkInModule(Module &Src, unsigned Flags,
                           const FunctionInfoIndex *Index,
@@ -805,20 +805,17 @@ bool Linker::linkInModule(Module &Src, unsigned Flags,
 /// true is returned and ErrorMsg (if not null) is set to indicate the problem.
 /// Upon failure, the Dest module could be in a modified state, and shouldn't be
 /// relied on to be consistent.
-bool Linker::linkModules(Module &Dest, Module &Src,
-                         DiagnosticHandlerFunction DiagnosticHandler,
-                         unsigned Flags) {
-  Linker L(Dest, DiagnosticHandler);
+bool Linker::linkModules(Module &Dest, Module &Src, unsigned Flags) {
+  Linker L(Dest);
   return L.linkInModule(Src, Flags);
 }
 
 std::unique_ptr<Module>
 llvm::renameModuleForThinLTO(std::unique_ptr<Module> &M,
-                             const FunctionInfoIndex *Index,
-                             DiagnosticHandlerFunction DiagnosticHandler) {
+                             const FunctionInfoIndex *Index) {
   std::unique_ptr<llvm::Module> RenamedModule(
       new llvm::Module(M->getModuleIdentifier(), M->getContext()));
-  Linker L(*RenamedModule.get(), DiagnosticHandler);
+  Linker L(*RenamedModule.get());
   if (L.linkInModule(*M.get(), llvm::Linker::Flags::None, Index))
     return nullptr;
   return RenamedModule;
@@ -828,19 +825,29 @@ llvm::renameModuleForThinLTO(std::unique_ptr<Module> &M,
 // C API.
 //===----------------------------------------------------------------------===//
 
+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 LLVMLinkModules(LLVMModuleRef Dest, LLVMModuleRef Src,
                          LLVMLinkerMode Unused, char **OutMessages) {
   Module *D = unwrap(Dest);
+  LLVMContext &Ctx = D->getContext();
+
+  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);
+
+  LLVMBool Result = Linker::linkModules(*D, *unwrap(Src));
 
-  LLVMBool Result = Linker::linkModules(
-      *D, *unwrap(Src), [&](const DiagnosticInfo &DI) { DI.print(DP); });
+  Ctx.setDiagnosticHandler(OldDiagnosticHandler, OldDiagnosticContext, true);
 
-  if (OutMessages && Result) {
-    Stream.flush();
+  if (OutMessages && Result)
     *OutMessages = strdup(Message.c_str());
-  }
   return Result;
 }
diff --git a/lib/Object/IRObjectFile.cpp b/lib/Object/IRObjectFile.cpp
index dcd385a1780..c35c413b3c3 100644
--- a/lib/Object/IRObjectFile.cpp
+++ b/lib/Object/IRObjectFile.cpp
@@ -309,7 +309,7 @@ llvm::object::IRObjectFile::create(MemoryBufferRef Object,
       MemoryBuffer::getMemBuffer(BCOrErr.get(), false));
 
   ErrorOr<std::unique_ptr<Module>> MOrErr =
-      getLazyBitcodeModule(std::move(Buff), Context, nullptr,
+      getLazyBitcodeModule(std::move(Buff), Context,
                            /*ShouldLazyLoadMetadata*/ true);
   if (std::error_code EC = MOrErr.getError())
     return EC;
diff --git a/lib/Transforms/IPO/FunctionImport.cpp b/lib/Transforms/IPO/FunctionImport.cpp
index 4d137e9fe58..fe58bdbed19 100644
--- a/lib/Transforms/IPO/FunctionImport.cpp
+++ b/lib/Transforms/IPO/FunctionImport.cpp
@@ -271,7 +271,7 @@ bool FunctionImporter::importFunctions(Module &DestModule) {
   /// Second step: for every call to an external function, try to import it.
 
   // Linker that will be used for importing function
-  Linker TheLinker(DestModule, DiagnosticHandler);
+  Linker TheLinker(DestModule);
 
   // Map of Module -> List of Function to import from the Module
   std::map<StringRef, std::pair<Module *, DenseSet<const GlobalValue *>>>
@@ -380,7 +380,7 @@ public:
     auto ModuleLoader = [&M](StringRef Identifier) {
       return loadFile(Identifier, M.getContext());
     };
-    FunctionImporter Importer(*Index, diagnosticHandler, ModuleLoader);
+    FunctionImporter Importer(*Index, ModuleLoader);
     return Importer.importFunctions(M);
 
     return false;
diff --git a/test/Linker/drop-debug.ll b/test/Linker/drop-debug.ll
index 9c1072a75de..06689872e12 100644
--- a/test/Linker/drop-debug.ll
+++ b/test/Linker/drop-debug.ll
@@ -3,4 +3,4 @@
 ;; drop-debug.bc was created from "void f(void) {}" with clang 3.5 and
 ; -gline-tables-only, so it contains old debug info.
 
-; CHECK: warning: ignoring debug info with an invalid version (1) in {{.*}}/Inputs/drop-debug.bc
+; CHECK: WARNING: ignoring debug info with an invalid version (1) in {{.*}}/Inputs/drop-debug.bc
diff --git a/tools/bugpoint/BugDriver.cpp b/tools/bugpoint/BugDriver.cpp
index 9edc242d470..39887d5d59d 100644
--- a/tools/bugpoint/BugDriver.cpp
+++ b/tools/bugpoint/BugDriver.cpp
@@ -15,7 +15,6 @@
 
 #include "BugDriver.h"
 #include "ToolRunner.h"
-#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/IRReader/IRReader.h"
@@ -113,12 +112,6 @@ std::unique_ptr<Module> llvm::parseInputFile(StringRef Filename,
   return Result;
 }
 
-static void diagnosticHandler(const DiagnosticInfo &DI) {
-  DiagnosticPrinterRawOStream DP(errs());
-  DI.print(DP);
-  errs() << '\n';
-}
-
 // This method takes the specified list of LLVM input files, attempts to load
 // them, either as assembly or bitcode, then link them together. It returns
 // true on failure (if, for example, an input bitcode file could not be
@@ -139,7 +132,7 @@ bool BugDriver::addSources(const std::vector<std::string> &Filenames) {
     if (!M.get()) return true;
 
     outs() << "Linking in input file: '" << Filenames[i] << "'\n";
-    if (Linker::linkModules(*Program, *M, diagnosticHandler))
+    if (Linker::linkModules(*Program, *M))
       return true;
   }
 
diff --git a/tools/bugpoint/Miscompilation.cpp b/tools/bugpoint/Miscompilation.cpp
index 5c9f0271cec..2c64e372256 100644
--- a/tools/bugpoint/Miscompilation.cpp
+++ b/tools/bugpoint/Miscompilation.cpp
@@ -18,7 +18,6 @@
 #include "llvm/Config/config.h"   // for HAVE_LINK_R
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
-#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
@@ -211,14 +210,6 @@ namespace {
   };
 }
 
-static void diagnosticHandler(const DiagnosticInfo &DI) {
-  DiagnosticPrinterRawOStream DP(errs());
-  DI.print(DP);
-  errs() << '\n';
-  if (DI.getSeverity() == DS_Error)
-    exit(1);
-}
-
 /// Given two modules, link them together and run the program, checking to see
 /// if the program matches the diff. If there is an error, return NULL. If not,
 /// return the merged module. The Broken argument will be set to true if the
@@ -230,7 +221,7 @@ static std::unique_ptr<Module> testMergedProgram(const BugDriver &BD,
                                                  std::unique_ptr<Module> M2,
                                                  std::string &Error,
                                                  bool &Broken) {
-  if (Linker::linkModules(*M1, *M2, diagnosticHandler))
+  if (Linker::linkModules(*M1, *M2))
     exit(1);
 
   // Execute the program.
@@ -396,8 +387,7 @@ static bool ExtractLoops(BugDriver &BD,
         MisCompFunctions.emplace_back(F->getName(), F->getFunctionType());
       }
 
-      if (Linker::linkModules(*ToNotOptimize, *ToOptimizeLoopExtracted,
-                              diagnosticHandler))
+      if (Linker::linkModules(*ToNotOptimize, *ToOptimizeLoopExtracted))
         exit(1);
 
       MiscompiledFunctions.clear();
@@ -424,8 +414,7 @@ static bool ExtractLoops(BugDriver &BD,
     // extraction both didn't break the program, and didn't mask the problem.
     // Replace the current program with the loop extracted version, and try to
     // extract another loop.
-    if (Linker::linkModules(*ToNotOptimize, *ToOptimizeLoopExtracted,
-                            diagnosticHandler))
+    if (Linker::linkModules(*ToNotOptimize, *ToOptimizeLoopExtracted))
       exit(1);
 
     // All of the Function*'s in the MiscompiledFunctions list are in the old
@@ -593,7 +582,7 @@ static bool ExtractBlocks(BugDriver &BD,
     if (!I->isDeclaration())
       MisCompFunctions.emplace_back(I->getName(), I->getFunctionType());
 
-  if (Linker::linkModules(*ProgClone, *Extracted, diagnosticHandler))
+  if (Linker::linkModules(*ProgClone, *Extracted))
     exit(1);
 
   // Set the new program and delete the old one.
diff --git a/tools/gold/gold-plugin.cpp b/tools/gold/gold-plugin.cpp
index e52606b8828..186097fccf6 100644
--- a/tools/gold/gold-plugin.cpp
+++ b/tools/gold/gold-plugin.cpp
@@ -879,7 +879,7 @@ static ld_plugin_status allSymbolsReadHook(raw_fd_ostream *ApiFile) {
   Context.setDiagnosticHandler(diagnosticHandlerForContext, nullptr, true);
 
   std::unique_ptr<Module> Combined(new Module("ld-temp.o", Context));
-  IRMover L(*Combined, diagnosticHandler);
+  IRMover L(*Combined);
 
   std::string DefaultTriple = sys::getDefaultTargetTriple();
 
diff --git a/tools/llvm-link/llvm-link.cpp b/tools/llvm-link/llvm-link.cpp
index 6f90f4056f7..8030f4c9037 100644
--- a/tools/llvm-link/llvm-link.cpp
+++ b/tools/llvm-link/llvm-link.cpp
@@ -141,6 +141,10 @@ static void diagnosticHandler(const DiagnosticInfo &DI) {
   errs() << '\n';
 }
 
+static void diagnosticHandlerWithContext(const DiagnosticInfo &DI, void *C) {
+  diagnosticHandler(DI);
+}
+
 /// Import any functions requested via the -import option.
 static bool importFunctions(const char *argv0, LLVMContext &Context,
                             Linker &L) {
@@ -265,11 +269,13 @@ int main(int argc, char **argv) {
   PrettyStackTraceProgram X(argc, argv);
 
   LLVMContext &Context = getGlobalContext();
+  Context.setDiagnosticHandler(diagnosticHandlerWithContext, nullptr, true);
+
   llvm_shutdown_obj Y;  // Call llvm_shutdown() on exit.
   cl::ParseCommandLineOptions(argc, argv, "llvm linker\n");
 
   auto Composite = make_unique<Module>("llvm-link", Context);
-  Linker L(*Composite, diagnosticHandler);
+  Linker L(*Composite);
 
   unsigned Flags = Linker::Flags::None;
   if (Internalize)
diff --git a/unittests/Linker/LinkModulesTest.cpp b/unittests/Linker/LinkModulesTest.cpp
index e56a692125e..ba8f4798172 100644
--- a/unittests/Linker/LinkModulesTest.cpp
+++ b/unittests/Linker/LinkModulesTest.cpp
@@ -71,7 +71,9 @@ protected:
   BasicBlock *ExitBB;
 };
 
-static void expectNoDiags(const DiagnosticInfo &DI) { EXPECT_TRUE(false); }
+static void expectNoDiags(const DiagnosticInfo &DI, void *C) {
+  EXPECT_TRUE(false);
+}
 
 TEST_F(LinkModuleTest, BlockAddress) {
   IRBuilder<> Builder(EntryBB);
@@ -95,7 +97,8 @@ TEST_F(LinkModuleTest, BlockAddress) {
   Builder.CreateRet(ConstantPointerNull::get(Type::getInt8PtrTy(Ctx)));
 
   Module *LinkedModule = new Module("MyModuleLinked", Ctx);
-  Linker::linkModules(*LinkedModule, *M, expectNoDiags);
+  Ctx.setDiagnosticHandler(expectNoDiags);
+  Linker::linkModules(*LinkedModule, *M);
 
   // Delete the original module.
   M.reset();
@@ -171,13 +174,15 @@ static Module *getInternal(LLVMContext &Ctx) {
 TEST_F(LinkModuleTest, EmptyModule) {
   std::unique_ptr<Module> InternalM(getInternal(Ctx));
   std::unique_ptr<Module> EmptyM(new Module("EmptyModule1", Ctx));
-  Linker::linkModules(*EmptyM, *InternalM, expectNoDiags);
+  Ctx.setDiagnosticHandler(expectNoDiags);
+  Linker::linkModules(*EmptyM, *InternalM);
 }
 
 TEST_F(LinkModuleTest, EmptyModule2) {
   std::unique_ptr<Module> InternalM(getInternal(Ctx));
   std::unique_ptr<Module> EmptyM(new Module("EmptyModule1", Ctx));
-  Linker::linkModules(*InternalM, *EmptyM, expectNoDiags);
+  Ctx.setDiagnosticHandler(expectNoDiags);
+  Linker::linkModules(*InternalM, *EmptyM);
 }
 
 TEST_F(LinkModuleTest, TypeMerge) {
@@ -192,7 +197,8 @@ TEST_F(LinkModuleTest, TypeMerge) {
                       "@t2 = weak global %t zeroinitializer\n";
   std::unique_ptr<Module> M2 = parseAssemblyString(M2Str, Err, C);
 
-  Linker::linkModules(*M1, *M2, [](const llvm::DiagnosticInfo &) {});
+  Ctx.setDiagnosticHandler(expectNoDiags);
+  Linker::linkModules(*M1, *M2);
 
   EXPECT_EQ(M1->getNamedGlobal("t1")->getType(),
             M1->getNamedGlobal("t2")->getType());
@@ -269,7 +275,8 @@ TEST_F(LinkModuleTest, MoveDistinctMDs) {
   // Link into destination module.
   auto Dst = llvm::make_unique<Module>("Linked", C);
   ASSERT_TRUE(Dst.get());
-  Linker::linkModules(*Dst, *Src, [](const llvm::DiagnosticInfo &) {});
+  Ctx.setDiagnosticHandler(expectNoDiags);
+  Linker::linkModules(*Dst, *Src);
 
   // Check that distinct metadata was moved, not cloned.  Even !4, the uniqued
   // node, should effectively be moved, since its only operand hasn't changed.