1 files changed, 60 insertions, 31 deletions

diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index b874308afa3..fc134e4b5d9 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -124,7 +124,7 @@ bool MetadataTracking::track(void *Ref, Metadata &MD, OwnerTy Owner) {
   assert(Ref && "Expected live reference");
   assert((Owner || *static_cast<Metadata **>(Ref) == &MD) &&
          "Reference without owner must be direct");
-  if (auto *R = ReplaceableMetadataImpl::get(MD)) {
+  if (auto *R = ReplaceableMetadataImpl::getOrCreate(MD)) {
     R->addRef(Ref, Owner);
     return true;
   }
@@ -133,7 +133,7 @@ bool MetadataTracking::track(void *Ref, Metadata &MD, OwnerTy Owner) {
 
 void MetadataTracking::untrack(void *Ref, Metadata &MD) {
   assert(Ref && "Expected live reference");
-  if (auto *R = ReplaceableMetadataImpl::get(MD))
+  if (auto *R = ReplaceableMetadataImpl::getIfExists(MD))
     R->dropRef(Ref);
 }
 
@@ -141,15 +141,17 @@ bool MetadataTracking::retrack(void *Ref, Metadata &MD, void *New) {
   assert(Ref && "Expected live reference");
   assert(New && "Expected live reference");
   assert(Ref != New && "Expected change");
-  if (auto *R = ReplaceableMetadataImpl::get(MD)) {
+  if (auto *R = ReplaceableMetadataImpl::getIfExists(MD)) {
     R->moveRef(Ref, New, MD);
     return true;
   }
+  assert(!isReplaceable(MD) &&
+         "Expected un-replaceable metadata, since we didn't move a reference");
   return false;
 }
 
 bool MetadataTracking::isReplaceable(const Metadata &MD) {
-  return ReplaceableMetadataImpl::get(const_cast<Metadata &>(MD));
+  return ReplaceableMetadataImpl::isReplaceable(MD);
 }
 
 void ReplaceableMetadataImpl::addRef(void *Ref, OwnerTy Owner) {
@@ -268,9 +270,21 @@ void ReplaceableMetadataImpl::resolveAllUses(bool ResolveUsers) {
   }
 }
 
-ReplaceableMetadataImpl *ReplaceableMetadataImpl::get(Metadata &MD) {
+ReplaceableMetadataImpl *ReplaceableMetadataImpl::getOrCreate(Metadata &MD) {
   if (auto *N = dyn_cast<MDNode>(&MD))
-    return N->Context.getReplaceableUses();
+    return N->isResolved() ? nullptr : N->Context.getOrCreateReplaceableUses();
+  return dyn_cast<ValueAsMetadata>(&MD);
+}
+
+ReplaceableMetadataImpl *ReplaceableMetadataImpl::getIfExists(Metadata &MD) {
+  if (auto *N = dyn_cast<MDNode>(&MD))
+    return N->isResolved() ? nullptr : N->Context.getReplaceableUses();
+  return dyn_cast<ValueAsMetadata>(&MD);
+}
+
+bool ReplaceableMetadataImpl::isReplaceable(const Metadata &MD) {
+  if (auto *N = dyn_cast<MDNode>(&MD))
+    return !N->isResolved();
   return dyn_cast<ValueAsMetadata>(&MD);
 }
 
@@ -452,16 +466,12 @@ MDNode::MDNode(LLVMContext &Context, unsigned ID, StorageType Storage,
   for (Metadata *MD : Ops2)
     setOperand(Op++, MD);
 
-  if (isDistinct())
+  if (!isUniqued())
     return;
 
-  if (isUniqued())
-    // Check whether any operands are unresolved, requiring re-uniquing.  If
-    // not, don't support RAUW.
-    if (!countUnresolvedOperands())
-      return;
-
-  this->Context.makeReplaceable(make_unique<ReplaceableMetadataImpl>(Context));
+  // Count the unresolved operands.  If there are any, RAUW support will be
+  // added lazily on first reference.
+  countUnresolvedOperands();
 }
 
 TempMDNode MDNode::clone() const {
@@ -481,10 +491,10 @@ static bool isOperandUnresolved(Metadata *Op) {
   return false;
 }
 
-unsigned MDNode::countUnresolvedOperands() {
+void MDNode::countUnresolvedOperands() {
   assert(NumUnresolved == 0 && "Expected unresolved ops to be uncounted");
+  assert(isUniqued() && "Expected this to be uniqued");
   NumUnresolved = std::count_if(op_begin(), op_end(), isOperandUnresolved);
-  return NumUnresolved;
 }
 
 void MDNode::makeUniqued() {
@@ -497,8 +507,11 @@ void MDNode::makeUniqued() {
 
   // Make this 'uniqued'.
   Storage = Uniqued;
-  if (!countUnresolvedOperands())
-    resolve();
+  countUnresolvedOperands();
+  if (!NumUnresolved) {
+    dropReplaceableUses();
+    assert(isResolved() && "Expected this to be resolved");
+  }
 
   assert(isUniqued() && "Expected this to be uniqued");
 }
@@ -507,9 +520,8 @@ void MDNode::makeDistinct() {
   assert(isTemporary() && "Expected this to be temporary");
   assert(!isResolved() && "Expected this to be unresolved");
 
-  // Pretend to be uniqued, resolve the node, and then store in distinct table.
-  Storage = Uniqued;
-  resolve();
+  // Drop RAUW support and store as a distinct node.
+  dropReplaceableUses();
   storeDistinctInContext();
 
   assert(isDistinct() && "Expected this to be distinct");
@@ -520,16 +532,22 @@ void MDNode::resolve() {
   assert(isUniqued() && "Expected this to be uniqued");
   assert(!isResolved() && "Expected this to be unresolved");
 
-  // Move the map, so that this immediately looks resolved.
-  auto Uses = Context.takeReplaceableUses();
   NumUnresolved = 0;
+  dropReplaceableUses();
+
   assert(isResolved() && "Expected this to be resolved");
+}
+
+void MDNode::dropReplaceableUses() {
+  assert(!NumUnresolved && "Unexpected unresolved operand");
 
-  // Drop RAUW support.
-  Uses->resolveAllUses();
+  // Drop any RAUW support.
+  if (Context.hasReplaceableUses())
+    Context.takeReplaceableUses()->resolveAllUses();
 }
 
 void MDNode::resolveAfterOperandChange(Metadata *Old, Metadata *New) {
+  assert(isUniqued() && "Expected this to be uniqued");
   assert(NumUnresolved != 0 && "Expected unresolved operands");
 
   // Check if an operand was resolved.
@@ -542,9 +560,17 @@ void MDNode::resolveAfterOperandChange(Metadata *Old, Metadata *New) {
 }
 
 void MDNode::decrementUnresolvedOperandCount() {
-  if (!--NumUnresolved)
-    // Last unresolved operand has just been resolved.
-    resolve();
+  assert(!isResolved() && "Expected this to be unresolved");
+  if (isTemporary())
+    return;
+
+  assert(isUniqued() && "Expected this to be uniqued");
+  if (--NumUnresolved)
+    return;
+
+  // Last unresolved operand has just been resolved.
+  dropReplaceableUses();
+  assert(isResolved() && "Expected this to become resolved");
 }
 
 void MDNode::resolveCycles() {
@@ -619,7 +645,7 @@ void MDTuple::recalculateHash() {
 void MDNode::dropAllReferences() {
   for (unsigned I = 0, E = NumOperands; I != E; ++I)
     setOperand(I, nullptr);
-  if (!isResolved()) {
+  if (Context.hasReplaceableUses()) {
     Context.getReplaceableUses()->resolveAllUses(/* ResolveUsers */ false);
     (void)Context.takeReplaceableUses();
   }
@@ -665,7 +691,8 @@ void MDNode::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);
-    Context.getReplaceableUses()->replaceAllUsesWith(Uniqued);
+    if (Context.hasReplaceableUses())
+      Context.getReplaceableUses()->replaceAllUsesWith(Uniqued);
     deleteAsSubclass();
     return;
   }
@@ -763,8 +790,10 @@ void MDNode::deleteTemporary(MDNode *N) {
 }
 
 void MDNode::storeDistinctInContext() {
-  assert(isResolved() && "Expected resolved nodes");
+  assert(!Context.hasReplaceableUses() && "Unexpected replaceable uses");
+  assert(!NumUnresolved && "Unexpected unresolved nodes");
   Storage = Distinct;
+  assert(isResolved() && "Expected this to be resolved");
 
   // Reset the hash.
   switch (getMetadataID()) {