[NFC] pull retvec logic to MemoryTaggingSupport.

[fmayer]

Summary

we will also need this for aarch64 stack tagging.

[eugenis]

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

73

That's a very vague function name for the common utils header.

How about getUntagLocationForFunctionExit and return Instruction*?

[fmayer]

change function signature.

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

73

Good point. getUntagLocationForFunctionExit is a bit of a mouthfull though, what do you think about this?

[eugenis]

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

73

I'd expect untagLocationIfFunctionExit to actually emit the untagging code.

[fmayer]

naming

[eugenis]

LGTM

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

 //===- MemoryTaggingSupport.h - helpers for memory tagging implementations ===//
 // 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 declares common infrastructure for HWAddressSanitizer and
 // Aarch64StackTagging.
 //
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 #define LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 
 namespace llvm {
if (NumCoveredExits == ReachableRetVec.size()) {
     return false;
   }
   return true;
 }
 
 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
                         const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
                         const DominatorTree *DT, size_t MaxLifetimes);
+
+Instruction *getUntagLocationIfFunctionExit(Instruction &Inst);
+
 } // namespace llvm
 
 #endif

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

for (auto &Inst : BB) {
           if (II->getIntrinsicID() == Intrinsic::lifetime_start)
             AllocasToInstrument[AI].LifetimeStart.push_back(II);
           else
             AllocasToInstrument[AI].LifetimeEnd.push_back(II);
           continue;
         }
       }
 
-      if (isa<ReturnInst>(Inst)) {
-        if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall())
-          RetVec.push_back(CI);
-        else
-          RetVec.push_back(&Inst);
-      } else if (isa<ResumeInst, CleanupReturnInst>(Inst)) {
-        RetVec.push_back(&Inst);
-      }
+      Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
+      if (ExitUntag)
+        RetVec.push_back(ExitUntag);
 
       if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
         for (Value *V : DVI->location_ops()) {
           if (auto *Alloca = dyn_cast_or_null<AllocaInst>(V))
             if (!AllocaDbgMap.count(Alloca) ||
                 AllocaDbgMap[Alloca].back() != DVI)
               AllocaDbgMap[Alloca].push_back(DVI);
         }

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
   // An alloca that has exactly one start and end in every possible execution.
   // If it has multiple ends, they have to be unreachable from each other, so
   // at most one of them is actually used for each execution of the function.
   return LifetimeStart.size() == 1 &&
          (LifetimeEnd.size() == 1 ||
           (LifetimeEnd.size() > 0 &&
            !maybeReachableFromEachOther(LifetimeEnd, DT, MaxLifetimes)));
 }
+
+Instruction *getUntagLocationIfFunctionExit(Instruction &Inst) {
+  if (isa<ReturnInst>(Inst)) {
+    if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall())
+      return CI;
+    return &Inst;
+  }
+  if (isa<ResumeInst, CleanupReturnInst>(Inst)) {
+    return &Inst;
+  }
+  return nullptr;
+}
 } // namespace llvm