[dfsan] Refactor loadShadow

[stephan.yichao.zhao]

Summary

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

[morehouse]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

546

Long name. WDYT about renaming to loadLegacyShadowFast?

552

WDYT about renaming to loadFast16ShadowFast?

1580

Please fix this lint.

[stephan.yichao.zhao]

shorten names
fixed lint warnings

[gbalats]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

549–551

This comment is more helpful in the implementation of the function. Same for the second sentence of the previous declaration.

[stephan.yichao.zhao]

moved some comments to method implementation.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

private:
   template <class AggregateType>
   Value *collapseAggregateShadow(AggregateType *AT, Value *Shadow,
                                  IRBuilder<> &IRB);
 
   Value *collapseToPrimitiveShadow(Value *Shadow, IRBuilder<> &IRB);
 
   /// Returns the shadow value of an argument A.
   Value *getShadowForTLSArgument(Argument *A);
+
+  /// The fast path of loading shadow in legacy mode.
+  Value *loadLegacyShadowFast(Value *ShadowAddr, uint64_t Size,
+                              Align ShadowAlign, Instruction *Pos);
+
+  /// The fast path of loading shadow in fast-16-label mode.
+  Value *loadFast16ShadowFast(Value *ShadowAddr, uint64_t Size,
+                              Align ShadowAlign, Instruction *Pos);
 };
 
 class DFSanVisitor : public InstVisitor<DFSanVisitor> {
 public:
   DFSanFunction &DFSF;
 
   DFSanVisitor(DFSanFunction &DFSF) : DFSF(DFSF) {}
 
   const DataLayout &getDataLayout() const {
     return DFSF.F->getParent()->getDataLayout();
   }

 }
 
 Value *DFSanVisitor::visitOperandShadowInst(Instruction &I) {
   Value *CombinedShadow = DFSF.combineOperandShadows(&I);
   DFSF.setShadow(&I, CombinedShadow);
   return CombinedShadow;
 }
 
