From 02dbed4b7f4ff70403496e3d9dacd441b1d49b9f Mon Sep 17 00:00:00 2001
From: Florian Mayer <fmayer@google.com>
Date: Thu, 3 Feb 2022 17:10:28 -0800
Subject: [PATCH] [mte] [nfc] do not keep Tag in AllocaInfo

this is to get the implementation closer to HWASan

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D118960
---
 .../Target/AArch64/AArch64StackTagging.cpp    | 64 ++++++++-----------
 1 file changed, 27 insertions(+), 37 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index b9452c7bb2b0..8b653e23be8f 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -296,7 +296,6 @@ class AArch64StackTagging : public FunctionPass {
     SmallVector<IntrinsicInst *, 2> LifetimeStart;
     SmallVector<IntrinsicInst *, 2> LifetimeEnd;
     SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;
-    int Tag; // -1 for non-tagged allocations
   };
 
   const bool MergeInit;
@@ -467,15 +466,13 @@ void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
 }
 
 Instruction *AArch64StackTagging::insertBaseTaggedPointer(
-    const MapVector<AllocaInst *, AllocaInfo> &Allocas,
+    const MapVector<AllocaInst *, AllocaInfo> &InterestingAllocas,
     const DominatorTree *DT) {
   BasicBlock *PrologueBB = nullptr;
   // Try sinking IRG as deep as possible to avoid hurting shrink wrap.
-  for (auto &I : Allocas) {
+  for (auto &I : InterestingAllocas) {
     const AllocaInfo &Info = I.second;
     AllocaInst *AI = Info.AI;
-    if (Info.Tag < 0)
-      continue;
     if (!PrologueBB) {
       PrologueBB = AI->getParent();
       continue;
@@ -539,7 +536,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   if (MergeInit)
     AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
-  MapVector<AllocaInst *, AllocaInfo> Allocas; // need stable iteration order
+  MapVector<AllocaInst *, AllocaInfo>
+      InterestingAllocas; // need stable iteration order
   SmallVector<Instruction *, 8> RetVec;
   SmallVector<Instruction *, 4> UnrecognizedLifetimes;
 
@@ -551,17 +549,20 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
       }
     }
     if (auto *AI = dyn_cast<AllocaInst>(&I)) {
-      Allocas[AI].AI = AI;
-      Allocas[AI].OldAI = AI;
+      if (isInterestingAlloca(*AI)) {
+        InterestingAllocas[AI].AI = AI;
+        InterestingAllocas[AI].OldAI = AI;
+      }
       continue;
     }
 
     if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I)) {
       for (Value *V : DVI->location_ops())
         if (auto *AI = dyn_cast_or_null<AllocaInst>(V))
-          if (Allocas[AI].DbgVariableIntrinsics.empty() ||
-              Allocas[AI].DbgVariableIntrinsics.back() != DVI)
-            Allocas[AI].DbgVariableIntrinsics.push_back(DVI);
+          if (isInterestingAlloca(*AI) &&
+              (InterestingAllocas[AI].DbgVariableIntrinsics.empty() ||
+               InterestingAllocas[AI].DbgVariableIntrinsics.back() != DVI))
+            InterestingAllocas[AI].DbgVariableIntrinsics.push_back(DVI);
       continue;
     }
 
@@ -573,10 +574,13 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
         UnrecognizedLifetimes.push_back(&I);
         continue;
       }
+      if (!isInterestingAlloca(*AI))
+        continue;
       if (II->getIntrinsicID() == Intrinsic::lifetime_start)
-        Allocas[AI].LifetimeStart.push_back(II);
+        InterestingAllocas[AI].LifetimeStart.push_back(II);
       else
-        Allocas[AI].LifetimeEnd.push_back(II);
+        InterestingAllocas[AI].LifetimeEnd.push_back(II);
+      continue;
     }
 
     Instruction *ExitUntag = getUntagLocationIfFunctionExit(I);
@@ -584,35 +588,21 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
       RetVec.push_back(ExitUntag);
   }
 
-  if (Allocas.empty())
+  if (InterestingAllocas.empty())
     return false;
 
-  int NextTag = 0;
-  int NumInterestingAllocas = 0;
-  for (auto &I : Allocas) {
+  for (auto &I : InterestingAllocas) {
     AllocaInfo &Info = I.second;
-    assert(Info.AI);
-
-    if (!isInterestingAlloca(*Info.AI)) {
-      Info.Tag = -1;
-      continue;
-    }
-
+    assert(Info.AI && isInterestingAlloca(*Info.AI));
     alignAndPadAlloca(Info);
-    NumInterestingAllocas++;
-    Info.Tag = NextTag;
-    NextTag = (NextTag + 1) % 16;
   }
 
-  if (NumInterestingAllocas == 0)
-    return false;
-
   std::unique_ptr<DominatorTree> DeleteDT;
   DominatorTree *DT = nullptr;
   if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
     DT = &P->getDomTree();
 
-  if (DT == nullptr && (NumInterestingAllocas > 1 ||
+  if (DT == nullptr && (InterestingAllocas.size() > 1 ||
                         !F->hasFnAttribute(Attribute::OptimizeNone))) {
     DeleteDT = std::make_unique<DominatorTree>(*F);
     DT = DeleteDT.get();
@@ -631,21 +621,21 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
-  Instruction *Base = insertBaseTaggedPointer(Allocas, DT);
+  Instruction *Base = insertBaseTaggedPointer(InterestingAllocas, DT);
 
-  for (auto &I : Allocas) {
+  int NextTag = 0;
+  for (auto &I : InterestingAllocas) {
     const AllocaInfo &Info = I.second;
     AllocaInst *AI = Info.AI;
-    if (Info.Tag < 0)
-      continue;
-
+    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 =
         IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base,
-                              ConstantInt::get(IRB.getInt64Ty(), Info.Tag)});
+                              ConstantInt::get(IRB.getInt64Ty(), Tag)});
     if (Info.AI->hasName())
       TagPCall->setName(Info.AI->getName() + ".tag");
     Info.AI->replaceAllUsesWith(TagPCall);