[Assignment Tracking][6/*] Add trackAssignments function

The Assignment Tracking debug-info feature is outlined in this RFC:

https://discourse.llvm.org/t/
rfc-assignment-tracking-a-better-way-of-specifying-variable-locations-in-ir

Add trackAssignments which adds assignment tracking metadata to a function for
a specified set of variables. The intended callers are the inliner and the
front end - those calls will be added in separate patches.

I've added a pass called declare-to-assign (AssignmentTrackingPass) that
converts dbg.declare intrinsics to dbg.assigns using trackAssignments so that
the function can be easily tested (see
llvm/test/DebugInfo/Generic/track-assignments.ll). The pass could also be used
by front ends to easily test out enabling assignment tracking.

Reviewed By: jmorse

Differential Revision: https://reviews.llvm.org/D132225

Author: OCHyams

llvm/include/llvm/IR/DIBuilder.h

@@ -37,6 +37,7 @@ namespace llvm {
   class LLVMContext;
   class Module;
   class Value;
+  class DbgAssignIntrinsic;
 
   class DIBuilder {
     Module &M;
@@ -47,6 +48,7 @@ namespace llvm {
     Function *ValueFn;       ///< llvm.dbg.value
     Function *LabelFn;       ///< llvm.dbg.label
     Function *AddrFn;        ///< llvm.dbg.addr
+    Function *AssignFn;      ///< llvm.dbg.assign
 
     SmallVector<TrackingMDNodeRef, 4> AllEnumTypes;
     /// Track the RetainTypes, since they can be updated later on.
@@ -917,6 +919,26 @@ namespace llvm {
                                DIExpression *Expr, const DILocation *DL,
                                BasicBlock *InsertAtEnd);
 
+    /// Insert a new llvm.dbg.assign intrinsic call.
+    /// \param LinkedInstr   Instruction with a DIAssignID to link with the new
+    ///                      intrinsic. The intrinsic will be inserted after
+    ///                      this instruction.
+    /// \param Val           The value component of this dbg.assign.
+    /// \param SrcVar        Variable's debug info descriptor.
+    /// \param ValExpr       A complex location expression to modify \p Val.
+    /// \param Addr          The address component (store destination).
+    /// \param AddrExpr      A complex location expression to modify \p Addr.
+    ///                      NOTE: \p ValExpr carries the FragInfo for the
+    ///                      variable.
+    /// \param DL            Debug info location, usually: (line: 0,
+    ///                      column: 0, scope: var-decl-scope). See
+    ///                      getDebugValueLoc.
+    DbgAssignIntrinsic *insertDbgAssign(Instruction *LinkedInstr, Value *Val,
+                                        DILocalVariable *SrcVar,
+                                        DIExpression *ValExpr, Value *Addr,
+                                        DIExpression *AddrExpr,
+                                        const DILocation *DL);
+
     /// Insert a new llvm.dbg.declare intrinsic call.
     /// \param Storage      llvm::Value of the variable
     /// \param VarInfo      Variable's debug info descriptor.

llvm/include/llvm/IR/DebugInfo.h

@@ -18,10 +18,13 @@
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/iterator_range.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
@@ -225,8 +228,75 @@ void RAUW(DIAssignID *Old, DIAssignID *New);
 /// Remove all Assignment Tracking related intrinsics and metadata from \p F.
 void deleteAll(Function *F);
 
+/// Helper struct for trackAssignments, below. We don't use the similar
+/// DebugVariable class because trackAssignments doesn't (yet?) understand
+/// partial variables (fragment info) as input and want to make that clear and
+/// explicit using types. In addition, eventually we will want to understand
+/// expressions that modify the base address too, which a DebugVariable doesn't
+/// capture.
+struct VarRecord {
+  DILocalVariable *Var;
+  DILocation *DL;
+
+  VarRecord(DbgVariableIntrinsic *DVI)
+      : Var(DVI->getVariable()), DL(getDebugValueLoc(DVI)) {}
+  VarRecord(DILocalVariable *Var, DILocation *DL) : Var(Var), DL(DL) {}
+  friend bool operator<(const VarRecord &LHS, const VarRecord &RHS) {
+    return std::tie(LHS.Var, LHS.DL) < std::tie(RHS.Var, RHS.DL);
+  }
+  friend bool operator==(const VarRecord &LHS, const VarRecord &RHS) {
+    return std::tie(LHS.Var, LHS.DL) == std::tie(RHS.Var, RHS.DL);
+  }
+};
+
+/// Map of backing storage to a set of variables that are stored to it.
+/// TODO: Backing storage shouldn't be limited to allocas only. Some local
+/// variables have their storage allocated by the calling function (addresses
+/// passed in with sret & byval parameters).
+using StorageToVarsMap = DenseMap<const AllocaInst *, SmallSet<VarRecord, 2>>;
+
+/// Track assignments to \p Vars between \p Start and \p End.
+
+void trackAssignments(Function::iterator Start, Function::iterator End,
+                      const StorageToVarsMap &Vars, const DataLayout &DL,
+                      bool DebugPrints = false);
+
+/// Describes properties of a store that has a static size and offset into a
+/// some base storage. Used by the getAssignmentInfo functions.
+struct AssignmentInfo {
+  AllocaInst const *Base;  ///< Base storage.
+  uint64_t OffsetInBits;   ///< Offset into Base.
+  uint64_t SizeInBits;     ///< Number of bits stored.
+  bool StoreToWholeAlloca; ///< SizeInBits equals the size of the base storage.
+
+  AssignmentInfo(const DataLayout &DL, AllocaInst const *Base,
+                 uint64_t OffsetInBits, uint64_t SizeInBits)
+      : Base(Base), OffsetInBits(OffsetInBits), SizeInBits(SizeInBits),
+        StoreToWholeAlloca(
+            OffsetInBits == 0 &&
+            SizeInBits == DL.getTypeSizeInBits(Base->getAllocatedType())) {}
+};
+
+Optional<AssignmentInfo> getAssignmentInfo(const DataLayout &DL,
+                                           const MemIntrinsic *I);
+Optional<AssignmentInfo> getAssignmentInfo(const DataLayout &DL,
+                                           const StoreInst *SI);
+Optional<AssignmentInfo> getAssignmentInfo(const DataLayout &DL,
+                                           const AllocaInst *AI);
+
 } // end namespace at
 
+/// Convert @llvm.dbg.declare intrinsics into sets of @llvm.dbg.assign
+/// intrinsics by treating stores to the dbg.declare'd address as assignments
+/// to the variable. Not all kinds of variables are supported yet; those will
+/// be left with their dbg.declare intrinsics.
+class AssignmentTrackingPass : public PassInfoMixin<AssignmentTrackingPass> {
+public:
+  void runOnFunction(Function &F);
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 /// Return true if assignment tracking is enabled.
 bool getEnableAssignmentTracking();
 } // end namespace llvm

llvm/lib/IR/DIBuilder.cpp

@@ -30,7 +30,7 @@ static cl::opt<bool>
 
 DIBuilder::DIBuilder(Module &m, bool AllowUnresolvedNodes, DICompileUnit *CU)
     : M(m), VMContext(M.getContext()), CUNode(CU), DeclareFn(nullptr),
-      ValueFn(nullptr), LabelFn(nullptr), AddrFn(nullptr),
+      ValueFn(nullptr), LabelFn(nullptr), AddrFn(nullptr), AssignFn(nullptr),
       AllowUnresolvedNodes(AllowUnresolvedNodes) {
   if (CUNode) {
     if (const auto &ETs = CUNode->getEnumTypes())
@@ -960,6 +960,36 @@ Instruction *DIBuilder::insertDeclare(Value *Storage, DILocalVariable *VarInfo,
   return insertDeclare(Storage, VarInfo, Expr, DL, InsertAtEnd, InsertBefore);
 }
 
+DbgAssignIntrinsic *
+DIBuilder::insertDbgAssign(Instruction *LinkedInstr, Value *Val,
+                           DILocalVariable *SrcVar, DIExpression *ValExpr,
+                           Value *Addr, DIExpression *AddrExpr,
+                           const DILocation *DL) {
+  LLVMContext &Ctx = LinkedInstr->getContext();
+  Module *M = LinkedInstr->getModule();
+  if (!AssignFn)
+    AssignFn = Intrinsic::getDeclaration(M, Intrinsic::dbg_assign);
+
+  auto *Link = LinkedInstr->getMetadata(LLVMContext::MD_DIAssignID);
+  assert(Link && "Linked instruction must have DIAssign metadata attached");
+
+  std::array<Value *, 6> Args = {
+      MetadataAsValue::get(Ctx, ValueAsMetadata::get(Val)),
+      MetadataAsValue::get(Ctx, SrcVar),
+      MetadataAsValue::get(Ctx, ValExpr),
+      MetadataAsValue::get(Ctx, Link),
+      MetadataAsValue::get(Ctx, ValueAsMetadata::get(Addr)),
+      MetadataAsValue::get(Ctx, AddrExpr),
+  };
+
+  IRBuilder<> B(Ctx);
+  B.SetCurrentDebugLocation(DL);
+
+  auto *DVI = cast<DbgAssignIntrinsic>(B.CreateCall(AssignFn, Args));
+  DVI->insertAfter(LinkedInstr);
+  return DVI;
+}
+
 Instruction *DIBuilder::insertLabel(DILabel *LabelInfo, const DILocation *DL,
                                     Instruction *InsertBefore) {
   return insertLabel(LabelInfo, DL,