[RemoveDIs][DebugInfo] Make DIAssignID always replaceable (#78300)

This patch is a necessary step to allowing the new non-intrinsic debug
info to replace llvm.dbg.assign intrinsics. DIAssignIDs must be able to
look up the debug assigns that refer to them, and this patch makes them
always be considered "replaceable", allowing us to track and replace uses
for non-temporary instances.

Author: SLTozer

llvm/include/llvm/IR/DebugInfoMetadata.h

@@ -323,6 +323,10 @@ class DIAssignID : public MDNode {
   // This node has no operands to replace.
   void replaceOperandWith(unsigned I, Metadata *New) = delete;
 
+  SmallVector<DPValue *> getAllDPValueUsers() {
+    return Context.getReplaceableUses()->getAllDPValueUsers();
+  }
+
   static DIAssignID *getDistinct(LLVMContext &Context) {
     return getImpl(Context, Distinct);
   }

llvm/include/llvm/IR/Metadata.h

@@ -1057,6 +1057,7 @@ struct TempMDNodeDeleter {
 class MDNode : public Metadata {
   friend class ReplaceableMetadataImpl;
   friend class LLVMContextImpl;
+  friend class DIAssignID;
 
   /// The header that is coallocated with an MDNode along with its "small"
   /// operands. It is located immediately before the main body of the node.
@@ -1239,7 +1240,8 @@ class MDNode : public Metadata {
   bool isDistinct() const { return Storage == Distinct; }
   bool isTemporary() const { return Storage == Temporary; }
 
-  bool isReplaceable() const { return isTemporary(); }
+  bool isReplaceable() const { return isTemporary() || isAlwaysReplaceable(); }
+  bool isAlwaysReplaceable() const { return getMetadataID() == DIAssignIDKind; }
 
   /// RAUW a temporary.
   ///

llvm/lib/IR/DebugInfo.cpp

@@ -1787,13 +1787,6 @@ void at::deleteAssignmentMarkers(const Instruction *Inst) {
 }
 
 void at::RAUW(DIAssignID *Old, DIAssignID *New) {
-  // Replace MetadataAsValue uses.
-  if (auto *OldIDAsValue =
-          MetadataAsValue::getIfExists(Old->getContext(), Old)) {
-    auto *NewIDAsValue = MetadataAsValue::get(Old->getContext(), New);
-    OldIDAsValue->replaceAllUsesWith(NewIDAsValue);
-  }
-
   // Replace attachments.
   AssignmentInstRange InstRange = getAssignmentInsts(Old);
   // Use intermediate storage for the instruction ptrs because the
@@ -1802,6 +1795,8 @@ void at::RAUW(DIAssignID *Old, DIAssignID *New) {
   SmallVector<Instruction *> InstVec(InstRange.begin(), InstRange.end());
   for (auto *I : InstVec)
     I->setMetadata(LLVMContext::MD_DIAssignID, New);
+
+  Old->replaceAllUsesWith(New);
 }
 
 void at::deleteAll(Function *F) {

llvm/lib/IR/Metadata.cpp

@@ -439,24 +439,30 @@ void ReplaceableMetadataImpl::resolveAllUses(bool ResolveUsers) {
 // commentry in DIArgList::handleChangedOperand for details. Hidden behind
 // conditional compilation to avoid a compile time regression.
 ReplaceableMetadataImpl *ReplaceableMetadataImpl::getOrCreate(Metadata &MD) {
-  if (auto *N = dyn_cast<MDNode>(&MD))
-    return N->isResolved() ? nullptr : N->Context.getOrCreateReplaceableUses();
+  if (auto *N = dyn_cast<MDNode>(&MD)) {
+    return !N->isResolved() || N->isAlwaysReplaceable()
+               ? N->Context.getOrCreateReplaceableUses()
+               : nullptr;
+  }
   if (auto ArgList = dyn_cast<DIArgList>(&MD))
     return ArgList;
   return dyn_cast<ValueAsMetadata>(&MD);
 }
 
 ReplaceableMetadataImpl *ReplaceableMetadataImpl::getIfExists(Metadata &MD) {
-  if (auto *N = dyn_cast<MDNode>(&MD))
-    return N->isResolved() ? nullptr : N->Context.getReplaceableUses();
+  if (auto *N = dyn_cast<MDNode>(&MD)) {
+    return !N->isResolved() || N->isAlwaysReplaceable()
+               ? N->Context.getReplaceableUses()
+               : nullptr;
+  }
   if (auto ArgList = dyn_cast<DIArgList>(&MD))
     return ArgList;
   return dyn_cast<ValueAsMetadata>(&MD);
 }
 
 bool ReplaceableMetadataImpl::isReplaceable(const Metadata &MD) {
   if (auto *N = dyn_cast<MDNode>(&MD))
-    return !N->isResolved();
+    return !N->isResolved() || N->isAlwaysReplaceable();
   return isa<ValueAsMetadata>(&MD) || isa<DIArgList>(&MD);
 }
 

llvm/lib/IR/Verifier.cpp

@@ -173,6 +173,11 @@ struct VerifierSupport {
     }
   }
 
+  void Write(const DPValue *V) {
+    if (V)
+      V->print(*OS, MST, false);
+  }
+
   void Write(const Metadata *MD) {
     if (!MD)
       return;
@@ -4740,6 +4745,12 @@ void Verifier::visitDIAssignIDMetadata(Instruction &I, MDNode *MD) {
                 "dbg.assign not in same function as inst", DAI, &I);
     }
   }
+  for (DPValue *DPV : cast<DIAssignID>(MD)->getAllDPValueUsers()) {
+    CheckDI(DPV->isDbgAssign(),
+            "!DIAssignID should only be used by Assign DPVs.", MD, DPV);
+    CheckDI(DPV->getFunction() == I.getFunction(),
+            "DPVAssign not in same function as inst", DPV, &I);
+  }
 }
 
 void Verifier::visitCallStackMetadata(MDNode *MD) {