2 files changed, 11 insertions, 10 deletions
diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index 93778c2..5b19638 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -401,6 +401,10 @@ MDNode::MDNode(LLVMContext &Context, unsigned ID, StorageType Storage,
       MDNodeSubclassData(0) {
   for (unsigned I = 0, E = MDs.size(); I != E; ++I)
     setOperand(I, MDs[I]);
+
+  if (Storage == Temporary)
+    this->Context.makeReplaceable(
+        make_unique<ReplaceableMetadataImpl>(Context));
 }
 
 bool MDNode::isResolved() const {
@@ -429,7 +433,7 @@ UniquableMDNode::UniquableMDNode(LLVMContext &C, unsigned ID,
   if (!NumUnresolved)
     return;
 
-  ReplaceableUses.reset(new ReplaceableMetadataImpl);
+  this->Context.makeReplaceable(make_unique<ReplaceableMetadataImpl>(C));
   SubclassData32 = NumUnresolved;
 }
 
@@ -437,7 +441,7 @@ void UniquableMDNode::resolve() {
   assert(!isResolved() && "Expected this to be unresolved");
 
   // Move the map, so that this immediately looks resolved.
-  auto Uses = std::move(ReplaceableUses);
+  auto Uses = Context.takeReplaceableUses();
   SubclassData32 = 0;
   assert(isResolved() && "Expected this to be resolved");
 
@@ -499,8 +503,8 @@ void MDNode::dropAllReferences() {
     setOperand(I, nullptr);
   if (auto *N = dyn_cast<UniquableMDNode>(this))
     if (!N->isResolved()) {
-      N->ReplaceableUses->resolveAllUses(/* ResolveUsers */ false);
-      N->ReplaceableUses.reset();
+      N->Context.getReplaceableUses()->resolveAllUses(/* ResolveUsers */ false);
+      (void)N->Context.takeReplaceableUses();
     }
 }
 
@@ -563,7 +567,7 @@ void UniquableMDNode::handleChangedOperand(void *Ref, Metadata *New) {
     // dropAllReferences(), but we still need the use-list).
     for (unsigned O = 0, E = getNumOperands(); O != E; ++O)
       setOperand(O, nullptr);
-    ReplaceableUses->replaceAllUsesWith(Uniqued);
+    Context.getReplaceableUses()->replaceAllUsesWith(Uniqued);
     deleteAsSubclass();
     return;
   }
diff --git a/llvm/lib/IR/MetadataTracking.cpp b/llvm/lib/IR/MetadataTracking.cpp
index ba97ca03..47f0b93 100644
--- a/llvm/lib/IR/MetadataTracking.cpp
+++ b/llvm/lib/IR/MetadataTracking.cpp
@@ -17,11 +17,8 @@
 using namespace llvm;
 
 ReplaceableMetadataImpl *ReplaceableMetadataImpl::get(Metadata &MD) {
-  if (auto *N = dyn_cast<MDNode>(&MD)) {
-    if (auto *U = dyn_cast<UniquableMDNode>(N))
-      return U->ReplaceableUses.get();
-    return cast<MDNodeFwdDecl>(N);
-  }
+  if (auto *N = dyn_cast<MDNode>(&MD))
+    return N->Context.getReplaceableUses();
   return dyn_cast<ValueAsMetadata>(&MD);
 }