1 files changed, 34 insertions, 16 deletions

diff --git a/include/llvm/Support/Allocator.h b/include/llvm/Support/Allocator.h
index c71759abd7d2..a5e662f4c588 100644
--- a/include/llvm/Support/Allocator.h
+++ b/include/llvm/Support/Allocator.h
@@ -1,4 +1,4 @@
-//===--- Allocator.h - Simple memory allocation abstraction -----*- C++ -*-===//
+//===- Allocator.h - Simple memory allocation abstraction -------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -144,19 +144,18 @@ public:
                 "that objects larger than a slab go into their own memory "
                 "allocation.");
 
-  BumpPtrAllocatorImpl()
-      : CurPtr(nullptr), End(nullptr), BytesAllocated(0), Allocator() {}
+  BumpPtrAllocatorImpl() = default;
+
   template <typename T>
   BumpPtrAllocatorImpl(T &&Allocator)
-      : CurPtr(nullptr), End(nullptr), BytesAllocated(0),
-        Allocator(std::forward<T &&>(Allocator)) {}
+      : Allocator(std::forward<T &&>(Allocator)) {}
 
   // Manually implement a move constructor as we must clear the old allocator's
   // slabs as a matter of correctness.
   BumpPtrAllocatorImpl(BumpPtrAllocatorImpl &&Old)
       : CurPtr(Old.CurPtr), End(Old.End), Slabs(std::move(Old.Slabs)),
         CustomSizedSlabs(std::move(Old.CustomSizedSlabs)),
-        BytesAllocated(Old.BytesAllocated),
+        BytesAllocated(Old.BytesAllocated), RedZoneSize(Old.RedZoneSize),
         Allocator(std::move(Old.Allocator)) {
     Old.CurPtr = Old.End = nullptr;
     Old.BytesAllocated = 0;
@@ -176,6 +175,7 @@ public:
     CurPtr = RHS.CurPtr;
     End = RHS.End;
     BytesAllocated = RHS.BytesAllocated;
+    RedZoneSize = RHS.RedZoneSize;
     Slabs = std::move(RHS.Slabs);
     CustomSizedSlabs = std::move(RHS.CustomSizedSlabs);
     Allocator = std::move(RHS.Allocator);
@@ -218,10 +218,16 @@ public:
     size_t Adjustment = alignmentAdjustment(CurPtr, Alignment);
     assert(Adjustment + Size >= Size && "Adjustment + Size must not overflow");
 
+    size_t SizeToAllocate = Size;
+#if LLVM_ADDRESS_SANITIZER_BUILD
+    // Add trailing bytes as a "red zone" under ASan.
+    SizeToAllocate += RedZoneSize;
+#endif
+
     // Check if we have enough space.
-    if (Adjustment + Size <= size_t(End - CurPtr)) {
+    if (Adjustment + SizeToAllocate <= size_t(End - CurPtr)) {
       char *AlignedPtr = CurPtr + Adjustment;
-      CurPtr = AlignedPtr + Size;
+      CurPtr = AlignedPtr + SizeToAllocate;
       // Update the allocation point of this memory block in MemorySanitizer.
       // Without this, MemorySanitizer messages for values originated from here
       // will point to the allocation of the entire slab.
@@ -232,7 +238,7 @@ public:
     }
 
     // If Size is really big, allocate a separate slab for it.
-    size_t PaddedSize = Size + Alignment - 1;
+    size_t PaddedSize = SizeToAllocate + Alignment - 1;
     if (PaddedSize > SizeThreshold) {
       void *NewSlab = Allocator.Allocate(PaddedSize, 0);
       // We own the new slab and don't want anyone reading anyting other than
@@ -251,10 +257,10 @@ public:
     // Otherwise, start a new slab and try again.
     StartNewSlab();
     uintptr_t AlignedAddr = alignAddr(CurPtr, Alignment);
-    assert(AlignedAddr + Size <= (uintptr_t)End &&
+    assert(AlignedAddr + SizeToAllocate <= (uintptr_t)End &&
            "Unable to allocate memory!");
     char *AlignedPtr = (char*)AlignedAddr;
-    CurPtr = AlignedPtr + Size;
+    CurPtr = AlignedPtr + SizeToAllocate;
     __msan_allocated_memory(AlignedPtr, Size);
     __asan_unpoison_memory_region(AlignedPtr, Size);
     return AlignedPtr;
@@ -283,6 +289,10 @@ public:
 
   size_t getBytesAllocated() const { return BytesAllocated; }
 
+  void setRedZoneSize(size_t NewSize) {
+    RedZoneSize = NewSize;
+  }
+
   void PrintStats() const {
     detail::printBumpPtrAllocatorStats(Slabs.size(), BytesAllocated,
                                        getTotalMemory());
@@ -292,10 +302,10 @@ private:
   /// \brief The current pointer into the current slab.
   ///
   /// This points to the next free byte in the slab.
-  char *CurPtr;
+  char *CurPtr = nullptr;
 
   /// \brief The end of the current slab.
-  char *End;
+  char *End = nullptr;
 
   /// \brief The slabs allocated so far.
   SmallVector<void *, 4> Slabs;
@@ -306,7 +316,11 @@ private:
   /// \brief How many bytes we've allocated.
   ///
   /// Used so that we can compute how much space was wasted.
-  size_t BytesAllocated;
+  size_t BytesAllocated = 0;
+
+  /// \brief The number of bytes to put between allocations when running under
+  /// a sanitizer.
+  size_t RedZoneSize = 1;
 
   /// \brief The allocator instance we use to get slabs of memory.
   AllocatorT Allocator;
@@ -357,7 +371,7 @@ private:
 };
 
 /// \brief The standard BumpPtrAllocator which just uses the default template
-/// paramaters.
+/// parameters.
 typedef BumpPtrAllocatorImpl<> BumpPtrAllocator;
 
 /// \brief A BumpPtrAllocator that allows only elements of a specific type to be
@@ -369,7 +383,11 @@ template <typename T> class SpecificBumpPtrAllocator {
   BumpPtrAllocator Allocator;
 
 public:
-  SpecificBumpPtrAllocator() = default;
+  SpecificBumpPtrAllocator() {
+    // Because SpecificBumpPtrAllocator walks the memory to call destructors,
+    // it can't have red zones between allocations.
+    Allocator.setRedZoneSize(0);
+  }
   SpecificBumpPtrAllocator(SpecificBumpPtrAllocator &&Old)
       : Allocator(std::move(Old.Allocator)) {}
   ~SpecificBumpPtrAllocator() { DestroyAll(); }