Reuse code for calculating allocation size for alloca

[bjope]

Summary

The AllocaInst class now got two methods for calculating the
allocation size for the alloca (in bits and in bytes).

This include minor refactoring in SafeStack, AddressSanitizer
and HWAddressSanitizer to use those methods.

[bjope]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

625

This method is used like this:

((!AI.isStaticAlloca()) || getAllocaSizeInBytes(AI) > 0)

So I am afraid this assert might blow (the old assert was guarded by AI.isArrayAllocation()).

[bjope]

This was an old idea that I never moved forward with. And I think the related code has been refactored since this review was created, so I'll just abandon the review.

include/llvm/IR/Instructions.h

public:
   const Value *getArraySize() const { return getOperand(0); }
   Value *getArraySize() { return getOperand(0); }
 
   /// Overload to return most specific pointer type.
   PointerType *getType() const {
     return cast<PointerType>(Instruction::getType());
   }
 
-  /// Get allocation size in bits. Returns None if size can't be determined,
-  /// e.g. in case of a VLA.
-  Optional<uint64_t> getAllocationSizeInBits(const DataLayout &DL) const;
+  /// Get allocation size. Returns None if the size can't be determined,
+  /// e.g. in case of a dynamic alloca.
+  Optional<uint64_t> getAllocationSize(const DataLayout &DL,
+                                       bool InBits = false) const;
+
+  /// Get allocation size in bits. Returns None if the size can't be determined,
+  /// e.g. in case of a dynamic alloca.
+  Optional<uint64_t> getAllocationSizeInBits(const DataLayout &DL) const {
+    return getAllocationSize(DL, true /* InBits */);
+  }
 
   /// Return the type that is being allocated by the instruction.
   Type *getAllocatedType() const { return AllocatedType; }
   /// for use only in special circumstances that need to generically
   /// transform a whole instruction (eg: IR linking and vectorization).
   void setAllocatedType(Type *Ty) { AllocatedType = Ty; }
 
   /// Return the alignment of the memory that is being allocated by the

lib/CodeGen/SafeStack.cpp

SafeStack(Function &F, const TargetLoweringBase &TL, const DataLayout &DL,
         Int8Ty(Type::getInt8Ty(F.getContext())) {}
 
   // Run the transformation on the associated function.
   // Returns whether the function was changed.
   bool run();
 };
 
 uint64_t SafeStack::getStaticAllocaAllocationSize(const AllocaInst* AI) {
-  uint64_t Size = DL.getTypeAllocSize(AI->getAllocatedType());
-  if (AI->isArrayAllocation()) {
-    auto C = dyn_cast<ConstantInt>(AI->getArraySize());
-    if (!C)
-      return 0;
-    Size *= C->getZExtValue();
-  }
-  return Size;
+  if (auto Size = AI->getAllocationSize(DL))
+    return *Size;
 }
 
 bool SafeStack::IsAccessSafe(Value *Addr, uint64_t AccessSize,
                              const Value *AllocaPtr, uint64_t AllocaSize) {
   AllocaOffsetRewriter Rewriter(SE, AllocaPtr);
   const SCEV *Expr = Rewriter.visit(SE.getSCEV(Addr));
 
   uint64_t BitWidth = SE.getTypeSizeInBits(Expr->getType());

lib/IR/Instructions.cpp

 
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
 //                            AllocaInst Class
 //===----------------------------------------------------------------------===//
 
 Optional<uint64_t>
-AllocaInst::getAllocationSizeInBits(const DataLayout &DL) const {
-  uint64_t Size = DL.getTypeAllocSizeInBits(getAllocatedType());
+AllocaInst::getAllocationSize(const DataLayout &DL, bool InBits) const {
+  if (!getAllocatedType()->isSized())
+    return None;
+  uint64_t Size = InBits ? DL.getTypeAllocSizeInBits(getAllocatedType())
+                         : DL.getTypeAllocSize(getAllocatedType());
   if (isArrayAllocation()) {
     auto C = dyn_cast<ConstantInt>(getArraySize());
     if (!C)
       return None;
     Size *= C->getZExtValue();
   }
   return Size;
 }

lib/Transforms/Instrumentation/AddressSanitizer.cpp

struct AddressSanitizer : public FunctionPass {
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
   }
 
   uint64_t getAllocaSizeInBytes(const AllocaInst &AI) const {
-    uint64_t ArraySize = 1;
-    if (AI.isArrayAllocation()) {
-      const ConstantInt *CI = dyn_cast<ConstantInt>(AI.getArraySize());
-      assert(CI && "non-constant array size");
-      ArraySize = CI->getZExtValue();
-    }
-    Type *Ty = AI.getAllocatedType();
-    uint64_t SizeInBytes =
-        AI.getModule()->getDataLayout().getTypeAllocSize(Ty);
-    return SizeInBytes * ArraySize;
+    const DataLayout &DL = AI.getModule()->getDataLayout();
+    Optional<uint64_t> Size = AI.getAllocationSize(DL);
+    assert(Size && "Unknown alloca size.");
+    return *Size;
   }
 
   /// Check if we want (and can) handle this alloca.
   bool isInterestingAlloca(const AllocaInst &AI);
 
   /// If it is an interesting memory access, return the PointerOperand
   /// and set IsWrite/Alignment. Otherwise return nullptr.
   /// MaybeMask is an output parameter for the mask Value, if we're looking at a

lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

IRB.CreateCall(HwasanMemoryAccessCallbackSized[IsWrite],
                    {AddrLong, ConstantInt::get(IntptrTy, TypeSize / 8)});
   }
   untagPointerOperand(I, Addr);
 
   return true;
 }
 
 static uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
-  uint64_t ArraySize = 1;
-  if (AI.isArrayAllocation()) {
-    const ConstantInt *CI = dyn_cast<ConstantInt>(AI.getArraySize());
-    assert(CI && "non-constant array size");
-    ArraySize = CI->getZExtValue();
-  }
-  Type *Ty = AI.getAllocatedType();
-  uint64_t SizeInBytes = AI.getModule()->getDataLayout().getTypeAllocSize(Ty);
-  return SizeInBytes * ArraySize;
+  const DataLayout &DL = AI.getModule()->getDataLayout();
+  Optional<uint64_t> Size = AI.getAllocationSize(DL);
+  assert(Size && "Unknown alloca size.");
+  return *Size;
 }
 
 bool HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI,
                                    Value *Tag) {
   size_t Size = (getAllocaSizeInBytes(*AI) + Mapping.getAllocaAlignment() - 1) &
                 ~(Mapping.getAllocaAlignment() - 1);
 
   Value *JustTag = IRB.CreateTrunc(Tag, IRB.getInt8Ty());