[dfsan] Refactor visitCallBase

[stephan.yichao.zhao]

Summary

To simplify the review of https://reviews.llvm.org/D95835.

[stephan.yichao.zhao]

tweak format

[morehouse]

Please fix all the lints while we're here.

[stephan.yichao.zhao]

fixed lint issues

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

public:
   void visitInsertElementInst(InsertElementInst &I);
   void visitShuffleVectorInst(ShuffleVectorInst &I);
   void visitExtractValueInst(ExtractValueInst &I);
   void visitInsertValueInst(InsertValueInst &I);
   void visitAllocaInst(AllocaInst &I);
   void visitSelectInst(SelectInst &I);
   void visitMemSetInst(MemSetInst &I);
   void visitMemTransferInst(MemTransferInst &I);
+
+private:
+  // Returns false when this is an invoke of a custom function.
+  bool visitWrappedCallBase(Function &F, CallBase &CB);
 };
 
 } // end anonymous namespace
 
 DataFlowSanitizer::DataFlowSanitizer(
     const std::vector<std::string> &ABIListFiles) {
   std::vector<std::string> AllABIListFiles(std::move(ABIListFiles));
   llvm::append_range(AllABIListFiles, ClABIListFiles);
case DataFlowSanitizer::IA_Args: {
           IRB.CreateInsertValue(InsVal, DFSF.getShadow(RI.getReturnValue()), 1);
       RI.setOperand(0, InsShadow);
       break;
     }
     }
   }
 }
 
-void DFSanVisitor::visitCallBase(CallBase &CB) {
-  Function *F = CB.getCalledFunction();
-  if ((F && F->isIntrinsic()) || CB.isInlineAsm()) {
-    visitOperandShadowInst(CB);
-    return;
-  }
-
-  // Calls to this function are synthesized in wrappers, and we shouldn't
-  // instrument them.
-  if (F == DFSF.DFS.DFSanVarargWrapperFn.getCallee()->stripPointerCasts())
-    return;
-
+bool DFSanVisitor::visitWrappedCallBase(Function &F, CallBase &CB) {
   IRBuilder<> IRB(&CB);
-
-  DenseMap<Value *, Function *>::iterator i =
-      DFSF.DFS.UnwrappedFnMap.find(CB.getCalledOperand());
-  if (i != DFSF.DFS.UnwrappedFnMap.end()) {
-    Function *F = i->second;
-    switch (DFSF.DFS.getWrapperKind(F)) {
-    case DataFlowSanitizer::WK_Warning:
-      CB.setCalledFunction(F);
-      IRB.CreateCall(DFSF.DFS.DFSanUnimplementedFn,
-                     IRB.CreateGlobalStringPtr(F->getName()));
-      DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
-      return;
-    case DataFlowSanitizer::WK_Discard:
-      CB.setCalledFunction(F);
-      DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
-      return;
-    case DataFlowSanitizer::WK_Functional:
-      CB.setCalledFunction(F);
-      visitOperandShadowInst(CB);
-      return;
-    case DataFlowSanitizer::WK_Custom:
-      // Don't try to handle invokes of custom functions, it's too complicated.
-      // Instead, invoke the dfsw$ wrapper, which will in turn call the __dfsw_
-      // wrapper.
-      if (CallInst *CI = dyn_cast<CallInst>(&CB)) {
-        FunctionType *FT = F->getFunctionType();
-        TransformedFunction CustomFn = DFSF.DFS.getCustomFunctionType(FT);
-        std::string CustomFName = "__dfsw_";
-        CustomFName += F->getName();
-        FunctionCallee CustomF = DFSF.DFS.Mod->getOrInsertFunction(
-            CustomFName, CustomFn.TransformedType);
-        if (Function *CustomFn = dyn_cast<Function>(CustomF.getCallee())) {
-          CustomFn->copyAttributesFrom(F);
+  switch (DFSF.DFS.getWrapperKind(&F)) {
+    CB.setCalledFunction(&F);
+                   IRB.CreateGlobalStringPtr(F.getName()));
+    return true;
+    CB.setCalledFunction(&F);
+    return true;
+    CB.setCalledFunction(&F);
+    return true;
+    CallInst *CI = dyn_cast<CallInst>(&CB);
+    if (!CI)
+      return false;
+
+    FunctionType *FT = F.getFunctionType();
+    CustomFName += F.getName();
+      CustomFn->copyAttributesFrom(&F);
 
-          // Custom functions returning non-void will write to the return label.
-          if (!FT->getReturnType()->isVoidTy()) {
-            CustomFn->removeAttributes(AttributeList::FunctionIndex,
-                                       DFSF.DFS.ReadOnlyNoneAttrs);
-          }
-        }
 
-        std::vector<Value *> Args;
 
-        auto i = CB.arg_begin();
-        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n) {
-          Type *T = (*i)->getType();
-          FunctionType *ParamFT;
-          if (isa<PointerType>(T) &&
-              (ParamFT = dyn_cast<FunctionType>(
-                   cast<PointerType>(T)->getElementType()))) {
-            std::string TName = "dfst";
-            TName += utostr(FT->getNumParams() - n);
-            TName += "$";
-            TName += F->getName();
-            Constant *T = DFSF.DFS.getOrBuildTrampolineFunction(ParamFT, TName);
-            Args.push_back(T);
-            Args.push_back(
-                IRB.CreateBitCast(*i, Type::getInt8PtrTy(*DFSF.DFS.Ctx)));
-          } else {
-            Args.push_back(*i);
-          }
-        }
+    // Adds non-variable arguments.
+    auto *I = CB.arg_begin();
+    for (unsigned n = FT->getNumParams(); n != 0; ++I, --n) {
+      Type *T = (*I)->getType();
+        TName += F.getName();
+            IRB.CreateBitCast(*I, Type::getInt8PtrTy(*DFSF.DFS.Ctx)));
+        Args.push_back(*I);
 
-        i = CB.arg_begin();
-        const unsigned ShadowArgStart = Args.size();
-        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
-          Args.push_back(
-              DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*i), &CB));
+    // Adds non-variable argument shadows.
+    I = CB.arg_begin();
+    for (unsigned N = FT->getNumParams(); N != 0; ++I, --N)
+      Args.push_back(DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*I), &CB));
 
