[musttail] Unify musttail call preceding return checking

[lxfind]

Summary

There is already an API in BasciBlock that checks and returns the musttail call if it precedes the return instruction.
Use it instead of manually checking in each place.

[rnk]

lgtm

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

 static uint64_t GetCtorAndDtorPriority(Triple &TargetTriple) {
   if (TargetTriple.isOSEmscripten()) {
     return kAsanEmscriptenCtorAndDtorPriority;
   } else {
     return kAsanCtorAndDtorPriority;
   }
 }
 
-// For a ret instruction followed by a musttail call, we cannot insert anything
-// in between. Instead we use the musttail call instruction as the insertion
-// point.
-static Instruction *adjustForMusttailCall(Instruction *I) {
-  ReturnInst *RI = dyn_cast<ReturnInst>(I);
-  if (!RI)
-    return I;
-  Instruction *Prev = RI->getPrevNode();
-  if (BitCastInst *BCI = dyn_cast_or_null<BitCastInst>(Prev))
-    Prev = BCI->getPrevNode();
-  if (CallInst *CI = dyn_cast_or_null<CallInst>(Prev))
-    if (CI->isMustTailCall())
-      return CI;
-  return RI;
-}
-
 namespace {
 
 /// Module analysis for getting various metadata about the module.
 class ASanGlobalsMetadataWrapperPass : public ModulePass {
 public:
   static char ID;
 
   ASanGlobalsMetadataWrapperPass() : ModulePass(ID) {
struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   // poisoned red zones around all of them.
   // Then unpoison everything back before the function returns.
   void processStaticAllocas();
   void processDynamicAllocas();
 
   void createDynamicAllocasInitStorage();
 
   // ----------------------- Visitors.
-  /// Collect all Ret instructions.
-  void visitReturnInst(ReturnInst &RI) { RetVec.push_back(&RI); }
+  /// Collect all Ret instructions, or the musttail call instruction if it
+  /// precedes the return instruction.
+  void visitReturnInst(ReturnInst &RI) {
+    if (CallInst *CI = RI.getParent()->getTerminatingMustTailCall())
+      RetVec.push_back(CI);
+    else
+      RetVec.push_back(&RI);
+  }
 
   /// Collect all Resume instructions.
   void visitResumeInst(ResumeInst &RI) { RetVec.push_back(&RI); }
 
   /// Collect all CatchReturnInst instructions.
   void visitCleanupReturnInst(CleanupReturnInst &CRI) { RetVec.push_back(&CRI); }
 
   void unpoisonDynamicAllocasBeforeInst(Instruction *InstBefore,
void unpoisonDynamicAllocasBeforeInst(Instruction *InstBefore,
     IRB.CreateCall(
         AsanAllocasUnpoisonFunc,
         {IRB.CreateLoad(IntptrTy, DynamicAllocaLayout), DynamicAreaPtr});
   }
 
   // Unpoison dynamic allocas redzones.
   void unpoisonDynamicAllocas() {
     for (Instruction *Ret : RetVec)
-      unpoisonDynamicAllocasBeforeInst(adjustForMusttailCall(Ret),
-                                       DynamicAllocaLayout);
+      unpoisonDynamicAllocasBeforeInst(Ret, DynamicAllocaLayout);
 
     for (Instruction *StackRestoreInst : StackRestoreVec)
       unpoisonDynamicAllocasBeforeInst(StackRestoreInst,
                                        StackRestoreInst->getOperand(0));
   }
 
   // Deploy and poison redzones around dynamic alloca call. To do this, we
   // should replace this call with another one with changed parameters and
if (!StaticAllocaPoisonCallVec.empty()) {
     }
   }
 
   SmallVector<uint8_t, 64> ShadowClean(ShadowAfterScope.size(), 0);
   SmallVector<uint8_t, 64> ShadowAfterReturn;
 
   // (Un)poison the stack before all ret instructions.
   for (Instruction *Ret : RetVec) {
-    Instruction *Adjusted = adjustForMusttailCall(Ret);
-    IRBuilder<> IRBRet(Adjusted);
+    IRBuilder<> IRBRet(Ret);
     // Mark the current frame as retired.
     IRBRet.CreateStore(ConstantInt::get(IntptrTy, kRetiredStackFrameMagic),
                        BasePlus0);
     if (DoStackMalloc) {
       assert(StackMallocIdx >= 0);
       // if FakeStack != 0  // LocalStackBase == FakeStack
       //     // In use-after-return mode, poison the whole stack frame.
       //     if StackMallocIdx <= 4
       //         // For small sizes inline the whole thing:
       //         memset(ShadowBase, kAsanStackAfterReturnMagic, ShadowSize);
       //         **SavedFlagPtr(FakeStack) = 0
       //     else
       //         __asan_stack_free_N(FakeStack, LocalStackSize)
       // else
       //     <This is not a fake stack; unpoison the redzones>
       Value *Cmp =
           IRBRet.CreateICmpNE(FakeStack, Constant::getNullValue(IntptrTy));
       Instruction *ThenTerm, *ElseTerm;
-      SplitBlockAndInsertIfThenElse(Cmp, Adjusted, &ThenTerm, &ElseTerm);
+      SplitBlockAndInsertIfThenElse(Cmp, Ret, &ThenTerm, &ElseTerm);
 
       IRBuilder<> IRBPoison(ThenTerm);
       if (StackMallocIdx <= 4) {
         int ClassSize = kMinStackMallocSize << StackMallocIdx;
         ShadowAfterReturn.resize(ClassSize / L.Granularity,
                                  kAsanStackUseAfterReturnMagic);
         copyToShadow(ShadowAfterReturn, ShadowAfterReturn, IRBPoison,
                      ShadowBase);

llvm/lib/Transforms/Utils/EntryExitInstrumenter.cpp

static bool runOnFunction(Function &F, bool PostInlining) {
 
   if (!ExitFunc.empty()) {
     for (BasicBlock &BB : F) {
       Instruction *T = BB.getTerminator();
       if (!isa<ReturnInst>(T))
         continue;
 
       // If T is preceded by a musttail call, that's the real terminator.
-      Instruction *Prev = T->getPrevNode();
-      if (BitCastInst *BCI = dyn_cast_or_null<BitCastInst>(Prev))
-        Prev = BCI->getPrevNode();
-      if (CallInst *CI = dyn_cast_or_null<CallInst>(Prev)) {
-        if (CI->isMustTailCall())
-          T = CI;
-      }
+      if (CallInst *CI = BB.getTerminatingMustTailCall())
 
       DebugLoc DL;
       if (DebugLoc TerminatorDL = T->getDebugLoc())
         DL = TerminatorDL;
       else if (auto SP = F.getSubprogram())
         DL = DebugLoc::get(0, 0, SP);
 
       insertCall(F, ExitFunc, T, DL);

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

while (StateBB != StateE) {
     BasicBlock *CurBB = &*StateBB++;
 
     // Branches and invokes do not escape, only unwind, resume, and return
     // do.
     Instruction *TI = CurBB->getTerminator();
     if (!isa<ReturnInst>(TI) && !isa<ResumeInst>(TI))
       continue;
 
-    // If the ret instruction is followed by a musttaill call,
-    // or a bitcast instruction and then a musttail call, we should return
-    // the musttail call as the insertion point to not break the musttail
-    // contract.
-    auto AdjustMustTailCall = [&](Instruction *I) -> Instruction * {
-      auto *RI = dyn_cast<ReturnInst>(I);
-      if (!RI || !RI->getPrevNode())
-        return I;
-      auto *CI = dyn_cast<CallInst>(RI->getPrevNode());
-      if (CI && CI->isMustTailCall())
-        return CI;
-      auto *BI = dyn_cast<BitCastInst>(RI->getPrevNode());
-      if (!BI || !BI->getPrevNode())
-        return I;
-      CI = dyn_cast<CallInst>(BI->getPrevNode());
-      if (CI && CI->isMustTailCall())
-        return CI;
-      return I;
-    };
-
-    Builder.SetInsertPoint(AdjustMustTailCall(TI));
+    if (CallInst *CI = CurBB->getTerminatingMustTailCall())
+      TI = CI;
+    Builder.SetInsertPoint(TI);
     return &Builder;
   }
 
   Done = true;
 
   if (!HandleExceptions)
     return nullptr;