2 files changed, 105 insertions, 0 deletions

diff --git a/llvm/unittests/Linker/LinkModulesTest.cpp b/llvm/unittests/Linker/LinkModulesTest.cpp
index 45f1308d3bd..904ba58ce48 100644
--- a/llvm/unittests/Linker/LinkModulesTest.cpp
+++ b/llvm/unittests/Linker/LinkModulesTest.cpp
@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/DataLayout.h"
@@ -219,4 +220,77 @@ TEST_F(LinkModuleTest, CAPIFailure) {
   LLVMDisposeMessage(errout);
 }
 
+TEST_F(LinkModuleTest, MoveDistinctMDs) {
+  LLVMContext C;
+  SMDiagnostic Err;
+
+  const char *SrcStr = "define void @foo() !attach !0 {\n"
+                       "entry:\n"
+                       "  call void @llvm.md(metadata !1)\n"
+                       "  ret void, !attach !2\n"
+                       "}\n"
+                       "declare void @llvm.md(metadata)\n"
+                       "!named = !{!3, !4}\n"
+                       "!0 = distinct !{}\n"
+                       "!1 = distinct !{}\n"
+                       "!2 = distinct !{}\n"
+                       "!3 = distinct !{}\n"
+                       "!4 = !{!3}\n";
+
+  std::unique_ptr<Module> Src = parseAssemblyString(SrcStr, Err, C);
+  assert(Src);
+  ASSERT_TRUE(Src.get());
+
+  // Get the addresses of the Metadata before merging.
+  Function *F = &*Src->begin();
+  ASSERT_EQ("foo", F->getName());
+  BasicBlock *BB = &F->getEntryBlock();
+  auto *CI = cast<CallInst>(&BB->front());
+  auto *RI = cast<ReturnInst>(BB->getTerminator());
+  NamedMDNode *NMD = &*Src->named_metadata_begin();
+
+  MDNode *M0 = F->getMetadata("attach");
+  MDNode *M1 =
+      cast<MDNode>(cast<MetadataAsValue>(CI->getArgOperand(0))->getMetadata());
+  MDNode *M2 = RI->getMetadata("attach");
+  MDNode *M3 = NMD->getOperand(0);
+  MDNode *M4 = NMD->getOperand(1);
+
+  // Confirm a few things about the IR.
+  EXPECT_TRUE(M0->isDistinct());
+  EXPECT_TRUE(M1->isDistinct());
+  EXPECT_TRUE(M2->isDistinct());
+  EXPECT_TRUE(M3->isDistinct());
+  EXPECT_TRUE(M4->isUniqued());
+  EXPECT_EQ(M3, M4->getOperand(0));
+
+  // Link into destination module.
+  auto Dst = llvm::make_unique<Module>("Linked", C);
+  ASSERT_TRUE(Dst.get());
+  Linker::LinkModules(Dst.get(), Src.get(),
+                      [](const llvm::DiagnosticInfo &) {});
+
+  // Check that distinct metadata was moved, not cloned.  Even !4, the uniqued
+  // node, should effectively be moved, since its only operand hasn't changed.
+  F = &*Dst->begin();
+  BB = &F->getEntryBlock();
+  CI = cast<CallInst>(&BB->front());
+  RI = cast<ReturnInst>(BB->getTerminator());
+  NMD = &*Dst->named_metadata_begin();
+
+  EXPECT_EQ(M0, F->getMetadata("attach"));
+  EXPECT_EQ(M1, cast<MetadataAsValue>(CI->getArgOperand(0))->getMetadata());
+  EXPECT_EQ(M2, RI->getMetadata("attach"));
+  EXPECT_EQ(M3, NMD->getOperand(0));
+  EXPECT_EQ(M4, NMD->getOperand(1));
+
+  // Confirm a few things about the IR.  This shouldn't have changed.
+  EXPECT_TRUE(M0->isDistinct());
+  EXPECT_TRUE(M1->isDistinct());
+  EXPECT_TRUE(M2->isDistinct());
+  EXPECT_TRUE(M3->isDistinct());
+  EXPECT_TRUE(M4->isUniqued());
+  EXPECT_EQ(M3, M4->getOperand(0));
+}
+
 } // end anonymous namespace
diff --git a/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp b/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp
index 137a2607c84..9dbe4dbc56d 100644
--- a/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp
+++ b/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp
@@ -24,4 +24,35 @@ TEST(ValueMapperTest, MapMetadataUnresolved) {
   EXPECT_EQ(T.get(), MapMetadata(T.get(), VM, RF_NoModuleLevelChanges));
 }
 
+TEST(ValueMapperTest, MapMetadataDistinct) {
+  LLVMContext Context;
+  auto *D = MDTuple::getDistinct(Context, None);
+
+  {
+    // The node should be cloned.
+    ValueToValueMapTy VM;
+    EXPECT_NE(D, MapMetadata(D, VM, RF_None));
+  }
+  {
+    // The node should be moved.
+    ValueToValueMapTy VM;
+    EXPECT_EQ(D, MapMetadata(D, VM, RF_MoveDistinctMDs));
+  }
+}
+
+TEST(ValueMapperTest, MapMetadataDistinctOperands) {
+  LLVMContext Context;
+  Metadata *Old = MDTuple::getDistinct(Context, None);
+  auto *D = MDTuple::getDistinct(Context, Old);
+  ASSERT_EQ(Old, D->getOperand(0));
+
+  Metadata *New = MDTuple::getDistinct(Context, None);
+  ValueToValueMapTy VM;
+  VM.MD()[Old].reset(New);
+
+  // Make sure operands are updated.
+  EXPECT_EQ(D, MapMetadata(D, VM, RF_MoveDistinctMDs));
+  EXPECT_EQ(New, D->getOperand(0));
+}
+
 }