[MTE] [HWASan] Unify VariableLocation replacement.

[fmayer]

Summary

HWASan and MTE used to replace the debug variable locations differently,
which made the code hard to understand. Unify the two in one easier to
understand method.

llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h

 
 private:
   StackInfo Info;
   std::function<bool(const AllocaInst &)> IsInterestingAlloca;
 };
 
 uint64_t getAllocaSizeInBytes(const AllocaInst &AI);
 bool alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Align);
+void replaceAllVariableLocationOps(AllocaInfo &Info, Value *Old, Value *New);
 
 } // namespace memtag
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

bool AArch64StackTagging::runOnFunction(Function &Fn) {
 
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
   Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
 
   int NextTag = 0;
   for (auto &I : SInfo.AllocasToInstrument) {
-    const memtag::AllocaInfo &Info = I.second;
+    memtag::AllocaInfo &Info = I.second;
     AllocaInst *AI = Info.AI;
     int Tag = NextTag;
     NextTag = (NextTag + 1) % 16;
     // Replace alloca with tagp(alloca).
     IRBuilder<> IRB(Info.AI->getNextNode());
     Function *TagP = Intrinsic::getDeclaration(
         F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()});
     Instruction *TagPCall =
if (SInfo.UnrecognizedLifetimes.empty() && StandardLifetime &&
       // Remove all lifetime intrinsics for this alloca.
       for (auto &II : Info.LifetimeStart)
         II->eraseFromParent();
       for (auto &II : Info.LifetimeEnd)
         II->eraseFromParent();
     }
 
     // Fixup debug intrinsics to point to the new alloca.
-    for (auto DVI : Info.DbgVariableIntrinsics)
-      DVI->replaceVariableLocationOp(Info.OldAI, Info.AI);
+    memtag::replaceAllVariableLocationOps(Info, Info.OldAI, Info.AI);
   }
 
   // If we have instrumented at least one alloca, all unrecognized lifetime
   // instrinsics have to go.
   for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
 
   return true;
 }

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

if (ShouldDetectUseAfterScope && StandardLifetime) {
       if (!StandardLifetime) {
         for (auto &II : Info.LifetimeStart)
           II->eraseFromParent();
         for (auto &II : Info.LifetimeEnd)
           II->eraseFromParent();
       }
     }
     if (memtag::alignAndPadAlloca(Info, Align(Mapping.getObjectAlignment()))) {
-      for (auto DVI : Info.DbgVariableIntrinsics) {
-        SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
-                                           DVI->location_ops().end());
-        for (Value *V : LocationOps) {
-          if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
-            if (V == AI)
-              DVI->replaceVariableLocationOp(V, Info.AI);
-          }
-        }
-      }
+      memtag::replaceAllVariableLocationOps(Info, AI, Info.AI);
       AI->eraseFromParent();
     }
   }
   for (auto &I : SInfo.UnrecognizedLifetimes)
     I->eraseFromParent();
   return true;
 }
 

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

bool alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) {
   NewAI->copyMetadata(*Info.AI);
 
   auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI);
   Info.AI->replaceAllUsesWith(NewPtr);
   Info.AI = NewAI;
   return true;
 }
 
+void replaceAllVariableLocationOps(AllocaInfo &Info, Value *Old, Value *New) {
+  for (auto DVI : Info.DbgVariableIntrinsics) {
+    for (size_t I = 0; I < DVI->getNumVariableLocationOps(); ++I) {
+      if (DVI->getVariableLocationOp(I) == Old)
+        DVI->replaceVariableLocationOp(I, New);
+    }
+  }
+}
+
 } // namespace memtag
 } // namespace llvm