[Assignment Tracking][5/*] Add core infrastructure for instruction reference

[Orlando]

Summary

The Assignment Tracking debug-info feature is outlined in this RFC. This first series of patches adds documentation, the changes necessary to start emitting and using the new metadata, and updates clang with an option to enable the feature. Working with the new metadata in the middle and back end will come later. There are still a few rough edges but I'm putting these patches up now hoping to get feedback on the design and implementation from the upstream community.

---

Overview

It's possible to find intrinsics linked to an instruction by looking at the MetadataAsValue uses of the attached DIAssignID. That covers instruction -> intrinsic(s) lookup. Add a global DIAssignID -> instruction(s) map which gives us the ability to perform intrinsic -> instruction(s) lookup. Add plumbing to keep the map up to date through optimisations and add utility functions including two that perform those lookups. Finally, add a unittest.

Details / patch tour

In llvm/lib/IR/LLVMContextImpl.h add AssignmentIDToInstrs which maps DIAssignID * attachments to Instruction *s. Because the DIAssignID * is the key we can't use a TrackingMDNodeRef for it, and therefore cannot easily update the mapping when a temporary DIAssignID is replaced.

Temporary DIAssignID's are only used in IR parsing to deal with metadata forward references. Update llvm/lib/AsmParser/LLParser.cpp to avoid using temporary DIAssignID's for attachments.

In llvm/lib/IR/Metadata.cpp add Instruction::updateDIAssignIDMapping which is called to remove or add an entry (or both) to AssignmentIDToInstrs. Call this from Instruction::setMetadata and add a call to setMetadata in Intruction's dtor that explicitly unsets the DIAssignID so that the mappging gets updated.

In llvm/lib/IR/DebugInfo.cpp and DebugInfo.h add utility functions:

getAssignmentInsts(const DbgAssignIntrinsic *DAI)
getAssignmentMarkers(const Instruction *Inst)
RAUW(DIAssignID *Old, DIAssignID *New)
deleteAll(Function *F)

These core utils are tested in llvm/unittests/IR/DebugInfoTest.cpp.

Notes / observations

This all needs to be looked at closely from a performance perspective once it's up and running. This mapping is obviously quite intrusive, but I'm not sure we can get around that.

[Orlando]

+ Add context to diff

[jmorse]

llvm/include/llvm/IR/DebugInfo.h

163

On the one hand, "llvm::at::" isn't the most descriptive of namespaces; but on the other hand, I don't believe llvm provides a stable C++ API, so if someone doesn't like it then it can just be changed. Good use of namespaces IMO.

163

NB: a high level comment saying "Utilities for enumerating storing instructions from an assignment ID", and then from line 175 an equivalent saying "For enumerating dbg.assigns..." would make this easier to read. (IMHO, YMMV).

llvm/include/llvm/IR/Instruction.h

532–533

(Three slashes)

llvm/include/llvm/IR/IntrinsicInst.h

288–291 ↗

(On Diff #456661)

IMO can be condensed to just "This method abstracts where the fragment is stored, for intrinsics with more than one expression" or something that doesn't name specific instructions. (The concern being that the comment will rot rapidly if it's too specific).

llvm/lib/IR/DebugInfo.cpp

1650–1653

Depending on how this API is to be used, would it be better to assert that an existing DIAssignID is always used? "Fail fast fail hard" works pretty well to detect deep errors.

1656

Spurious ;

1667–1669

Is there a risk that the SmallVector providing storage to the range being iterated on, is re-allocated / invalidates-iterators during the call to setMetadata, seeing how setMetadata -> updateDIAssignIDMapping -> erases / clears parts of the SmallVector?

If so, might be best replaced with something that modifies AssignmentIDToInstrs directly.

llvm/lib/IR/IntrinsicInst.cpp

86 ↗

(On Diff #456661)

This has a side-effect, but is called in an assertion -- that'll mean the side-effect only happens when LLVM is built with assertions, which is presumably undesirable?

94–95 ↗

(On Diff #456661)

DbgAssignAddrReplaced should be called? (It's missing parentheses)

179–184 ↗

(On Diff #456661)

AFAIUI, and I might be very wrong here, replaceUsesOfWith isn't declared virtual and so only code that casts an instruction to DbgAssignIntrinsic will take this code path. Also, the base replaceUsesOfWith already iterates over all operands and updates them, so this might not be necessary.

llvm/lib/IR/LLVMContextImpl.h

1490–1492

(Three slashes)

llvm/lib/IR/Metadata.cpp

1334

Could we make this DIAssignID * instead of auto *?

(I have this twitchy feeling because every variable in this function is auto; and all the rest of them are totally justifiable, they're container references and iterators. Except this!).

1342

Clearer assertion message please

1346

IMO: can / should be an assertion, right?

1349

llvm/lib/IR/Verifier.cpp

4438–4441

Possibly stupid question, but isn't the set iterated over here the same as in the users list above? Is it worth putting the AssertDI inside the loop above? (Feels like a massive nit pick over a tiny performance thing, feel free to ignore).

5670–5677

Feels better to use OpAddress etc rather than hard coded operand indexes.

llvm/unittests/IR/DebugInfoTest.cpp

405

std::next preferred I think (what if begin() returns a reference that gets mutated?)

[Orlando]

+ Address review comments
+ Fix verifier failure in unittest (wrong subprogram for variable)
+ Add test for verifier check added in this patch

llvm/include/llvm/IR/IntrinsicInst.h

288–291 ↗

(On Diff #456661)

That change leaked into this patch from the previous patch in the series (D132223) that added the dbg.assign boilerplate, sorry for the noise.

llvm/lib/IR/DebugInfo.cpp

1650–1653

This isn't a error state - you can get into the situation where you have a DIAssignID attachment on a function which is not used by any llvm.dbg.assign intrinsics. For instance, if a llvm.dbg.assign has been deleted due to living in a dead block -- whether or not _that_ is desirable needs to be looked at on a case-by-case basis IMO, so an assert would be too broad.

1667–1669

Good point, and I've added some "Iterators invalidated by ..." comments to the getAssignment... functions.

llvm/lib/IR/IntrinsicInst.cpp

86 ↗

(On Diff #456661)

This stuff is from the previous patch in the stack, sorry!

llvm/lib/IR/Verifier.cpp

4438–4441

Not a stupid question, they are the same. I don't mind either way so I've changed it to your preference. I've also updated the verifier tests added in recent updates to earlier patches to check this codepath.

5670–5677

This also comes from the previous patch in the stack, but I will apply this suggestion to it over there.

[jmorse]

LGTM

llvm/lib/IR/DebugInfo.cpp

1650–1653

SGTM

[Orlando]

Thanks!

llvm/lib/IR/DebugInfo.cpp

1650–1653

I can't seem to edit inline comments. For posterity: in the first sentence, when I said "on a function" I meant "on an instruction".

[chrisjackson]

llvm/include/llvm/IR/DebugInfo.h

163

I'm bikeshedding but i think the namespace name is too terse. Also possibly not nice to read because it's a word and an acronym

llvm/lib/AsmParser/LLParser.cpp

803

Possibly need to assert on ToReplace?

805

nit, Is assert text really a question? Can you just state instead?

[Orlando]

Hi @chrisjackson, thank you for taking a look at these patches! :-)

llvm/include/llvm/IR/DebugInfo.h

163

Any suggestions or thoughts on alternatives?
astr (or does the 'ast' or possible 'str' stand out too much?)
astra (too fun?)
asst (not a fan)
atra ... etc.

llvm/lib/AsmParser/LLParser.cpp

803

Original loop didn't see fit to do so, so I think we're okay not to too?

805

Sure, will do before landing if there are no major changes to make.

[aprantl]

This patch added a cyclic dependency that breaks the module build. Could you please fix/revert it?

https://green.lab.llvm.org/green/view/LLDB/job/lldb-cmake/48197/consoleFull#-69937453049ba4694-19c4-4d7e-bec5-911270d8a58c

In file included from <module-includes>:1:
/Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/include/llvm/IR/Argument.h:18:10: fatal error: cyclic dependency in module 'LLVM_IR': LLVM_IR -> LLVM_intrinsic_gen -> LLVM_IR
#include "llvm/IR/Value.h"
         ^
While building module 'LLVM_MC' imported from /Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/lib/MC/MCAsmInfoCOFF.cpp:14:
While building module 'LLVM_IR' imported from /Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/include/llvm/MC/MCPseudoProbe.h:57:
In file included from <module-includes>:12:
/Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/include/llvm/IR/DebugInfo.h:24:10: fatal error: could not build module 'LLVM_intrinsic_gen'
#include "llvm/IR/IntrinsicInst.h"
 ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
While building module 'LLVM_MC' imported from /Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/lib/MC/MCAsmInfoCOFF.cpp:14:
In file included from <module-includes>:15:
In file included from /Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/include/llvm/MC/MCContext.h:23:
/Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/include/llvm/MC/MCPseudoProbe.h:57:10: fatal error: could not build module 'LLVM_IR'
#include "llvm/IR/PseudoProbe.h"
 ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~
/Users/buildslave/jenkins/workspace/lldb-cmake/llvm-project/llvm/lib/MC/MCAsmInfoCOFF.cpp:14:10: fatal error: could not build module 'LLVM_MC'
#include "llvm/MC/MCAsmInfoCOFF.h"
 ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
4 errors generated.

[rastogishubham]

This patch also breaks an LLDB test on Apple Silicon:

https://ci.swift.org/view/LLDB/job/llvm-org-lldb-release-debuginfo/658/consoleFull

******************** TEST 'lldb-shell :: ScriptInterpreter/Python/Crashlog/app_specific_backtrace_crashlog.test' FAILED ********************
Script:
--
: 'RUN: at line 3';   mkdir -p /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/tools/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/Output/app_specific_backtrace_crashlog.test.tmp.dir
: 'RUN: at line 4';   /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/bin/yaml2obj /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/llvm-project/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/Inputs/application_specific_info/asi.yaml > /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/tools/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/Output/app_specific_backtrace_crashlog.test.tmp.dir/asi
: 'RUN: at line 5';   /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/bin/lldb --no-lldbinit -S /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/tools/lldb/test/Shell/lit-lldb-init-quiet -o 'command script import lldb.macosx.crashlog'  -o 'crashlog -a -i -t /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/tools/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/Output/app_specific_backtrace_crashlog.test.tmp.dir/asi /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/llvm-project/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/Inputs/application_specific_info/asi.ips'  -o "thread list" -o "bt all" 2>&1 | /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/bin/FileCheck /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/llvm-project/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/app_specific_backtrace_crashlog.test
--
Exit Code: 1

Command Output (stderr):
--
/Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/llvm-project/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/app_specific_backtrace_crashlog.test:47:15: error: CHECK-NEXT: expected string not found in input
# CHECK-NEXT: frame #6: 0x00000001a05d3e4f dyld`start{{.*}}
              ^
<stdin>:42:45: note: scanning from here
 frame #5: 0x00000001047e3ecf asi`main + 127
                                            ^
<stdin>:43:2: note: possible intended match here
 frame #6: 0x00000001a05d3e4f
 ^

Input file: <stdin>
Check file: /Users/ec2-user/jenkins/workspace/llvm-org-lldb-release-debuginfo/llvm-project/lldb/test/Shell/ScriptInterpreter/Python/Crashlog/app_specific_backtrace_crashlog.test

-dump-input=help explains the following input dump.

Input was:
<<<<<<
           .
           .
           .
          37:  frame #0: 0x00000001a0a58418 
          38:  frame #1: 0x00000001a05a2ea7 
          39:  frame #2: 0x00000001a0b3dcc3 
          40:  frame #3: 0x00000001a0b46af3 
          41:  frame #4: 0x00000001a09a12a3 
          42:  frame #5: 0x00000001047e3ecf asi`main + 127 
next:47'0                                                 X error: no match found
          43:  frame #6: 0x00000001a05d3e4f 
next:47'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
next:47'1      ?                             possible intended match
          44: (lldb) thread list 
next:47'0     ~~~~~~~~~~~~~~~~~~~
          45: Process 96535 stopped 
next:47'0     ~~~~~~~~~~~~~~~~~~~~~~
          46: * thread #1: tid = 0x1af8f3, 0x00000001a08c7224, queue = 'com.apple.main-thread', stop reason = EXC_CRASH (code=0, subcode=0x0) 
next:47'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
          47: (lldb) bt all 
next:47'0     ~~~~~~~~~~~~~~
          48: * thread #1, queue = 'com.apple.main-thread', stop reason = EXC_CRASH (code=0, subcode=0x0) 
next:47'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           .
           .
           .
>>>>>>

--

********************

Could you please revert the patch and fix the issues?

[rastogishubham]

I am reverting this patch because of the bot failures Adrian mentioned above. I am sorry for any inconvenience!

[Orlando]

In D132224#3913559, @rastogishubham wrote:

I am reverting this patch because of the bot failures Adrian mentioned above. I am sorry for any inconvenience!

No worries, this was after office hours for me. Thanks for the revert(s).

[Orlando]

Re-landed in 26382a4412d29e1c31fc3cda5071d5d60832b69c in which I've updated llvm/include/llvm/module.modulemap. I'm not familiar with how module maps work but this looks like it has fixed it.

I've folded D133576 into this commit (as was the original plan).

[Orlando]

This patch added a cyclic dependency that breaks the module build. Could you please fix/revert it?

That should be fixed now.

In D132224#3913190, @rastogishubham wrote:

This patch also breaks an LLDB test on Apple Silicon:
...
Could you please revert the patch and fix the issues?

I don't think my patch was responsible for this test failure though given that a run with the patch reverted (https://ci.swift.org/view/LLDB/job/llvm-org-lldb-release-debuginfo/673/) still has that failure AFAICT.

llvm/include/llvm/AsmParser/LLParser.h

private:
     // Module being parsed, null if we are only parsing summary index.
     Module *M;
     // Summary index being parsed, null if we are only parsing Module.
     ModuleSummaryIndex *Index;
     SlotMapping *Slots;
 
     SmallVector<Instruction*, 64> InstsWithTBAATag;
 
+    /// DIAssignID metadata does not support temporary RAUW so we cannot use
+    /// the normal metadata forward reference resolution method. Instead,
+    /// non-temporary DIAssignID are attached to instructions (recorded here)
+    /// then replaced later.
+    DenseMap<MDNode *, SmallVector<Instruction *, 2>> TempDIAssignIDAttachments;
+
     // Type resolution handling data structures.  The location is set when we
     // have processed a use of the type but not a definition yet.
     StringMap<std::pair<Type*, LocTy> > NamedTypes;
     std::map<unsigned, std::pair<Type*, LocTy> > NumberedTypes;
 
     std::map<unsigned, TrackingMDNodeRef> NumberedMetadata;
     std::map<unsigned, std::pair<TempMDTuple, LocTy>> ForwardRefMDNodes;
 

llvm/include/llvm/IR/DebugInfo.h

 #ifndef LLVM_IR_DEBUGINFO_H
 #define LLVM_IR_DEBUGINFO_H
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/iterator_range.h"
-#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/IntrinsicInst.h"
 
 namespace llvm {
 
 class DbgDeclareInst;
 class DbgValueInst;
 class DbgVariableIntrinsic;
 class Instruction;
 class Module;
private:
   SmallVector<DICompileUnit *, 8> CUs;
   SmallVector<DISubprogram *, 8> SPs;
   SmallVector<DIGlobalVariableExpression *, 8> GVs;
   SmallVector<DIType *, 8> TYs;
   SmallVector<DIScope *, 8> Scopes;
   SmallPtrSet<const MDNode *, 32> NodesSeen;
 };
 
+/// Assignment Tracking (at).
+namespace at {
+//
+// Utilities for enumerating storing instructions from an assignment ID.
+//
+/// A range of instructions.
+using AssignmentInstRange =
+    iterator_range<SmallVectorImpl<Instruction *>::iterator>;
+/// Return a range of instructions (typically just one) that have \p ID
+/// as an attachment.
+/// Iterators invalidated by adding or removing DIAssignID metadata to/from any
+/// instruction (including by deleting or cloning instructions).
+AssignmentInstRange getAssignmentInsts(DIAssignID *ID);
+/// Return a range of instructions (typically just one) that perform the
+/// assignment that \p DAI encodes.
+/// Iterators invalidated by adding or removing DIAssignID metadata to/from any
+/// instruction (including by deleting or cloning instructions).
+inline AssignmentInstRange getAssignmentInsts(const DbgAssignIntrinsic *DAI) {
+  return getAssignmentInsts(cast<DIAssignID>(DAI->getAssignId()));
+}
+
+//
+// Utilities for enumerating llvm.dbg.assign intrinsic from an assignment ID.
+//
+/// High level: this is an iterator for llvm.dbg.assign intrinsics.
+/// Implementation details: this is a wrapper around Value's User iterator that
+/// dereferences to a DbgAssignIntrinsic ptr rather than a User ptr.
+class DbgAssignIt
+    : public iterator_adaptor_base<DbgAssignIt, Value::user_iterator,
+                                   typename std::iterator_traits<
+                                       Value::user_iterator>::iterator_category,
+                                   DbgAssignIntrinsic *, std::ptrdiff_t,
+                                   DbgAssignIntrinsic **,
+                                   DbgAssignIntrinsic *&> {
+public:
+  DbgAssignIt(Value::user_iterator It) : iterator_adaptor_base(It) {}
+  DbgAssignIntrinsic *operator*() const { return cast<DbgAssignIntrinsic>(*I); }
+};
+/// A range of llvm.dbg.assign intrinsics.
+using AssignmentMarkerRange = iterator_range<DbgAssignIt>;
+/// Return a range of dbg.assign intrinsics which use \ID as an operand.
+/// Iterators invalidated by deleting an intrinsic contained in this range.
+AssignmentMarkerRange getAssignmentMarkers(DIAssignID *ID);
+/// Return a range of dbg.assign intrinsics for which \p Inst performs the
+/// assignment they encode.
+/// Iterators invalidated by deleting an intrinsic contained in this range.
+inline AssignmentMarkerRange getAssignmentMarkers(const Instruction *Inst) {
+  if (auto *ID = Inst->getMetadata(LLVMContext::MD_DIAssignID))
+    return getAssignmentMarkers(cast<DIAssignID>(ID));
+  else
+    return make_range(Value::user_iterator(), Value::user_iterator());
+}
+
+/// Replace all uses (and attachments) of \p Old with \p New.
+void RAUW(DIAssignID *Old, DIAssignID *New);
+
+/// Remove all Assignment Tracking related intrinsics and metadata from \p F.
+void deleteAll(Function *F);
+
+} // end namespace at
+
 /// Return true if assignment tracking is enabled.
 bool getEnableAssignmentTracking();
 } // end namespace llvm
 
 #endif // LLVM_IR_DEBUGINFO_H

llvm/include/llvm/IR/Instruction.h

 
 private:
   // These are all implemented in Metadata.cpp.
   MDNode *getMetadataImpl(unsigned KindID) const;
   MDNode *getMetadataImpl(StringRef Kind) const;
   void
   getAllMetadataImpl(SmallVectorImpl<std::pair<unsigned, MDNode *>> &) const;
 
+  /// Update the LLVMContext ID-to-Instruction(s) mapping. If \p ID is nullptr
+  /// then clear the mapping for this instruction.
+  void updateDIAssignIDMapping(DIAssignID *ID);
+
 public:
   //===--------------------------------------------------------------------===//
   // Predicates and helper methods.
   //===--------------------------------------------------------------------===//
 
   /// Return true if the instruction is associative:
   ///
   ///   Associative operators satisfy:  x op (y op z) === (x op y) op z

llvm/lib/AsmParser/LLParser.cpp

if (parseSpecializedMDNode(Init, IsDistinct))
       return true;
   } else if (parseToken(lltok::exclaim, "Expected '!' here") ||
              parseMDTuple(Init, IsDistinct))
     return true;
 
   // See if this was forward referenced, if so, handle it.
   auto FI = ForwardRefMDNodes.find(MetadataID);
   if (FI != ForwardRefMDNodes.end()) {
-    FI->second.first->replaceAllUsesWith(Init);
+    auto *ToReplace = FI->second.first.get();
+    // DIAssignID has its own special forward-reference "replacement" for
+    // attachments (the temporary attachments are never actually attached).
+    if (isa<DIAssignID>(Init)) {
+      for (auto *Inst : TempDIAssignIDAttachments[ToReplace]) {
+        assert(!Inst->getMetadata(LLVMContext::MD_DIAssignID) &&
+               "Inst already has DIAssignID?");
+        Inst->setMetadata(LLVMContext::MD_DIAssignID, Init);
+      }
+    }
+
+    ToReplace->replaceAllUsesWith(Init);
     ForwardRefMDNodes.erase(FI);
 
     assert(NumberedMetadata[MetadataID] == Init && "Tracking VH didn't work");
   } else {
     if (NumberedMetadata.count(MetadataID))
       return tokError("Metadata id is already used");
     NumberedMetadata[MetadataID].reset(Init);
   }
do {
     if (Lex.getKind() != lltok::MetadataVar)
       return tokError("expected metadata after comma");
 
     unsigned MDK;
     MDNode *N;
     if (parseMetadataAttachment(MDK, N))
       return true;
 
+    if (MDK == LLVMContext::MD_DIAssignID)
+      TempDIAssignIDAttachments[N].push_back(&Inst);
+    else
       Inst.setMetadata(MDK, N);
+
     if (MDK == LLVMContext::MD_tbaa)
       InstsWithTBAATag.push_back(&Inst);
 
     // If this is the end of the list, we're done.
   } while (EatIfPresent(lltok::comma));
   return false;
 }
 

llvm/lib/IR/DebugInfo.cpp

 //===- DebugInfo.cpp - Debug Information Helper Classes -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the helper classes used to build and interpret debug
 // information in LLVM IR form.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm-c/DebugInfo.h"
+#include "LLVMContextImpl.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"

 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Casting.h"
 #include <algorithm>
 #include <cassert>
 #include <utility>
 
 using namespace llvm;
+using namespace llvm::at;
 using namespace llvm::dwarf;
 
 static cl::opt<bool>
     ExperimentalAssignmentTracking("experimental-assignment-tracking",
                                    cl::init(false));
 bool llvm::getEnableAssignmentTracking() {
   return ExperimentalAssignmentTracking;
 }

 #define HANDLE_METADATA_LEAF(CLASS) \
   case Metadata::CLASS##Kind: \
     return (LLVMMetadataKind)LLVM##CLASS##MetadataKind;
 #include "llvm/IR/Metadata.def"
   default:
     return (LLVMMetadataKind)LLVMGenericDINodeMetadataKind;
   }
 }
+
+AssignmentInstRange at::getAssignmentInsts(DIAssignID *ID) {
+  assert(ID && "Expected non-null ID");
+  LLVMContext &Ctx = ID->getContext();
+  auto &Map = Ctx.pImpl->AssignmentIDToInstrs;
+
+  auto MapIt = Map.find(ID);
+  if (MapIt == Map.end())
+    return make_range(nullptr, nullptr);
+
+  return make_range(MapIt->second.begin(), MapIt->second.end());
+}
+
+AssignmentMarkerRange at::getAssignmentMarkers(DIAssignID *ID) {
+  assert(ID && "Expected non-null ID");
+  LLVMContext &Ctx = ID->getContext();
+
+  auto *IDAsValue = MetadataAsValue::getIfExists(Ctx, ID);
+
+  // The ID is only used wrapped in MetadataAsValue(ID), so lets check that
+  // one of those already exists first.
+  if (!IDAsValue)
+    return make_range(Value::user_iterator(), Value::user_iterator());
+
+  return make_range(IDAsValue->user_begin(), IDAsValue->user_end());
+}
+
+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
+  // getAssignmentInsts range iterators will be invalidated by adding and
+  // removing DIAssignID attachments.
+  SmallVector<Instruction *> InstVec(InstRange.begin(), InstRange.end());
+  for (auto *I : InstVec)
+    I->setMetadata(LLVMContext::MD_DIAssignID, New);
+}
+
+void at::deleteAll(Function *F) {
+  SmallVector<DbgAssignIntrinsic *, 12> ToDelete;
+  for (BasicBlock &BB : *F) {
+    for (Instruction &I : BB) {
+      if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(&I))
+        ToDelete.push_back(DAI);
+      else
+        I.setMetadata(LLVMContext::MD_DIAssignID, nullptr);
+    }
+  }
+  for (auto *DAI : ToDelete)
+    DAI->eraseFromParent();
+}

llvm/lib/IR/Instruction.cpp

Instruction::~Instruction() {
   //   uses to an empty ValueAsMetadata node. This makes extant dbg.value uses
   //   trivially dead (i.e. fair game for deletion in many passes), leading to
   //   stale dbg.values being in effect for too long.
   // - Call salvageDebugInfoOrMarkUndef. Not needed to make instruction removal
   //   correct. OTOH results in wasted work in some common cases (e.g. when all
   //   instructions in a BasicBlock are deleted).
   if (isUsedByMetadata())
     ValueAsMetadata::handleRAUW(this, UndefValue::get(getType()));
+
+  // Explicitly remove DIAssignID metadata to clear up ID -> Instruction(s)
+  // mapping in LLVMContext.
+  setMetadata(LLVMContext::MD_DIAssignID, nullptr);
 }
 
 
 void Instruction::setParent(BasicBlock *P) {
   Parent = P;
 }
 
 const Module *Instruction::getModule() const {

llvm/lib/IR/LLVMContextImpl.h

#include "llvm/IR/Metadata.def"
   ValueHandlesTy ValueHandles;
 
   /// CustomMDKindNames - Map to hold the metadata string to ID mapping.
   StringMap<unsigned> CustomMDKindNames;
 
   /// Collection of metadata used in this context.
   DenseMap<const Value *, MDAttachments> ValueMetadata;
 
+  /// Map DIAssignID -> Instructions with that attachment.
+  /// Managed by Instruction via Instruction::updateDIAssignIDMapping.
+  /// Query using the at:: functions defined in DebugInfo.h.
+  DenseMap<DIAssignID *, SmallVector<Instruction *, 1>> AssignmentIDToInstrs;
+
   /// Collection of per-GlobalObject sections used in this context.
   DenseMap<const GlobalObject *, StringRef> GlobalObjectSections;
 
   /// Collection of per-GlobalValue partitions used in this context.
   DenseMap<const GlobalValue *, StringRef> GlobalValuePartitions;
 
   /// DiscriminatorTable - This table maps file:line locations to an
   /// integer representing the next DWARF path discriminator to assign to

llvm/lib/IR/Metadata.cpp

void Instruction::dropUnknownNonDebugMetadata(ArrayRef<unsigned> KnownIDs) {
   });
 
   if (Info.empty()) {
     // Drop our entry at the store.
     clearMetadata();
   }
 }
 
+void Instruction::updateDIAssignIDMapping(DIAssignID *ID) {
+  auto &IDToInstrs = getContext().pImpl->AssignmentIDToInstrs;
+  if (const DIAssignID *CurrentID =
+          cast_or_null<DIAssignID>(getMetadata(LLVMContext::MD_DIAssignID))) {
+    // Nothing to do if the ID isn't changing.
+    if (ID == CurrentID)
+      return;
+
+    // Unmap this instruction from its current ID.
+    auto InstrsIt = IDToInstrs.find(CurrentID);
+    assert(InstrsIt != IDToInstrs.end() &&
+           "Expect existing attachment to be mapped");
+
+    auto &InstVec = InstrsIt->second;
+    auto *InstIt = std::find(InstVec.begin(), InstVec.end(), this);
+    assert(InstIt != InstVec.end() &&
+           "Expect instruction to be mapped to attachment");
+    // The vector contains a ptr to this. If this is the only element in the
+    // vector, remove the ID:vector entry, otherwise just remove the
+    // instruction from the vector.
+    if (InstVec.size() == 1)
+      IDToInstrs.erase(InstrsIt);
+    else
+      InstVec.erase(InstIt);
+  }
+
+  // Map this instruction to the new ID.
+  if (ID)
+    IDToInstrs[ID].push_back(this);
+}
+
 void Instruction::setMetadata(unsigned KindID, MDNode *Node) {
   if (!Node && !hasMetadata())
     return;
 
   // Handle 'dbg' as a special case since it is not stored in the hash table.
   if (KindID == LLVMContext::MD_dbg) {
     DbgLoc = DebugLoc(Node);
     return;
   }
 
+  // Update DIAssignID to Instruction(s) mapping.
+  if (KindID == LLVMContext::MD_DIAssignID) {
+    // The DIAssignID tracking infrastructure doesn't support RAUWing temporary
+    // nodes with DIAssignIDs. The cast_or_null below would also catch this, but
+    // having a dedicated assert helps make this obvious.
+    assert((!Node || !Node->isTemporary()) && "Temporary DIAssignIDs are invalid");
+    updateDIAssignIDMapping(cast_or_null<DIAssignID>(Node));
+  }
+
   Value::setMetadata(KindID, Node);
 }
 
 void Instruction::addAnnotationMetadata(StringRef Name) {
   MDBuilder MDB(getContext());
 
   auto *Existing = getMetadata(LLVMContext::MD_annotation);
   SmallVector<Metadata *, 4> Names;

llvm/lib/IR/Verifier.cpp

 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Comdat.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
void Verifier::visitDIAssignIDMetadata(Instruction &I, MDNode *MD) {
 
   // Iterate over the MetadataAsValue uses of the DIAssignID - these should
   // only be found as DbgAssignIntrinsic operands.
   if (auto *AsValue = MetadataAsValue::getIfExists(Context, MD)) {
     for (auto *User : AsValue->users()) {
       AssertDI(isa<DbgAssignIntrinsic>(User),
                "!DIAssignID should only be used by llvm.dbg.assign intrinsics",
                MD, User);
+      // All of the dbg.assign intrinsics should be in the same function as I.
+      if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(User))
+        AssertDI(DAI->getFunction() == I.getFunction(),
+                 "dbg.assign not in same function as inst", DAI, &I);
     }
   }
 }
 
 void Verifier::visitAnnotationMetadata(MDNode *Annotation) {
   Assert(isa<MDTuple>(Annotation), "annotation must be a tuple");
   Assert(Annotation->getNumOperands() >= 1,
          "annotation must have at least one operand");
   for (const MDOperand &Op : Annotation->operands())
     Assert(isa<MDString>(Op.get()), "operands must be strings");
AssertDI(isa<DILocalVariable>(DII.getRawVariable()),
          DII.getRawVariable());
   AssertDI(isa<DIExpression>(DII.getRawExpression()),
          "invalid llvm.dbg." + Kind + " intrinsic expression", &DII,
          DII.getRawExpression());
 
   if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(&DII)) {
     AssertDI(isa<DIAssignID>(DAI->getAssignId()),
              "invalid llvm.dbg.assign intrinsic ID (4th arg)", &DII,
              DAI->getAssignId());
     auto *Addr = cast<MetadataAsValue>(DII.getArgOperand(4))->getMetadata();
     AssertDI(isa<ValueAsMetadata>(Addr),
              "invalid llvm.dbg.assign intrinsic address (5th arg)", &DII, Addr);
     auto *AddrExpr = cast<MetadataAsValue>(DII.getArgOperand(5))->getMetadata();
     AssertDI(isa<DIExpression>(AddrExpr),
              "invalid llvm.dbg.assign intrinsic address expression (6th arg)",
              &DII, AddrExpr);
+
+    // All of the linked instructions should be in the same function as DII.
+    for (Instruction *I : at::getAssignmentInsts(DAI))
+      AssertDI(DAI->getFunction() == I->getFunction(),
+               "inst not in same function as dbg.assign", I, DAI);
   }
 
   // Ignore broken !dbg attachments; they're checked elsewhere.
   if (MDNode *N = DII.getDebugLoc().getAsMDNode())
     if (!isa<DILocation>(N))
       return;
 
   BasicBlock *BB = DII.getParent();

llvm/test/DebugInfo/Generic/assignment-tracking/parse-and-verify/verify.ll

 ; RUN: opt %s -S -verify -experimental-assignment-tracking 2>&1 \
 ; RUN: | FileCheck %s
 
 ;; Check that badly formed assignment tracking metadata is caught either
 ;; while parsing or by the verifier.
 ;;
 ;; Checks for this one are inline.
 
+define dso_local void @fun2() !dbg !15 {
+  ;; DIAssignID copied here from @fun() where it is used by intrinsics.
+  ; CHECK: dbg.assign not in same function as inst
+  %x = alloca i32, align 4, !DIAssignID !14
+  ret void
+}
+
 define dso_local void @fun() !dbg !7 {
 entry:
   %a = alloca i32, align 4, !DIAssignID !14
   ; Here something other than a dbg.assign intrinsic is using a DIAssignID.
   ; CHECK: !DIAssignID should only be used by llvm.dbg.assign intrinsics
   call void @llvm.dbg.value(metadata !14, metadata !10, metadata !DIExpression()), !dbg !13
 
   ; Each following dbg.assign has an argument of the incorrect type.

 !6 = !{!"clang version 14.0.0"}
 !7 = distinct !DISubprogram(name: "fun", scope: !1, file: !1, line: 1, type: !8, scopeLine: 1, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)
 !8 = !DISubroutineType(types: !9)
 !9 = !{null}
 !10 = !DILocalVariable(name: "local", scope: !7, file: !1, line: 2, type: !11)
 !11 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
 !13 = !DILocation(line: 1, column: 1, scope: !7)
 !14 = distinct !DIAssignID()
+!15 = distinct !DISubprogram(name: "fun2", scope: !1, file: !1, line: 1, type: !8, scopeLine: 1, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)

llvm/unittests/IR/DebugInfoTest.cpp

)");
   EXPECT_EQ(MD0Local->getValue(), Alloca);
   auto *MD1 = cast<MetadataAsValue>(Inst->getOperand(1))->getMetadata();
   EXPECT_EQ(MD1->getMetadataID(), Metadata::MetadataKind::DILocalVariableKind);
   auto *MD2 = cast<MetadataAsValue>(Inst->getOperand(2))->getMetadata();
   auto *MDExp = cast<DIExpression>(MD2);
   EXPECT_EQ(MDExp->getNumElements(), 0u);
 }
 
+TEST(AssignmentTrackingTest, Utils) {
+  // Test the assignment tracking utils defined in DebugInfo.h namespace at {}.
+  // This includes:
+  //     getAssignmentInsts
+  //     getAssignmentMarkers
+  //     RAUW
+  //     deleteAll
+  //
+  // The input IR includes two functions, fun1 and fun2. Both contain an alloca
+  // with a DIAssignID tag. fun1's alloca is linked to two llvm.dbg.assign
+  // intrinsics, one of which is for an inlined variable and appears before the
+  // alloca.
+
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"(
+    define dso_local void @fun1() !dbg !7 {
+    entry:
+      call void @llvm.dbg.assign(metadata i32 undef, metadata !10, metadata !DIExpression(), metadata !12, metadata i32 undef, metadata !DIExpression()), !dbg !13
+      %local = alloca i32, align 4, !DIAssignID !12
+      call void @llvm.dbg.assign(metadata i32 undef, metadata !16, metadata !DIExpression(), metadata !12, metadata i32 undef, metadata !DIExpression()), !dbg !15
+      ret void, !dbg !15
+    }
+
+    define dso_local void @fun2() !dbg !17 {
+    entry:
+      %local = alloca i32, align 4, !DIAssignID !20
+      call void @llvm.dbg.assign(metadata i32 undef, metadata !18, metadata !DIExpression(), metadata !20, metadata i32 undef, metadata !DIExpression()), !dbg !19
+      ret void, !dbg !19
+    }
+
+    declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata)
+
+    !llvm.dbg.cu = !{!0}
+    !llvm.module.flags = !{!3, !4, !5}
+    !llvm.ident = !{!6}
+
+    !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 14.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, splitDebugInlining: false, nameTableKind: None)
+    !1 = !DIFile(filename: "test.c", directory: "/")
+    !2 = !{}
+    !3 = !{i32 7, !"Dwarf Version", i32 4}
+    !4 = !{i32 2, !"Debug Info Version", i32 3}
+    !5 = !{i32 1, !"wchar_size", i32 4}
+    !6 = !{!"clang version 14.0.0"}
+    !7 = distinct !DISubprogram(name: "fun1", scope: !1, file: !1, line: 1, type: !8, scopeLine: 1, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)
+    !8 = !DISubroutineType(types: !9)
+    !9 = !{null}
+    !10 = !DILocalVariable(name: "local3", scope: !14, file: !1, line: 2, type: !11)
+    !11 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+    !12 = distinct !DIAssignID()
+    !13 = !DILocation(line: 5, column: 1, scope: !14, inlinedAt: !15)
+    !14 = distinct !DISubprogram(name: "inline", scope: !1, file: !1, line: 1, type: !8, scopeLine: 1, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)
+    !15 = !DILocation(line: 3, column: 1, scope: !7)
+    !16 = !DILocalVariable(name: "local1", scope: !7, file: !1, line: 2, type: !11)
+    !17 = distinct !DISubprogram(name: "fun2", scope: !1, file: !1, line: 1, type: !8, scopeLine: 1, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)
+    !18 = !DILocalVariable(name: "local2", scope: !17, file: !1, line: 2, type: !11)
+    !19 = !DILocation(line: 4, column: 1, scope: !17)
+    !20 = distinct !DIAssignID()
+    )");
+
+  // Check the test IR isn't malformed.
+  ASSERT_TRUE(M);
+
+  Function &Fun1 = *M->getFunction("fun1");
+  Instruction &Alloca = *Fun1.getEntryBlock().getFirstNonPHIOrDbg();
+
+  // 1. Check the Instruction <-> Intrinsic mappings work in fun1.
+  //
+  // Check there are two llvm.dbg.assign intrinsics linked to Alloca.
+  auto CheckFun1Mapping = [&Alloca]() {
+    auto Markers = at::getAssignmentMarkers(&Alloca);
+    EXPECT_TRUE(std::distance(Markers.begin(), Markers.end()) == 2);
+    // Check those two entries are distinct.
+    DbgAssignIntrinsic *First = *Markers.begin();
+    DbgAssignIntrinsic *Second = *std::next(Markers.begin());
+    EXPECT_NE(First, Second);
+
+    // Check that we can get back to Alloca from each llvm.dbg.assign.
+    for (auto *DAI : Markers) {
+      auto Insts = at::getAssignmentInsts(DAI);
+      // Check there is exactly one instruction linked to each intrinsic. Use
+      // ASSERT_TRUE because we're going to dereference the begin iterator.
+      ASSERT_TRUE(std::distance(Insts.begin(), Insts.end()) == 1);
+      EXPECT_FALSE(Insts.empty());
+      // Check the linked instruction is Alloca.
+      Instruction *LinkedInst = *Insts.begin();
+      EXPECT_EQ(LinkedInst, &Alloca);
+    }
+  };
+  CheckFun1Mapping();
+
+  // 2. Check DIAssignID RAUW replaces attachments and uses.
+  //
+  DIAssignID *Old =
+      cast_or_null<DIAssignID>(Alloca.getMetadata(LLVMContext::MD_DIAssignID));
+  DIAssignID *New = DIAssignID::getDistinct(C);
+  ASSERT_TRUE(Old && New && New != Old);
+  at::RAUW(Old, New);
+  // Check fun1's alloca and intrinsics have been updated and the mapping still
+  // works.
+  EXPECT_EQ(New, cast_or_null<DIAssignID>(
+                     Alloca.getMetadata(LLVMContext::MD_DIAssignID)));
+  CheckFun1Mapping();
+
+  // Check that fun2's alloca and intrinsic have not not been updated.
+  Instruction &Fun2Alloca =
+      *M->getFunction("fun2")->getEntryBlock().getFirstNonPHIOrDbg();
+  DIAssignID *Fun2ID = cast_or_null<DIAssignID>(
+      Fun2Alloca.getMetadata(LLVMContext::MD_DIAssignID));
+  EXPECT_NE(New, Fun2ID);
+  auto Fun2Markers = at::getAssignmentMarkers(&Fun2Alloca);
+  ASSERT_TRUE(std::distance(Fun2Markers.begin(), Fun2Markers.end()) == 1);
+  auto Fun2Insts = at::getAssignmentInsts(*Fun2Markers.begin());
+  ASSERT_TRUE(std::distance(Fun2Insts.begin(), Fun2Insts.end()) == 1);
+  EXPECT_EQ(*Fun2Insts.begin(), &Fun2Alloca);
+
+  // 3. Check that deleting works and applies only to the target function.
+  at::deleteAll(&Fun1);
+  // There should now only be the alloca and ret in fun1.
+  EXPECT_EQ(Fun1.begin()->size(), 2);
+  // fun2's alloca should have the same DIAssignID and remain linked to its
+  // llvm.dbg.assign.
+  EXPECT_EQ(Fun2ID, cast_or_null<DIAssignID>(
+                        Fun2Alloca.getMetadata(LLVMContext::MD_DIAssignID)));
+  EXPECT_FALSE(at::getAssignmentMarkers(&Fun2Alloca).empty());
+}
+
 } // end namespace