[RemoveDI] Handle DPValues in SROA

In order to reduce duplication, the migrate-debug-info loop has been changed to
a generic lambda with some helper function overloads, which is called for
dbg.declares, dbg.assigns, and DPValues alike.

Author: OCHyams

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -4838,6 +4838,35 @@ AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
   return NewAI;
 }
 
+static void insertNewDbgInst(DIBuilder &DIB, DbgDeclareInst *Orig,
+                             AllocaInst *NewAddr, DIExpression *NewFragmentExpr,
+                             Instruction *BeforeInst) {
+  DIB.insertDeclare(NewAddr, Orig->getVariable(), NewFragmentExpr,
+                    Orig->getDebugLoc(), BeforeInst);
+}
+static void insertNewDbgInst(DIBuilder &DIB, DbgAssignIntrinsic *Orig,
+                             AllocaInst *NewAddr, DIExpression *NewFragmentExpr,
+                             Instruction *BeforeInst) {
+  (void)BeforeInst;
+  if (!NewAddr->hasMetadata(LLVMContext::MD_DIAssignID)) {
+    NewAddr->setMetadata(LLVMContext::MD_DIAssignID,
+                         DIAssignID::getDistinct(NewAddr->getContext()));
+  }
+  auto *NewAssign = DIB.insertDbgAssign(
+      NewAddr, Orig->getValue(), Orig->getVariable(), NewFragmentExpr, NewAddr,
+      Orig->getAddressExpression(), Orig->getDebugLoc());
+  LLVM_DEBUG(dbgs() << "Created new assign intrinsic: " << *NewAssign << "\n");
+}
+static void insertNewDbgInst(DIBuilder &DIB, DPValue *Orig, AllocaInst *NewAddr,
+                             DIExpression *NewFragmentExpr,
+                             Instruction *BeforeInst) {
+  (void)DIB;
+  DPValue *New = new DPValue(ValueAsMetadata::get(NewAddr), Orig->getVariable(),
+                             NewFragmentExpr, Orig->getDebugLoc(),
+                             DPValue::LocationType::Declare);
+  BeforeInst->getParent()->insertDPValueBefore(New, BeforeInst->getIterator());
+}
+
 /// Walks the slices of an alloca and form partitions based on them,
 /// rewriting each of their uses.
 bool SROA::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
@@ -4939,15 +4968,7 @@ bool SROA::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
 
   // Migrate debug information from the old alloca to the new alloca(s)
   // and the individual partitions.
-  TinyPtrVector<DbgVariableIntrinsic *> DbgVariables;
-  SmallVector<DbgDeclareInst *, 1> DbgDeclares;
-  findDbgDeclares(DbgDeclares, &AI);
-  for (auto *DbgDeclare : DbgDeclares)
-    DbgVariables.push_back(DbgDeclare);
-  for (auto *DbgAssign : at::getAssignmentMarkers(&AI))
-    DbgVariables.push_back(DbgAssign);
-
-  for (DbgVariableIntrinsic *DbgVariable : DbgVariables) {
+  auto MigrateOne = [&](auto *DbgVariable) {
     auto *Expr = DbgVariable->getExpression();
     DIBuilder DIB(*AI.getModule(), /*AllowUnresolved*/ false);
     uint64_t AllocaSize =
@@ -5001,37 +5022,33 @@ bool SROA::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
       // Remove any existing intrinsics on the new alloca describing
       // the variable fragment.
       SmallVector<DbgDeclareInst *, 1> FragDbgDeclares;
-      findDbgDeclares(FragDbgDeclares, Fragment.Alloca);
-      for (DbgDeclareInst *OldDII : FragDbgDeclares) {
-        auto SameVariableFragment = [](const DbgVariableIntrinsic *LHS,
-                                       const DbgVariableIntrinsic *RHS) {
+      SmallVector<DPValue *, 1> FragDPVs;
+      findDbgDeclares(FragDbgDeclares, Fragment.Alloca, &FragDPVs);
+      auto RemoveOne = [DbgVariable](auto *OldDII) {
+        auto SameVariableFragment = [](const auto *LHS, const auto *RHS) {
           return LHS->getVariable() == RHS->getVariable() &&
                  LHS->getDebugLoc()->getInlinedAt() ==
                      RHS->getDebugLoc()->getInlinedAt();
         };
         if (SameVariableFragment(OldDII, DbgVariable))
           OldDII->eraseFromParent();
-      }
+      };
+      for_each(FragDbgDeclares, RemoveOne);
+      for_each(FragDPVs, RemoveOne);
 
-      if (auto *DbgAssign = dyn_cast<DbgAssignIntrinsic>(DbgVariable)) {
-        if (!Fragment.Alloca->hasMetadata(LLVMContext::MD_DIAssignID)) {
-          Fragment.Alloca->setMetadata(
-              LLVMContext::MD_DIAssignID,
-              DIAssignID::getDistinct(AI.getContext()));
-        }
-        auto *NewAssign = DIB.insertDbgAssign(
-            Fragment.Alloca, DbgAssign->getValue(), DbgAssign->getVariable(),
-            FragmentExpr, Fragment.Alloca, DbgAssign->getAddressExpression(),
-            DbgAssign->getDebugLoc());
-        NewAssign->setDebugLoc(DbgAssign->getDebugLoc());
-        LLVM_DEBUG(dbgs() << "Created new assign intrinsic: " << *NewAssign
-                          << "\n");
-      } else {
-        DIB.insertDeclare(Fragment.Alloca, DbgVariable->getVariable(),
-                          FragmentExpr, DbgVariable->getDebugLoc(), &AI);
-      }
+      insertNewDbgInst(DIB, DbgVariable, Fragment.Alloca, FragmentExpr, &AI);
     }
-  }
+  };
+
+  // Migrate debug information from the old alloca to the new alloca(s)
+  // and the individual partitions.
+  SmallVector<DbgDeclareInst *, 1> DbgDeclares;
+  SmallVector<DPValue *, 1> DPValues;
+  findDbgDeclares(DbgDeclares, &AI, &DPValues);
+  for_each(DbgDeclares, MigrateOne);
+  for_each(DPValues, MigrateOne);
+  for_each(at::getAssignmentMarkers(&AI), MigrateOne);
+
   return Changed;
 }
 
