[Assignment Tracking] Fix replaceVariableLocationOp for dbg.assign with DIArgList

[Orlando]

Summary

The last time this function was updated DIArgLists were not supported for dbg.assigns. Without this patch it's possible to now dereference an invalid (the end) iterator. This occurs when the Address component gets replaced and there's a DIArgList for the Value component (the contained values are irrelevant).

The added unittest crashes without the code change in this patch.

[Orlando]

This fixes https://lab.llvm.org/buildbot/#/builders/123/builds/18283 when enabling assignment tracking.

[jmorse]

LGTM, an annoying edge case.

llvm/lib/IR/IntrinsicInst.cpp

auto ReplaceDbgAssignAddress = [this, OldValue, NewValue]() -> bool {
     return true;
   };
   bool DbgAssignAddrReplaced = ReplaceDbgAssignAddress();
   (void)DbgAssignAddrReplaced;
 
   assert(NewValue && "Values must be non-null");
   auto Locations = location_ops();
   auto OldIt = find(Locations, OldValue);
-  assert((OldIt != Locations.end() || DbgAssignAddrReplaced) &&
-         "OldValue must be a current location");
-  if (!hasArgList()) {
-    // Additional check necessary to avoid unconditionally replacing this
-    // operand when a dbg.assign address is replaced (DbgAssignAddrReplaced is
-    // true).
-    if (OldValue != getVariableLocationOp(0))
-      return;
+  if (OldIt == Locations.end()) {
+    assert(DbgAssignAddrReplaced &&
+           "OldValue must be dbg.assign addr if unused in DIArgList");
+  }
+
+  assert(OldIt != Locations.end() && "OldValue must be a current location");
+  if (!hasArgList()) {
     Value *NewOperand = isa<MetadataAsValue>(NewValue)
                             ? NewValue
                             : MetadataAsValue::get(
                                   getContext(), ValueAsMetadata::get(NewValue));
     return setArgOperand(0, NewOperand);
   }
   SmallVector<ValueAsMetadata *, 4> MDs;
   ValueAsMetadata *NewOperand = getAsMetadata(NewValue);

llvm/unittests/IR/DebugInfoTest.cpp

#define TEST_REPLACE(Old, New, ExpectedValue, ExpectedAddr)                    \
 
   // Replace address only.
   TEST_REPLACE(/*Old*/ P1, /*New*/ P2, /*Value*/ V1, /*Address*/ P2);
   // Replace value only.
   TEST_REPLACE(/*Old*/ V1, /*New*/ P2, /*Value*/ P2, /*Address*/ P2);
   // Replace both.
   TEST_REPLACE(/*Old*/ P2, /*New*/ P1, /*Value*/ P1, /*Address*/ P1);
 
+  // Replace address only, value uses a DIArgList.
+  // Value = {DIArgList(V1)}, Addr = P1.
+  DAI->setRawLocation(DIArgList::get(C, ValueAsMetadata::get(V1)));
+  DAI->setExpression(DIExpression::get(
+      C, {dwarf::DW_OP_LLVM_arg, 0, dwarf::DW_OP_stack_value}));
+  TEST_REPLACE(/*Old*/ P1, /*New*/ P2, /*Value*/ V1, /*Address*/ P2);
 #undef TEST_REPLACE
 }
 
 TEST(AssignmentTrackingTest, Utils) {
   // Test the assignment tracking utils defined in DebugInfo.h namespace at {}.
   // This includes:
   //     getAssignmentInsts
   //     getAssignmentMarkers