IR: Add MDCompositeTypeBase::replace*()

Add `replaceElements()`, `replaceVTableHolder()`, and
`replaceTemplateParams()` to `MDCompositeTypeBase`.  Included an
assertion in `replaceElements()` to match the one in
`DICompositeType::replaceArrays()`.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@229744 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 64 insertions, 0 deletions

diff --git a/include/llvm/IR/DebugInfoMetadata.h b/include/llvm/IR/DebugInfoMetadata.h
index a8a39260401e..fd513d462247 100644
--- a/include/llvm/IR/DebugInfoMetadata.h
+++ b/include/llvm/IR/DebugInfoMetadata.h
@@ -576,6 +576,29 @@ public:
 
   MDString *getRawIdentifier() const { return getOperandAs<MDString>(7); }
 
+  /// \brief Replace operands.
+  ///
+  /// If this \a isUniqued() and not \a isResolved(), on a uniquing collision
+  /// this will be RAUW'ed and deleted.  Use a \a TrackingMDRef to keep track
+  /// of its movement if necessary.
+  /// @{
+  void replaceElements(MDTuple *Elements) {
+#ifndef NDEBUG
+    if (auto *Old = cast_or_null<MDTuple>(getElements()))
+      for (const auto &Op : Old->operands())
+        assert(std::find(Elements->op_begin(), Elements->op_end(), Op) &&
+               "Lost a member during member list replacement");
+#endif
+    replaceOperandWith(4, Elements);
+  }
+  void replaceVTableHolder(Metadata *VTableHolder) {
+    replaceOperandWith(5, VTableHolder);
+  }
+  void replaceTemplateParams(MDTuple *TemplateParams) {
+    replaceOperandWith(6, TemplateParams);
+  }
+  /// @}
+
   static bool classof(const Metadata *MD) {
     return MD->getMetadataID() == MDCompositeTypeKind ||
            MD->getMetadataID() == MDSubroutineTypeKind;
diff --git a/unittests/IR/MetadataTest.cpp b/unittests/IR/MetadataTest.cpp
index a5d64b116c49..fc2ed73d89c0 100644
--- a/unittests/IR/MetadataTest.cpp
+++ b/unittests/IR/MetadataTest.cpp
@@ -897,6 +897,47 @@ TEST_F(MDCompositeTypeTest, get) {
   EXPECT_EQ(N, MDNode::replaceWithUniqued(std::move(Temp)));
 }
 
+TEST_F(MDCompositeTypeTest, replaceOperands) {
+  unsigned Tag = dwarf::DW_TAG_structure_type;
+  StringRef Name = "some name";
+  Metadata *File = MDTuple::getDistinct(Context, None);
+  unsigned Line = 1;
+  Metadata *Scope = MDTuple::getDistinct(Context, None);
+  Metadata *BaseType = MDTuple::getDistinct(Context, None);
+  unsigned SizeInBits = 2;
+  unsigned AlignInBits = 3;
+  unsigned OffsetInBits = 4;
+  unsigned Flags = 5;
+  unsigned RuntimeLang = 6;
+  StringRef Identifier = "some id";
+
+  auto *N = MDCompositeType::get(Context, Tag, Name, File, Line, Scope,
+                                 BaseType, SizeInBits, AlignInBits,
+                                 OffsetInBits, Flags, nullptr, RuntimeLang,
+                                 nullptr, nullptr, Identifier);
+
+  auto *Elements = MDTuple::getDistinct(Context, None);
+  EXPECT_EQ(nullptr, N->getElements());
+  N->replaceElements(Elements);
+  EXPECT_EQ(Elements, N->getElements());
+  N->replaceElements(nullptr);
+  EXPECT_EQ(nullptr, N->getElements());
+
+  auto *VTableHolder = MDTuple::getDistinct(Context, None);
+  EXPECT_EQ(nullptr, N->getVTableHolder());
+  N->replaceVTableHolder(VTableHolder);
+  EXPECT_EQ(VTableHolder, N->getVTableHolder());
+  N->replaceVTableHolder(nullptr);
+  EXPECT_EQ(nullptr, N->getVTableHolder());
+
+  auto *TemplateParams = MDTuple::getDistinct(Context, None);
+  EXPECT_EQ(nullptr, N->getTemplateParams());
+  N->replaceTemplateParams(TemplateParams);
+  EXPECT_EQ(TemplateParams, N->getTemplateParams());
+  N->replaceTemplateParams(nullptr);
+  EXPECT_EQ(nullptr, N->getTemplateParams());
+}
+
 typedef MetadataTest MDSubroutineTypeTest;
 
 TEST_F(MDSubroutineTypeTest, get) {