-// Generates IR to load shadow corresponding to bytes [Addr, Addr+Size), where
-// Addr has alignment Align, and take the union of each of those shadows. The
-// returned shadow always has primitive type.
-Value *DFSanFunction::loadShadow(Value *Addr, uint64_t Size, uint64_t Align,
+Value *DFSanFunction::loadFast16ShadowFast(Value *ShadowAddr, uint64_t Size,
+                                           Align ShadowAlign,
                                            Instruction *Pos) {
-  if (AllocaInst *AI = dyn_cast<AllocaInst>(Addr)) {
-    const auto i = AllocaShadowMap.find(AI);
-    if (i != AllocaShadowMap.end()) {
-      IRBuilder<> IRB(Pos);
-      return IRB.CreateLoad(DFS.PrimitiveShadowTy, i->second);
-    }
-  }
-
-  const llvm::Align ShadowAlign(Align * DFS.ShadowWidthBytes);
-  SmallVector<const Value *, 2> Objs;
-  getUnderlyingObjects(Addr, Objs);
-  bool AllConstants = true;
-  for (const Value *Obj : Objs) {
-    if (isa<Function>(Obj) || isa<BlockAddress>(Obj))
-      continue;
-    if (isa<GlobalVariable>(Obj) && cast<GlobalVariable>(Obj)->isConstant())
-      continue;
-
-    AllConstants = false;
-    break;
-  }
-  if (AllConstants)
-    return DFS.ZeroPrimitiveShadow;
-
-  Value *ShadowAddr = DFS.getShadowAddress(Addr, Pos);
-  switch (Size) {
-  case 0:
-    return DFS.ZeroPrimitiveShadow;
-  case 1: {
-    LoadInst *LI = new LoadInst(DFS.PrimitiveShadowTy, ShadowAddr, "", Pos);
-    LI->setAlignment(ShadowAlign);
-    return LI;
-  }
-  case 2: {
-    IRBuilder<> IRB(Pos);
-    Value *ShadowAddr1 = IRB.CreateGEP(DFS.PrimitiveShadowTy, ShadowAddr,
-                                       ConstantInt::get(DFS.IntptrTy, 1));
-    return combineShadows(
-        IRB.CreateAlignedLoad(DFS.PrimitiveShadowTy, ShadowAddr, ShadowAlign),
-        IRB.CreateAlignedLoad(DFS.PrimitiveShadowTy, ShadowAddr1, ShadowAlign),
-        Pos);
-  }
-  }
-
-  if (ClFast16Labels && Size % (64 / DFS.ShadowWidthBits) == 0) {
-    // First OR all the WideShadows, then OR individual shadows within the
-    // combined WideShadow.  This is fewer instructions than ORing shadows
-    // individually.
-    IRBuilder<> IRB(Pos);
-    Value *WideAddr =
-        IRB.CreateBitCast(ShadowAddr, Type::getInt64PtrTy(*DFS.Ctx));
-    Value *CombinedWideShadow =
-        IRB.CreateAlignedLoad(IRB.getInt64Ty(), WideAddr, ShadowAlign);
-    for (uint64_t Ofs = 64 / DFS.ShadowWidthBits; Ofs != Size;
-         Ofs += 64 / DFS.ShadowWidthBits) {
-      WideAddr = IRB.CreateGEP(Type::getInt64Ty(*DFS.Ctx), WideAddr,
-                               ConstantInt::get(DFS.IntptrTy, 1));
-      Value *NextWideShadow =
-          IRB.CreateAlignedLoad(IRB.getInt64Ty(), WideAddr, ShadowAlign);
-      CombinedWideShadow = IRB.CreateOr(CombinedWideShadow, NextWideShadow);
-    }
-    for (unsigned Width = 32; Width >= DFS.ShadowWidthBits; Width >>= 1) {
-      Value *ShrShadow = IRB.CreateLShr(CombinedWideShadow, Width);
-      CombinedWideShadow = IRB.CreateOr(CombinedWideShadow, ShrShadow);
-    }
-    return IRB.CreateTrunc(CombinedWideShadow, DFS.PrimitiveShadowTy);
-  }
-  if (!AvoidNewBlocks && Size % (64 / DFS.ShadowWidthBits) == 0) {
-    // Fast path for the common case where each byte has identical shadow: load
-    // shadow 64 bits at a time, fall out to a __dfsan_union_load call if any
-    // shadow is non-equal.
-    BasicBlock *FallbackBB = BasicBlock::Create(*DFS.Ctx, "", F);
-    IRBuilder<> FallbackIRB(FallbackBB);
-    CallInst *FallbackCall = FallbackIRB.CreateCall(
-        DFS.DFSanUnionLoadFn,
-        {ShadowAddr, ConstantInt::get(DFS.IntptrTy, Size)});
-    FallbackCall->addAttribute(AttributeList::ReturnIndex, Attribute::ZExt);
+  // combined WideShadow. This is fewer instructions than ORing shadows
+
+Value *DFSanFunction::loadLegacyShadowFast(Value *ShadowAddr, uint64_t Size,
+                                           Align ShadowAlign,
+                                           Instruction *Pos) {
+      DFS.DFSanUnionLoadFn, {ShadowAddr, ConstantInt::get(DFS.IntptrTy, Size)});
 
-    // Compare each of the shadows stored in the loaded 64 bits to each other,
-    // by computing (WideShadow rotl ShadowWidthBits) == WideShadow.
-    IRBuilder<> IRB(Pos);
-    Value *WideAddr =
-        IRB.CreateBitCast(ShadowAddr, Type::getInt64PtrTy(*DFS.Ctx));
-    Value *WideShadow =
-        IRB.CreateAlignedLoad(IRB.getInt64Ty(), WideAddr, ShadowAlign);
-    Value *TruncShadow = IRB.CreateTrunc(WideShadow, DFS.PrimitiveShadowTy);
-    Value *ShlShadow = IRB.CreateShl(WideShadow, DFS.ShadowWidthBits);
-    Value *ShrShadow = IRB.CreateLShr(WideShadow, 64 - DFS.ShadowWidthBits);
-    Value *RotShadow = IRB.CreateOr(ShlShadow, ShrShadow);
-    Value *ShadowsEq = IRB.CreateICmpEQ(WideShadow, RotShadow);
 
-    BasicBlock *Head = Pos->getParent();
-    BasicBlock *Tail = Head->splitBasicBlock(Pos->getIterator());
 
-    if (DomTreeNode *OldNode = DT.getNode(Head)) {
-      std::vector<DomTreeNode *> Children(OldNode->begin(), OldNode->end());
 
-      DomTreeNode *NewNode = DT.addNewBlock(Tail, Head);
-      for (auto Child : Children)
-        DT.changeImmediateDominator(Child, NewNode);
-    }
+    for (auto *Child : Children)
 
-    // In the following code LastBr will refer to the previous basic block's
-    // conditional branch instruction, whose true successor is fixed up to point
-    // to the next block during the loop below or to the tail after the final
-    // iteration.
-    BranchInst *LastBr = BranchInst::Create(FallbackBB, FallbackBB, ShadowsEq);
-    ReplaceInstWithInst(Head->getTerminator(), LastBr);
-    DT.addNewBlock(FallbackBB, Head);
 
-    for (uint64_t Ofs = 64 / DFS.ShadowWidthBits; Ofs != Size;
-         Ofs += 64 / DFS.ShadowWidthBits) {
-      BasicBlock *NextBB = BasicBlock::Create(*DFS.Ctx, "", F);
-      DT.addNewBlock(NextBB, LastBr->getParent());
-      IRBuilder<> NextIRB(NextBB);
-      WideAddr = NextIRB.CreateGEP(Type::getInt64Ty(*DFS.Ctx), WideAddr,
-                                   ConstantInt::get(DFS.IntptrTy, 1));
-      Value *NextWideShadow = NextIRB.CreateAlignedLoad(NextIRB.getInt64Ty(),
-                                                        WideAddr, ShadowAlign);
-      ShadowsEq = NextIRB.CreateICmpEQ(WideShadow, NextWideShadow);
-      LastBr->setSuccessor(0, NextBB);
-      LastBr = NextIRB.CreateCondBr(ShadowsEq, FallbackBB, FallbackBB);
-    }
+    Value *NextWideShadow =
+        NextIRB.CreateAlignedLoad(NextIRB.getInt64Ty(), WideAddr, ShadowAlign);
 
-    LastBr->setSuccessor(0, Tail);
-    FallbackIRB.CreateBr(Tail);
-    PHINode *Shadow =
-        PHINode::Create(DFS.PrimitiveShadowTy, 2, "", &Tail->front());
-    Shadow->addIncoming(FallbackCall, FallbackBB);
-    Shadow->addIncoming(TruncShadow, LastBr->getParent());
-    return Shadow;
-  }
 
+// Generates IR to load shadow corresponding to bytes [Addr, Addr+Size), where
+// Addr has alignment Align, and take the union of each of those shadows. The
+// returned shadow always has primitive type.
+Value *DFSanFunction::loadShadow(Value *Addr, uint64_t Size, uint64_t Align,
+                                 Instruction *Pos) {
+  if (AllocaInst *AI = dyn_cast<AllocaInst>(Addr)) {
+    const auto i = AllocaShadowMap.find(AI);
+    if (i != AllocaShadowMap.end()) {
+      IRBuilder<> IRB(Pos);
+      return IRB.CreateLoad(DFS.PrimitiveShadowTy, i->second);
+    }
+  }
+
+  const llvm::Align ShadowAlign(Align * DFS.ShadowWidthBytes);
+  SmallVector<const Value *, 2> Objs;
+  getUnderlyingObjects(Addr, Objs);
+  bool AllConstants = true;
+  for (const Value *Obj : Objs) {
+    if (isa<Function>(Obj) || isa<BlockAddress>(Obj))
+      continue;
+    if (isa<GlobalVariable>(Obj) && cast<GlobalVariable>(Obj)->isConstant())
+      continue;
+
+    AllConstants = false;
+    break;
+  }
+  if (AllConstants)
+    return DFS.ZeroPrimitiveShadow;
+
+  Value *ShadowAddr = DFS.getShadowAddress(Addr, Pos);
+  switch (Size) {
+  case 0:
+    return DFS.ZeroPrimitiveShadow;
+  case 1: {
+    LoadInst *LI = new LoadInst(DFS.PrimitiveShadowTy, ShadowAddr, "", Pos);
+    LI->setAlignment(ShadowAlign);
+    return LI;
+  }
+  case 2: {
+    IRBuilder<> IRB(Pos);
+    Value *ShadowAddr1 = IRB.CreateGEP(DFS.PrimitiveShadowTy, ShadowAddr,
+                                       ConstantInt::get(DFS.IntptrTy, 1));
+    return combineShadows(
+        IRB.CreateAlignedLoad(DFS.PrimitiveShadowTy, ShadowAddr, ShadowAlign),
+        IRB.CreateAlignedLoad(DFS.PrimitiveShadowTy, ShadowAddr1, ShadowAlign),
+        Pos);
+  }
+  }
+
+  if (ClFast16Labels && Size % (64 / DFS.ShadowWidthBits) == 0)
+    return loadFast16ShadowFast(ShadowAddr, Size, ShadowAlign, Pos);
+
+  if (!AvoidNewBlocks && Size % (64 / DFS.ShadowWidthBits) == 0)
+    return loadLegacyShadowFast(ShadowAddr, Size, ShadowAlign, Pos);
+
   IRBuilder<> IRB(Pos);
   FunctionCallee &UnionLoadFn =
       ClFast16Labels ? DFS.DFSanUnionLoadFast16LabelsFn : DFS.DFSanUnionLoadFn;
   CallInst *FallbackCall = IRB.CreateCall(
       UnionLoadFn, {ShadowAddr, ConstantInt::get(DFS.IntptrTy, Size)});
   FallbackCall->addAttribute(AttributeList::ReturnIndex, Attribute::ZExt);
   return FallbackCall;
 }