3 files changed, 32 insertions, 1 deletions

diff --git a/llvm/include/llvm/IR/LeakDetector.h b/llvm/include/llvm/IR/LeakDetector.h
index fb3202e902a..b014fe081f0 100644
--- a/llvm/include/llvm/IR/LeakDetector.h
+++ b/llvm/include/llvm/IR/LeakDetector.h
@@ -28,6 +28,7 @@ namespace llvm {
 
 class LLVMContext;
 class Value;
+class MDNode;
 
 struct LeakDetector {
   /// addGarbageObject - Add a pointer to the internal set of "garbage" object
@@ -72,6 +73,18 @@ private:
   static void addGarbageObjectImpl(const Value *Object);
   static void removeGarbageObjectImpl(const Value *Object);
 
+  /// Overload the normal methods to work better with MDNode* to improve error
+  /// messages.
+  ///
+  /// For better or worse, this hides errors when other types are added as
+  /// garbage, deleted without being removed, and an MDNode is allocated in the
+  /// same spot.
+  ///
+  /// \note Only handle \a MDNode for now, since we can't always get access to
+  /// an \a LLVMContext for other \a Metadata types.
+  static void addGarbageObjectImpl(const MDNode *Object);
+  static void removeGarbageObjectImpl(const MDNode *Object);
+
   static void addGarbageObjectImpl(void *Object);
   static void removeGarbageObjectImpl(void *Object);
   static void checkForGarbageImpl(LLVMContext &C, const std::string &Message);
diff --git a/llvm/lib/IR/LeakDetector.cpp b/llvm/lib/IR/LeakDetector.cpp
index 6f71627fcf9..be07a481883 100644
--- a/llvm/lib/IR/LeakDetector.cpp
+++ b/llvm/lib/IR/LeakDetector.cpp
@@ -14,6 +14,7 @@
 #include "llvm/IR/LeakDetector.h"
 #include "LLVMContextImpl.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ManagedStatic.h"
@@ -39,6 +40,11 @@ void LeakDetector::addGarbageObjectImpl(const Value *Object) {
   pImpl->LLVMObjects.addGarbage(Object);
 }
 
+void LeakDetector::addGarbageObjectImpl(const MDNode *Object) {
+  LLVMContextImpl *pImpl = Object->getContext().pImpl;
+  pImpl->LLVMMDObjects.addGarbage(Object);
+}
+
 void LeakDetector::removeGarbageObjectImpl(void *Object) {
   sys::SmartScopedLock<true> Lock(*ObjectsLock);
   Objects->removeGarbage(Object);
@@ -49,6 +55,11 @@ void LeakDetector::removeGarbageObjectImpl(const Value *Object) {
   pImpl->LLVMObjects.removeGarbage(Object);
 }
 
+void LeakDetector::removeGarbageObjectImpl(const MDNode *Object) {
+  LLVMContextImpl *pImpl = Object->getContext().pImpl;
+  pImpl->LLVMMDObjects.removeGarbage(Object);
+}
+
 void LeakDetector::checkForGarbageImpl(LLVMContext &Context, 
                                        const std::string &Message) {
   LLVMContextImpl *pImpl = Context.pImpl;
@@ -56,10 +67,12 @@ void LeakDetector::checkForGarbageImpl(LLVMContext &Context,
   
   Objects->setName("GENERIC");
   pImpl->LLVMObjects.setName("LLVM");
+  pImpl->LLVMMDObjects.setName("LLVM-MD");
   
   // use non-short-circuit version so that both checks are performed
   if (Objects->hasGarbage(Message) |
-      pImpl->LLVMObjects.hasGarbage(Message))
+      pImpl->LLVMObjects.hasGarbage(Message) |
+      pImpl->LLVMMDObjects.hasGarbage(Message))
     errs() << "\nThis is probably because you removed an object, but didn't "
            << "delete it.  Please check your code for memory leaks.\n";
 
diff --git a/llvm/lib/IR/LeaksContext.h b/llvm/lib/IR/LeaksContext.h
index 3e485abdfd1..47704fa1a76 100644
--- a/llvm/lib/IR/LeaksContext.h
+++ b/llvm/lib/IR/LeaksContext.h
@@ -16,6 +16,7 @@
 #define LLVM_LIB_IR_LEAKSCONTEXT_H
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -31,6 +32,10 @@ struct PrinterTrait<Value> {
   static void print(const Value* P) { errs() << *P; }
 };
 
+template <> struct PrinterTrait<Metadata> {
+  static void print(const Metadata *P) { P->print(errs()); }
+};
+
 template <typename T>
 struct LeakDetectorImpl {
   explicit LeakDetectorImpl(const char* const name = "") :