@@ -5153,9 +5170,12 @@ bool SROA::deleteDeadInstructions(
     if (AllocaInst *AI = dyn_cast<AllocaInst>(I)) {
       DeletedAllocas.insert(AI);
       SmallVector<DbgDeclareInst *, 1> DbgDeclares;
-      findDbgDeclares(DbgDeclares, AI);
+      SmallVector<DPValue *, 1> DPValues;
+      findDbgDeclares(DbgDeclares, AI, &DPValues);
       for (DbgDeclareInst *OldDII : DbgDeclares)
         OldDII->eraseFromParent();
+      for (DPValue *OldDII : DPValues)
+        OldDII->eraseFromParent();
     }
 
     at::deleteAssignmentMarkers(I);

llvm/test/DebugInfo/ARM/sroa-complex.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes='sroa' -S -o - %s | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes='sroa' -S -o - %s | FileCheck %s
 target datalayout = "e-m:o-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
 target triple = "thumbv7-apple-unknown-macho"
 

llvm/test/DebugInfo/Generic/sroa-larger.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes='sroa' -S -o - %s | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes='sroa' -S -o - %s | FileCheck %s
 ; Generated from clang -c  -O2 -g -target x86_64-pc-windows-msvc
 ; struct A {
 ;   int _Myval2;

llvm/test/DebugInfo/Generic/sroa-samesize.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes='sroa' -S -o - %s | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes='sroa' -S -o - %s | FileCheck %s
 ; Generated from clang -c  -O2 -g -target x86_64-pc-windows-msvc
 ; struct A { double x1[]; };
 ; struct x2 {

llvm/test/DebugInfo/X86/sroa-after-inlining.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -passes='cgscc(function(sroa,instcombine),inline),function(instcombine,sroa),verify' -S -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes='cgscc(function(sroa,instcombine),inline),function(instcombine,sroa),verify' -S -o - | FileCheck %s
 ;
 ; This test checks that SROA pass processes debug info correctly if applied twice.
 ; Specifically, after SROA works first time, instcombine converts dbg.declare

llvm/test/DebugInfo/X86/sroasplit-1.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -passes='sroa,verify' -S -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes='sroa,verify' -S -o - | FileCheck %s
 ;
 ; Test that we can partial emit debug info for aggregates repeatedly
 ; split up by SROA.

llvm/test/DebugInfo/X86/sroasplit-2.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -passes='sroa,verify' -S -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes='sroa,verify' -S -o - | FileCheck %s
 ;
 ; Test that we can partial emit debug info for aggregates repeatedly
 ; split up by SROA.

llvm/test/DebugInfo/X86/sroasplit-3.ll

@@ -1,4 +1,6 @@
 ; RUN: opt %s -passes='sroa,verify' -S -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes='sroa,verify' -S -o - | FileCheck %s
+
 ; ModuleID = 'test.c'
 ; Test that SROA updates the debug info correctly if an alloca was rewritten but
 ; not partitioned into multiple allocas.

llvm/test/DebugInfo/X86/sroasplit-4.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes='sroa' < %s -S -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes='sroa' < %s -S -o - | FileCheck %s
 ;
 ; Test that recursively splitting an alloca updates the debug info correctly.
 ; CHECK: %[[T:.*]] = load i64, ptr @t, align 8

llvm/test/DebugInfo/X86/sroasplit-5.ll

@@ -1,4 +1,6 @@
 ; RUN: opt %s -passes='sroa,verify' -S -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes='sroa,verify' -S -o - | FileCheck %s
+
 ; From:
 ; struct prog_src_register {
 ;   unsigned : 4;

llvm/test/DebugInfo/X86/sroasplit-dbg-declare.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -S -passes='sroa' -o - %s | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -S -passes='sroa' -o - %s | FileCheck %s
 
 ; SROA should split the alloca in two new ones, each with its own dbg.declare.
 ; The original alloca and dbg.declare should be removed.
@@ -55,3 +56,4 @@ attributes #1 = { nounwind readnone speculatable }
 ; CHECK-NOT:  = alloca [9 x i32]
 ; CHECK-NOT:  call void @llvm.dbg.declare(metadata ptr
 
+; CHECK: declare void @llvm.dbg.declare(metadata, metadata, metadata)

llvm/test/DebugInfo/salvage-overflow.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -passes='sroa,early-cse' -S | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes='sroa,early-cse' -S | FileCheck %s
 ; CHECK: DIExpression(DW_OP_constu, 9223372036854775808, DW_OP_minus, DW_OP_stack_value)
 ; Created from the following C input (and then delta-reduced the IR):
 ;

llvm/test/Transforms/Util/dbg-user-of-aext.ll

@@ -2,6 +2,7 @@
 ; (here exposed through the SROA) pass refers to [s|z]exts of values (as
 ; opposed to the operand of a [s|z]ext).
 ; RUN: opt -S -passes='sroa' %s | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -S -passes='sroa' %s | FileCheck %s
 
 ; Built from:
 ; struct foo { bool b; long i; };