IR: Allow temporary nodes to become uniqued or distinct

Add `MDNode::replaceWithUniqued()` and `MDNode::replaceWithDistinct()`,
which mutate temporary nodes to become uniqued or distinct.  On uniquing
collisions, the unique version is returned and the node is deleted.

This takes advantage of temporary nodes being folded back in, and should
let me clean up some awkward logic in `MapMetadata()`.

llvm-svn: 226510

include/llvm/IR/Metadata.h

@@ -744,6 +744,24 @@ public:
     Context.getReplaceableUses()->replaceAllUsesWith(MD);
   }
 
+  /// \brief Replace a temporary node with a uniqued one.
+  ///
+  /// Create a uniqued version of \c N -- in place, if possible -- and return
+  /// it.  Takes ownership of the temporary node.
+  template <class T>
+  static typename std::enable_if<std::is_base_of<UniquableMDNode, T>::value,
+                                 T *>::type
+  replaceWithUniqued(std::unique_ptr<T, TempMDNodeDeleter> N);
+
+  /// \brief Replace a temporary node with a distinct one.
+  ///
+  /// Create a distinct version of \c N -- in place, if possible -- and return
+  /// it.  Takes ownership of the temporary node.
+  template <class T>
+  static typename std::enable_if<std::is_base_of<UniquableMDNode, T>::value,
+                                 T *>::type
+  replaceWithDistinct(std::unique_ptr<T, TempMDNodeDeleter> N);
+
 protected:
   /// \brief Set an operand.
   ///
@@ -788,6 +806,30 @@ public:
   static MDNode *getMostGenericRange(MDNode *A, MDNode *B);
 };
 
+template <class NodeTy>
+typename std::enable_if<std::is_base_of<UniquableMDNode, NodeTy>::value,
+                        NodeTy *>::type
+MDNode::replaceWithUniqued(std::unique_ptr<NodeTy, TempMDNodeDeleter> Node) {
+  // Try to uniquify in place.
+  UniquableMDNode *UniquedNode = Node->uniquify();
+  if (UniquedNode == Node.get()) {
+    Node->makeUniqued();
+    return Node.release();
+  }
+
+  // Collision, so RAUW instead.
+  Node->replaceAllUsesWith(UniquedNode);
+  return cast<NodeTy>(UniquedNode);
+}
+
+template <class NodeTy>
+typename std::enable_if<std::is_base_of<UniquableMDNode, NodeTy>::value,
+                        NodeTy *>::type
+MDNode::replaceWithDistinct(std::unique_ptr<NodeTy, TempMDNodeDeleter> Node) {
+  Node->makeDistinct();
+  return Node.release();
+}
+
 /// \brief Uniquable metadata node.
 ///
 /// A uniquable metadata node.  This contains the basic functionality
@@ -844,6 +886,19 @@ private:
   void decrementUnresolvedOperandCount();
   unsigned countUnresolvedOperands() const;
 
+  /// \brief Mutate this to be "uniqued".
+  ///
+  /// Mutate this so that \a isUniqued().
+  /// \pre \a isTemporary().
+  /// \pre already added to uniquing set.
+  void makeUniqued();
+
+  /// \brief Mutate this to be "distinct".
+  ///
+  /// Mutate this so that \a isDistinct().
+  /// \pre \a isTemporary().
+  void makeDistinct();
+
   void deleteAsSubclass();
   UniquableMDNode *uniquify();
   void eraseFromStore();

lib/IR/Metadata.cpp

@@ -436,6 +436,33 @@ unsigned UniquableMDNode::countUnresolvedOperands() const {
   return NumUnresolved;
 }
 
