1 files changed, 56 insertions, 11 deletions

diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index fa1f302568f..13844313925 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -525,8 +525,8 @@ void UniquableMDNode::handleChangedOperand(void *Ref, Metadata *New) {
     return;
   }
 
-  auto &Store = getContext().pImpl->MDTuples;
-  Store.erase(cast<MDTuple>(this));
+  // This node is uniqued.
+  eraseFromStore();
 
   Metadata *Old = getOperand(Op);
   setOperand(Op, New);
@@ -540,15 +540,10 @@ void UniquableMDNode::handleChangedOperand(void *Ref, Metadata *New) {
   }
 
   // Re-unique the node.
-  cast<MDTuple>(this)->recalculateHash();
-  MDTupleInfo::KeyTy Key(cast<MDTuple>(this));
-  auto I = Store.find_as(Key);
-  if (I == Store.end()) {
-    Store.insert(cast<MDTuple>(this));
-
+  auto *Uniqued = uniquify();
+  if (Uniqued == this) {
     if (!isResolved())
       resolveAfterOperandChange(Old, New);
-
     return;
   }
 
@@ -560,8 +555,8 @@ 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(*I);
-    delete cast<MDTuple>(this);
+    ReplaceableUses->replaceAllUsesWith(Uniqued);
+    deleteAsSubclass();
     return;
   }
 
@@ -569,6 +564,41 @@ void UniquableMDNode::handleChangedOperand(void *Ref, Metadata *New) {
   storeDistinctInContext();
 }
 
+void UniquableMDNode::deleteAsSubclass() {
+  switch (getMetadataID()) {
+  default:
+    llvm_unreachable("Invalid subclass of UniquableMDNode");
+#define HANDLE_UNIQUABLE_LEAF(CLASS)                                           \
+  case CLASS##Kind:                                                            \
+    delete cast<CLASS>(this);                                                  \
+    break;
+#include "llvm/IR/Metadata.def"
+  }
+}
+
+UniquableMDNode *UniquableMDNode::uniquify() {
+  switch (getMetadataID()) {
+  default:
+    llvm_unreachable("Invalid subclass of UniquableMDNode");
+#define HANDLE_UNIQUABLE_LEAF(CLASS)                                           \
+  case CLASS##Kind:                                                            \
+    return cast<CLASS>(this)->uniquifyImpl();
+#include "llvm/IR/Metadata.def"
+  }
+}
+
+void UniquableMDNode::eraseFromStore() {
+  switch (getMetadataID()) {
+  default:
+    llvm_unreachable("Invalid subclass of UniquableMDNode");
+#define HANDLE_UNIQUABLE_LEAF(CLASS)                                           \
+  case CLASS##Kind:                                                            \
+    cast<CLASS>(this)->eraseFromStoreImpl();                                   \
+    break;
+#include "llvm/IR/Metadata.def"
+  }
+}
+
 MDTuple *MDTuple::getImpl(LLVMContext &Context, ArrayRef<Metadata *> MDs,
                           bool ShouldCreate) {
   MDTupleInfo::KeyTy Key(MDs);
@@ -593,6 +623,21 @@ MDTuple *MDTuple::getDistinct(LLVMContext &Context, ArrayRef<Metadata *> MDs) {
   return N;
 }
 
+MDTuple *MDTuple::uniquifyImpl() {
+  recalculateHash();
+  MDTupleInfo::KeyTy Key(this);
+
+  auto &Store = getContext().pImpl->MDTuples;
+  auto I = Store.find_as(Key);
+  if (I == Store.end()) {
+    Store.insert(this);
+    return this;
+  }
+  return *I;
+}
+
+void MDTuple::eraseFromStoreImpl() { getContext().pImpl->MDTuples.erase(this); }
+
 MDNodeFwdDecl *MDNode::getTemporary(LLVMContext &Context,
                                     ArrayRef<Metadata *> MDs) {
   return MDNodeFwdDecl::get(Context, MDs);