-        if (FT->isVarArg()) {
-          auto *LabelVATy = ArrayType::get(DFSF.DFS.PrimitiveShadowTy,
-                                           CB.arg_size() - FT->getNumParams());
-          auto *LabelVAAlloca = new AllocaInst(
-              LabelVATy, getDataLayout().getAllocaAddrSpace(),
+    // Adds variable argument shadows.
+      auto *LabelVAAlloca =
+          new AllocaInst(LabelVATy, getDataLayout().getAllocaAddrSpace(),
                          "labelva", &DFSF.F->getEntryBlock().front());
 
-          for (unsigned n = 0; i != CB.arg_end(); ++i, ++n) {
-            auto LabelVAPtr = IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, n);
-            IRB.CreateStore(
-                DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*i), &CB),
+      for (unsigned N = 0; I != CB.arg_end(); ++I, ++N) {
+        auto *LabelVAPtr = IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, N);
+        IRB.CreateStore(DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*I), &CB),
                         LabelVAPtr);
-          }
 
-          Args.push_back(IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, 0));
-        }
 
-        if (!FT->getReturnType()->isVoidTy()) {
-          if (!DFSF.LabelReturnAlloca) {
-            DFSF.LabelReturnAlloca =
-                new AllocaInst(DFSF.DFS.PrimitiveShadowTy,
-                               getDataLayout().getAllocaAddrSpace(),
-                               "labelreturn", &DFSF.F->getEntryBlock().front());
-          }
-          Args.push_back(DFSF.LabelReturnAlloca);
-        }
+    // Adds the return value shadow.
+        DFSF.LabelReturnAlloca = new AllocaInst(
+            DFSF.DFS.PrimitiveShadowTy, getDataLayout().getAllocaAddrSpace(),
 
-        append_range(Args, drop_begin(CB.args(), FT->getNumParams()));
+    // Adds variable arguments.
 
-        CallInst *CustomCI = IRB.CreateCall(CustomF, Args);
-        CustomCI->setCallingConv(CI->getCallingConv());
-        CustomCI->setAttributes(TransformFunctionAttributes(CustomFn,
-            CI->getContext(), CI->getAttributes()));
+    CustomCI->setAttributes(TransformFunctionAttributes(
+        CustomFn, CI->getContext(), CI->getAttributes()));
 
-        // Update the parameter attributes of the custom call instruction to
-        // zero extend the shadow parameters. This is required for targets
-        // which consider PrimitiveShadowTy an illegal type.
-        for (unsigned n = 0; n < FT->getNumParams(); n++) {
-          const unsigned ArgNo = ShadowArgStart + n;
-          if (CustomCI->getArgOperand(ArgNo)->getType() ==
-              DFSF.DFS.PrimitiveShadowTy)
-            CustomCI->addParamAttr(ArgNo, Attribute::ZExt);
-        }
+    for (unsigned N = 0; N < FT->getNumParams(); N++) {
+      const unsigned ArgNo = ShadowArgStart + N;
 
-        if (!FT->getReturnType()->isVoidTy()) {
-          LoadInst *LabelLoad = IRB.CreateLoad(DFSF.DFS.PrimitiveShadowTy,
-                                               DFSF.LabelReturnAlloca);
-          DFSF.setShadow(CustomCI, DFSF.expandFromPrimitiveShadow(
-                                       FT->getReturnType(), LabelLoad, &CB));
-        }
+    // Loads the return value shadow.
+      LoadInst *LabelLoad =
+          IRB.CreateLoad(DFSF.DFS.PrimitiveShadowTy, DFSF.LabelReturnAlloca);
 
-        CI->replaceAllUsesWith(CustomCI);
-        CI->eraseFromParent();
-        return;
-      }
-      break;
-    }
+    return true;
+  return false;
+
+void DFSanVisitor::visitCallBase(CallBase &CB) {
+  Function *F = CB.getCalledFunction();
+  if ((F && F->isIntrinsic()) || CB.isInlineAsm()) {
+    visitOperandShadowInst(CB);
+    return;
   }
 
+  // Calls to this function are synthesized in wrappers, and we shouldn't
+  // instrument them.
+  if (F == DFSF.DFS.DFSanVarargWrapperFn.getCallee()->stripPointerCasts())
+    return;
+
+  DenseMap<Value *, Function *>::iterator i =
+      DFSF.DFS.UnwrappedFnMap.find(CB.getCalledOperand());
+  if (i != DFSF.DFS.UnwrappedFnMap.end())
+    if (visitWrappedCallBase(*i->second, CB))
+      return;
+
+  IRBuilder<> IRB(&CB);
+
   FunctionType *FT = CB.getFunctionType();
   if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
+    // Stores argument shadows.
     unsigned ArgOffset = 0;
     const DataLayout &DL = getDataLayout();
     for (unsigned I = 0, N = FT->getNumParams(); I != N; ++I) {
       unsigned Size =
           DL.getTypeAllocSize(DFSF.DFS.getShadowTy(FT->getParamType(I)));
       // Stop storing if arguments' size overflows. Inside a function, arguments
       // after overflow have zero shadow values.
       if (ArgOffset + Size > kArgTLSSize)
if (InvokeInst *II = dyn_cast<InvokeInst>(&CB)) {
         Next = &NewBB->front();
       }
     } else {
       assert(CB.getIterator() != CB.getParent()->end());
       Next = CB.getNextNode();
     }
 
     if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
+      // Loads the return value shadow.
       IRBuilder<> NextIRB(Next);
       const DataLayout &DL = getDataLayout();
       unsigned Size = DL.getTypeAllocSize(DFSF.DFS.getShadowTy(&CB));
       if (Size > kRetvalTLSSize) {
         // Set overflowed return shadow to be zero.
         DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
       } else {
         LoadInst *LI = NextIRB.CreateAlignedLoad(