3 files changed, 26 insertions, 22 deletions

diff --git a/include/llvm/Metadata.h b/include/llvm/Metadata.h
index e5363228a3..cecb7dadaf 100644
--- a/include/llvm/Metadata.h
+++ b/include/llvm/Metadata.h
@@ -108,9 +108,6 @@ class MDNode : public Value, public FoldingSetNode {
   /// node with T.
   void replaceOperand(MDNodeOperand *Op, Value *NewVal);
   ~MDNode();
-  /// replaceAllOperandsWithNull - This is used while destroying llvm context to 
-  /// gracefully delete all nodes. This method replaces all operands with null.
-  void replaceAllOperandsWithNull();
 
 protected:
   explicit MDNode(LLVMContext &C, Value *const *Vals, unsigned NumVals,
@@ -166,9 +163,7 @@ private:
   bool isNotUniqued() const { 
     return (getSubclassDataFromValue() & NotUniquedBit) != 0;
   }
-  void setIsNotUniqued() {
-    setValueSubclassData(getSubclassDataFromValue() | NotUniquedBit);
-  }
+  void setIsNotUniqued();
   
   // Shadow Value::setValueSubclassData with a private forwarding method so that
   // any future subclasses cannot accidentally use it.
diff --git a/lib/VMCore/LLVMContextImpl.h b/lib/VMCore/LLVMContextImpl.h
index 9887f28821..9978f40ed0 100644
--- a/lib/VMCore/LLVMContextImpl.h
+++ b/lib/VMCore/LLVMContextImpl.h
@@ -105,6 +105,11 @@ public:
   StringMap<MDString*> MDStringCache;
   
   FoldingSet<MDNode> MDNodeSet;
+  // MDNodes may be uniqued or not uniqued.  When they're not uniqued, they
+  // aren't in the MDNodeSet, but they're still shared between objects, so no
+  // one object can destroy them.  This set allows us to at least destroy them
+  // on Context destruction.
+  SmallPtrSet<MDNode*, 1> NonUniquedMDNodes;
   
   ConstantUniqueMap<char, Type, ConstantAggregateZero> AggZeroConstants;
 
@@ -235,17 +240,21 @@ public:
       (*I)->AbstractTypeUsers.clear();
       delete *I;
     }
-    // Destroy MDNode operands first.
+    // Destroy MDNodes.  ~MDNode can move and remove nodes between the MDNodeSet
+    // and the NonUniquedMDNodes sets, so copy the values out first.
+    SmallVector<MDNode*, 8> MDNodes;
+    MDNodes.reserve(MDNodeSet.size() + NonUniquedMDNodes.size());
     for (FoldingSetIterator<MDNode> I = MDNodeSet.begin(), E = MDNodeSet.end();
-         I != E;) {
-      MDNode *N = &(*I);
-      ++I;
-      N->replaceAllOperandsWithNull();
+         I != E; ++I) {
+      MDNodes.push_back(&*I);
     }
-    while (!MDNodeSet.empty()) {
-      MDNode *N = &(*MDNodeSet.begin());
-      N->destroy();
+    MDNodes.append(NonUniquedMDNodes.begin(), NonUniquedMDNodes.end());
+    for (SmallVector<MDNode*, 8>::iterator I = MDNodes.begin(),
+           E = MDNodes.end(); I != E; ++I) {
+      (*I)->destroy();
     }
+    assert(MDNodeSet.empty() && NonUniquedMDNodes.empty() &&
+           "Destroying all MDNodes didn't empty the Context's sets.");
     // Destroy MDStrings.
     for (StringMap<MDString*>::iterator I = MDStringCache.begin(),
            E = MDStringCache.end(); I != E; ++I) {
diff --git a/lib/VMCore/Metadata.cpp b/lib/VMCore/Metadata.cpp
index a08c45480b..faf83e6588 100644
--- a/lib/VMCore/Metadata.cpp
+++ b/lib/VMCore/Metadata.cpp
@@ -110,8 +110,10 @@ MDNode::MDNode(LLVMContext &C, Value *const *Vals, unsigned NumVals,
 MDNode::~MDNode() {
   assert((getSubclassDataFromValue() & DestroyFlag) != 0 &&
          "Not being destroyed through destroy()?");
-  if (!isNotUniqued()) {
-    LLVMContextImpl *pImpl = getType()->getContext().pImpl;
+  LLVMContextImpl *pImpl = getType()->getContext().pImpl;
+  if (isNotUniqued()) {
+    pImpl->NonUniquedMDNodes.erase(this);
+  } else {
     pImpl->MDNodeSet.RemoveNode(this);
   }
 
@@ -257,12 +259,10 @@ void MDNode::Profile(FoldingSetNodeID &ID) const {
     ID.AddPointer(getOperand(i));
 }
 
-// replaceAllOperandsWithNull - This is used while destroying llvm context to 
-// gracefully delete all nodes. This method replaces all operands with null.
-void MDNode::replaceAllOperandsWithNull() {
-  for (MDNodeOperand *Op = getOperandPtr(this, 0), *E = Op+NumOperands;
-       Op != E; ++Op)
-    replaceOperand(Op, 0);
+void MDNode::setIsNotUniqued() {
+  setValueSubclassData(getSubclassDataFromValue() | NotUniquedBit);
+  LLVMContextImpl *pImpl = getType()->getContext().pImpl;
+  pImpl->NonUniquedMDNodes.insert(this);
 }
 
 // Replace value from this node's operand list.