+void UniquableMDNode::makeUniqued() {
+  assert(isTemporary() && "Expected this to be temporary");
+  assert(!isResolved() && "Expected this to be unresolved");
+
+  // Make this 'uniqued'.
+  Storage = Uniqued;
+  if (unsigned NumUnresolved = countUnresolvedOperands())
+    SubclassData32 = NumUnresolved;
+  else
+    resolve();
+
+  assert(isUniqued() && "Expected this to be uniqued");
+}
+
+void UniquableMDNode::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();
+  storeDistinctInContext();
+
+  assert(isDistinct() && "Expected this to be distinct");
+  assert(isResolved() && "Expected this to be resolved");
+}
+
 void UniquableMDNode::resolve() {
   assert(isUniqued() && "Expected this to be uniqued");
   assert(!isResolved() && "Expected this to be unresolved");

unittests/IR/MetadataTest.cpp

@@ -438,6 +438,85 @@ TEST_F(MDNodeTest, replaceResolvedOperand) {
   Temp->replaceAllUsesWith(nullptr);
 }
 
+TEST_F(MDNodeTest, replaceWithUniqued) {
+  auto *Empty = MDTuple::get(Context, None);
+  MDTuple *FirstUniqued;
+  {
+    Metadata *Ops[] = {Empty};
+    auto Temp = MDTuple::getTemporary(Context, Ops);
+    EXPECT_TRUE(Temp->isTemporary());
+
+    // Don't expect a collision.
+    auto *Current = Temp.get();
+    FirstUniqued = MDNode::replaceWithUniqued(std::move(Temp));
+    EXPECT_TRUE(FirstUniqued->isUniqued());
+    EXPECT_TRUE(FirstUniqued->isResolved());
+    EXPECT_EQ(Current, FirstUniqued);
+  }
+  {
+    Metadata *Ops[] = {Empty};
+    auto Temp = MDTuple::getTemporary(Context, Ops);
+    EXPECT_TRUE(Temp->isTemporary());
+
+    // Should collide with Uniqued above this time.
+    auto *Uniqued = MDNode::replaceWithUniqued(std::move(Temp));
+    EXPECT_TRUE(Uniqued->isUniqued());
+    EXPECT_TRUE(Uniqued->isResolved());
+    EXPECT_EQ(FirstUniqued, Uniqued);
+  }
+  {
+    auto Unresolved = MDTuple::getTemporary(Context, None);
+    Metadata *Ops[] = {Unresolved.get()};
+    auto Temp = MDTuple::getTemporary(Context, Ops);
+    EXPECT_TRUE(Temp->isTemporary());
+
+    // Shouldn't be resolved.
+    auto *Uniqued = MDNode::replaceWithUniqued(std::move(Temp));
+    EXPECT_TRUE(Uniqued->isUniqued());
+    EXPECT_FALSE(Uniqued->isResolved());
+
+    // Should be a different node.
+    EXPECT_NE(FirstUniqued, Uniqued);
+
+    // Should resolve when we update its node (note: be careful to avoid a
+    // collision with any other nodes above).
+    Uniqued->replaceOperandWith(0, nullptr);
+    EXPECT_TRUE(Uniqued->isResolved());
+  }
+}
+
+TEST_F(MDNodeTest, replaceWithDistinct) {
+  {
+    auto *Empty = MDTuple::get(Context, None);
+    Metadata *Ops[] = {Empty};
+    auto Temp = MDTuple::getTemporary(Context, Ops);
+    EXPECT_TRUE(Temp->isTemporary());
+
+    // Don't expect a collision.
+    auto *Current = Temp.get();
+    auto *Distinct = MDNode::replaceWithDistinct(std::move(Temp));
+    EXPECT_TRUE(Distinct->isDistinct());
+    EXPECT_TRUE(Distinct->isResolved());
+    EXPECT_EQ(Current, Distinct);
+  }
+  {
+    auto Unresolved = MDTuple::getTemporary(Context, None);
+    Metadata *Ops[] = {Unresolved.get()};
+    auto Temp = MDTuple::getTemporary(Context, Ops);
+    EXPECT_TRUE(Temp->isTemporary());
+
+    // Don't expect a collision.
+    auto *Current = Temp.get();
+    auto *Distinct = MDNode::replaceWithDistinct(std::move(Temp));
+    EXPECT_TRUE(Distinct->isDistinct());
+    EXPECT_TRUE(Distinct->isResolved());
+    EXPECT_EQ(Current, Distinct);
+
+    // Cleanup; required for teardown.
+    Unresolved->replaceAllUsesWith(nullptr);
+  }
+}
+
 typedef MetadataTest MDLocationTest;
 
 TEST_F(MDLocationTest, Overflow) {