Re-committing changes from r76825 to BumpPtrAllocator with a fix and tests for

an off-by-one error.

llvm-svn: 76891

3 files changed, 320 insertions, 110 deletions

diff --git a/llvm/include/llvm/Support/Allocator.h b/llvm/include/llvm/Support/Allocator.h
index c0414f970a2..4c848788c73 100644
--- a/llvm/include/llvm/Support/Allocator.h
+++ b/llvm/include/llvm/Support/Allocator.h
@@ -15,6 +15,8 @@
 #define LLVM_SUPPORT_ALLOCATOR_H
 
 #include "llvm/Support/AlignOf.h"
+#include "llvm/Support/DataTypes.h"
+#include <cassert>
 #include <cstdlib>
 
 namespace llvm {
@@ -41,21 +43,104 @@ public:
   void PrintStats() const {}
 };
 
-/// BumpPtrAllocator - This allocator is useful for containers that need very
-/// simple memory allocation strategies.  In particular, this just keeps
+/// MemSlab - This structure lives at the beginning of every slab allocated by
+/// the bump allocator.
+class MemSlab {
+public:
+  size_t Size;
+  MemSlab *NextPtr;
+};
+
+/// SlabAllocator - This class can be used to parameterize the underlying
+/// allocation strategy for the bump allocator.  In particular, this is used
+/// by the JIT to allocate contiguous swathes of executable memory.  The
+/// interface uses MemSlab's instead of void *'s so that the allocator
+/// doesn't have to remember the size of the pointer it allocated.
+class SlabAllocator {
+public:
+  virtual ~SlabAllocator();
+  virtual MemSlab *Allocate(size_t Size) = 0;
+  virtual void Deallocate(MemSlab *Slab) = 0;
+};
+
+/// MallocSlabAllocator - The default slab allocator for the bump allocator
+/// is an adapter class for MallocAllocator that just forwards the method
+/// calls and translates the arguments.
+class MallocSlabAllocator : public SlabAllocator {
+  /// Allocator - The underlying allocator that we forward to.
+  ///
+  MallocAllocator Allocator;
+
+public:
+  MallocSlabAllocator() : Allocator() { }
+  virtual ~MallocSlabAllocator();
+  virtual MemSlab *Allocate(size_t Size);
+  virtual void Deallocate(MemSlab *Slab);
+};
+
+/// BumpPtrAllocator - This allocator is useful for containers that need
+/// very simple memory allocation strategies.  In particular, this just keeps
 /// allocating memory, and never deletes it until the entire block is dead. This
 /// makes allocation speedy, but must only be used when the trade-off is ok.
 class BumpPtrAllocator {
   BumpPtrAllocator(const BumpPtrAllocator &); // do not implement
   void operator=(const BumpPtrAllocator &);   // do not implement
 
-  void *TheMemory;
+  /// SlabSize - Allocate data into slabs of this size unless we get an
+  /// allocation above SizeThreshold.
+  size_t SlabSize;
+
+  /// SizeThreshold - For any allocation larger than this threshold, we should
+  /// allocate a separate slab.
+  size_t SizeThreshold;
+
+  /// Allocator - The underlying allocator we use to get slabs of memory.  This
+  /// defaults to MallocSlabAllocator, which wraps malloc, but it could be
+  /// changed to use a custom allocator.
+  SlabAllocator &Allocator;
+
+  /// CurSlab - The slab that we are currently allocating into.
+  ///
+  MemSlab *CurSlab;
+
+  /// CurPtr - The current pointer into the current slab.  This points to the
+  /// next free byte in the slab.
+  char *CurPtr;
+
+  /// End - The end of the current slab.
+  ///
+  char *End;
+
+  /// BytesAllocated - This field tracks how many bytes we've allocated, so
+  /// that we can compute how much space was wasted.
+  size_t BytesAllocated;
+
+  /// AlignPtr - Align Ptr to Alignment bytes, rounding up.  Alignment should
+  /// be a power of two.  This method rounds up, so AlignPtr(7, 4) == 8 and
+  /// AlignPtr(8, 4) == 8.
+  static char *AlignPtr(char *Ptr, size_t Alignment);
+
+  /// StartNewSlab - Allocate a new slab and move the bump pointers over into
+  /// the new slab.  Modifies CurPtr and End.
+  void StartNewSlab();
+
+  /// DeallocateSlabs - Deallocate all memory slabs after and including this
+  /// one.
+  void DeallocateSlabs(MemSlab *Slab);
+
+  static MallocSlabAllocator DefaultSlabAllocator;
+
 public:
-  BumpPtrAllocator();
+  BumpPtrAllocator(size_t size = 4096, size_t threshold = 4096,
+                   SlabAllocator &allocator = DefaultSlabAllocator);
   ~BumpPtrAllocator();
 
+  /// Reset - Deallocate all but the current slab and reset the current pointer
+  /// to the beginning of it, freeing all memory allocated so far.
   void Reset();
 
+  /// Allocate - Allocate space at the specified alignment.
+  ///
   void *Allocate(size_t Size, size_t Alignment);
 
   /// Allocate space, but do not construct, one object.
@@ -83,9 +168,11 @@ public:
 
   void Deallocate(const void * /*Ptr*/) {}
 
+  unsigned GetNumSlabs() const;
+
   void PrintStats() const;
 };
 
 }  // end namespace llvm
 
-#endif
+#endif // LLVM_SUPPORT_ALLOCATOR_H
diff --git a/llvm/lib/Support/Allocator.cpp b/llvm/lib/Support/Allocator.cpp
index db0d8f31e55..4e4a75ee58c 100644
--- a/llvm/lib/Support/Allocator.cpp
+++ b/llvm/lib/Support/Allocator.cpp
@@ -12,130 +12,158 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/Allocator.h"
-#include "llvm/Support/Recycler.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/Recycler.h"
 #include "llvm/Support/Streams.h"
-#include <ostream>
-using namespace llvm;
+#include <cstring>
 
-//===----------------------------------------------------------------------===//
-// MemRegion class implementation
-//===----------------------------------------------------------------------===//
+namespace llvm {
 
-namespace {
-/// MemRegion - This is one chunk of the BumpPtrAllocator.
-class MemRegion {
-  unsigned RegionSize;
-  MemRegion *Next;
-  char *NextPtr;
-public:
-  void Init(unsigned size, unsigned Alignment, MemRegion *next) {
-    RegionSize = size;
-    Next = next;
-    NextPtr = (char*)(this+1);
-    
-    // Align NextPtr.
-    NextPtr = (char*)((intptr_t)(NextPtr+Alignment-1) &
-                      ~(intptr_t)(Alignment-1));
-  }
-  
-  const MemRegion *getNext() const { return Next; }
-  unsigned getNumBytesAllocated() const {
-    return NextPtr-(const char*)this;
-  }
-  
-  /// Allocate - Allocate and return at least the specified number of bytes.
-  ///
-  void *Allocate(size_t AllocSize, size_t Alignment, MemRegion **RegPtr) {
-    
-    char* Result = (char*) (((uintptr_t) (NextPtr+Alignment-1)) 
-                            & ~((uintptr_t) Alignment-1));
-
-    // Speculate the new value of NextPtr.
-    char* NextPtrTmp = Result + AllocSize;
-    
-    // If we are still within the current region, return Result.
-    if (unsigned (NextPtrTmp - (char*) this) <= RegionSize) {
-      NextPtr = NextPtrTmp;
-      return Result;
-    }
-    
-    // Otherwise, we have to allocate a new chunk.  Create one twice as big as
-    // this one.
-    MemRegion *NewRegion = (MemRegion *)malloc(RegionSize*2);
-    NewRegion->Init(RegionSize*2, Alignment, this);
-
-    // Update the current "first region" pointer  to point to the new region.
-    *RegPtr = NewRegion;
-    
-    // Try allocating from it now.
-    return NewRegion->Allocate(AllocSize, Alignment, RegPtr);
-  }
-  
-  /// Deallocate - Recursively release all memory for this and its next regions
-  /// to the system.
-  void Deallocate() {
-    MemRegion *next = Next;
-    free(this);
-    if (next)
-      next->Deallocate();
-  }
+BumpPtrAllocator::BumpPtrAllocator(size_t size, size_t threshold,
+                                   SlabAllocator &allocator)
+    : SlabSize(size), SizeThreshold(threshold), Allocator(allocator),
+      CurSlab(0), BytesAllocated(0) {
+  StartNewSlab();
+}
 
-  /// DeallocateAllButLast - Recursively release all memory for this and its
-  /// next regions to the system stopping at the last region in the list.
-  /// Returns the pointer to the last region.
-  MemRegion *DeallocateAllButLast() {
-    MemRegion *next = Next;
-    if (!next)
-      return this;
-    free(this);
-    return next->DeallocateAllButLast();
-  }
-};
+BumpPtrAllocator::~BumpPtrAllocator() {
+  DeallocateSlabs(CurSlab);
 }
 
-//===----------------------------------------------------------------------===//
-// BumpPtrAllocator class implementation
-//===----------------------------------------------------------------------===//
+/// AlignPtr - Align Ptr to Alignment bytes, rounding up.  Alignment should
+/// be a power of two.  This method rounds up, so AlignPtr(7, 4) == 8 and
+/// AlignPtr(8, 4) == 8.
+char *BumpPtrAllocator::AlignPtr(char *Ptr, size_t Alignment) {
+  assert(Alignment && (Alignment & (Alignment - 1)) == 0 &&
+         "Alignment is not a power of two!");
 
-BumpPtrAllocator::BumpPtrAllocator() {
-  TheMemory = malloc(4096);
-  ((MemRegion*)TheMemory)->Init(4096, 1, 0);
+  // Do the alignment.
+  return (char*)(((uintptr_t)Ptr + Alignment - 1) &
+                 ~(uintptr_t)(Alignment - 1));
 }
 
-BumpPtrAllocator::~BumpPtrAllocator() {
-  ((MemRegion*)TheMemory)->Deallocate();
+/// StartNewSlab - Allocate a new slab and move the bump pointers over into
+/// the new slab.  Modifies CurPtr and End.
+void BumpPtrAllocator::StartNewSlab() {
+  MemSlab *NewSlab = Allocator.Allocate(SlabSize);
+  NewSlab->NextPtr = CurSlab;
+  CurSlab = NewSlab;
+  CurPtr = (char*)(CurSlab + 1);
+  End = ((char*)CurSlab) + CurSlab->Size;
+}
+
+/// DeallocateSlabs - Deallocate all memory slabs after and including this
+/// one.
+void BumpPtrAllocator::DeallocateSlabs(MemSlab *Slab) {
+  while (Slab) {
+    MemSlab *NextSlab = Slab->NextPtr;
+#ifndef NDEBUG
+    // Poison the memory so stale pointers crash sooner.  Note we must
+    // preserve the Size and NextPtr fields at the beginning.
+    memset(Slab + 1, 0xCD, Slab->Size - sizeof(MemSlab));
+#endif
+    Allocator.Deallocate(Slab);
+    Slab = NextSlab;
+  }
 }
 
+/// Reset - Deallocate all but the current slab and reset the current pointer
+/// to the beginning of it, freeing all memory allocated so far.
 void BumpPtrAllocator::Reset() {
-  MemRegion *MRP = (MemRegion*)TheMemory;
-  MRP = MRP->DeallocateAllButLast();
-  MRP->Init(4096, 1, 0);
-  TheMemory = MRP;
+  DeallocateSlabs(CurSlab->NextPtr);
+  CurSlab->NextPtr = 0;
+  CurPtr = (char*)(CurSlab + 1);
+  End = ((char*)CurSlab) + CurSlab->Size;
 }
 
-void *BumpPtrAllocator::Allocate(size_t Size, size_t Align) {
-  MemRegion *MRP = (MemRegion*)TheMemory;
-  void *Ptr = MRP->Allocate(Size, Align, &MRP);
-  TheMemory = MRP;
+/// Allocate - Allocate space at the specified alignment.
+///
+void *BumpPtrAllocator::Allocate(size_t Size, size_t Alignment) {
+  // Keep track of how many bytes we've allocated.
+  BytesAllocated += Size;
+
+  // 0-byte alignment means 1-byte alignment.
+  if (Alignment == 0) Alignment = 1;
+
+  // Allocate the aligned space, going forwards from CurPtr.
+  char *Ptr = AlignPtr(CurPtr, Alignment);
+
+  // Check if we can hold it.
+  if (Ptr + Size <= End) {
+    CurPtr = Ptr + Size;
+    return Ptr;
+  }
+
+  // If Size is really big, allocate a separate slab for it.
+  if (Size > SizeThreshold) {
+    size_t PaddedSize = Size + sizeof(MemSlab) + Alignment - 1;
+    MemSlab *NewSlab = Allocator.Allocate(PaddedSize);
+
+    // Put the new slab after the current slab, since we are not allocating
+    // into it.
+    NewSlab->NextPtr = CurSlab->NextPtr;
+    CurSlab->NextPtr = NewSlab;
+
+    Ptr = AlignPtr((char*)(NewSlab + 1), Alignment);
+    assert((uintptr_t)Ptr + Size <= (uintptr_t)NewSlab + NewSlab->Size);
+    return Ptr;
+  }
+
+  // Otherwise, start a new slab and try again.
+  StartNewSlab();
+  Ptr = AlignPtr(CurPtr, Alignment);
+  CurPtr = Ptr + Size;
+  assert(CurPtr <= End && "Unable to allocate memory!");
   return Ptr;
 }
 
+unsigned BumpPtrAllocator::GetNumSlabs() const {
+  unsigned NumSlabs = 0;
+  for (MemSlab *Slab = CurSlab; Slab != 0; Slab = Slab->NextPtr) {
+    ++NumSlabs;
+  }
+  return NumSlabs;
+}
+
 void BumpPtrAllocator::PrintStats() const {
-  unsigned BytesUsed = 0;
-  unsigned NumRegions = 0;
-  const MemRegion *R = (MemRegion*)TheMemory;
-  for (; R; R = R->getNext(), ++NumRegions)
-    BytesUsed += R->getNumBytesAllocated();
-
-  cerr << "\nNumber of memory regions: " << NumRegions << "\n";
-  cerr << "Bytes allocated: " << BytesUsed << "\n";
+  unsigned NumSlabs = 0;
+  size_t TotalMemory = 0;
+  for (MemSlab *Slab = CurSlab; Slab != 0; Slab = Slab->NextPtr) {
+    TotalMemory += Slab->Size;
+    ++NumSlabs;
+  }
+
+  cerr << "\nNumber of memory regions: " << NumSlabs << '\n'
+       << "Bytes used: " << BytesAllocated << '\n'
+       << "Bytes allocated: " << TotalMemory << '\n'
+       << "Bytes wasted: " << (TotalMemory - BytesAllocated)
+       << " (includes alignment, etc)\n";
+}
+
+MallocSlabAllocator BumpPtrAllocator::DefaultSlabAllocator =
+  MallocSlabAllocator();
+
+SlabAllocator::~SlabAllocator() { }
+
+MallocSlabAllocator::~MallocSlabAllocator() { }
+
+MemSlab *MallocSlabAllocator::Allocate(size_t Size) {
+  MemSlab *Slab = (MemSlab*)Allocator.Allocate(Size, 0);
+  Slab->Size = Size;
+  Slab->NextPtr = 0;
+  return Slab;
+}
+
+void MallocSlabAllocator::Deallocate(MemSlab *Slab) {
+  Allocator.Deallocate(Slab);
+}
+
+void PrintRecyclerStats(size_t Size,
+                        size_t Align,
+                        size_t FreeListSize) {
+  cerr << "Recycler element size: " << Size << '\n'
+       << "Recycler element alignment: " << Align << '\n'
+       << "Number of elements free for recycling: " << FreeListSize << '\n';
 }
 
-void llvm::PrintRecyclerStats(size_t Size,
-                              size_t Align,
-                              size_t FreeListSize) {
-  cerr << "Recycler element size: " << Size << '\n';
-  cerr << "Recycler element alignment: " << Align << '\n';
-  cerr << "Number of elements free for recycling: " << FreeListSize << '\n';
 }
diff --git a/llvm/unittests/Support/AllocatorTest.cpp b/llvm/unittests/Support/AllocatorTest.cpp
index e69de29bb2d..cc3296a8d01 100644
--- a/llvm/unittests/Support/AllocatorTest.cpp
+++ b/llvm/unittests/Support/AllocatorTest.cpp
@@ -0,0 +1,95 @@
+//===- llvm/unittest/Support/AllocatorTest.cpp - BumpPtrAllocator tests ---===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Allocator.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+TEST(AllocatorTest, Basics) {
+  BumpPtrAllocator Alloc;
+  int *a = (int*)Alloc.Allocate(sizeof(int), 0);
+  int *b = (int*)Alloc.Allocate(sizeof(int) * 10, 0);
+  int *c = (int*)Alloc.Allocate(sizeof(int), 0);
+  *a = 1;
+  b[0] = 2;
+  b[9] = 2;
+  *c = 3;
+  EXPECT_EQ(1, *a);
+  EXPECT_EQ(2, b[0]);
+  EXPECT_EQ(2, b[9]);
+  EXPECT_EQ(3, *c);
+  EXPECT_EQ(1U, Alloc.GetNumSlabs());
+}
+
+// Allocate enough bytes to create three slabs.
+TEST(AllocatorTest, ThreeSlabs) {
+  BumpPtrAllocator Alloc(4096, 4096);
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(1U, Alloc.GetNumSlabs());
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(2U, Alloc.GetNumSlabs());
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(3U, Alloc.GetNumSlabs());
+}
+
+// Allocate enough bytes to create two slabs, reset the allocator, and do it
+// again.
+TEST(AllocatorTest, TestReset) {
+  BumpPtrAllocator Alloc(4096, 4096);
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(1U, Alloc.GetNumSlabs());
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(2U, Alloc.GetNumSlabs());
+  Alloc.Reset();
+  EXPECT_EQ(1U, Alloc.GetNumSlabs());
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(1U, Alloc.GetNumSlabs());
+  Alloc.Allocate(3000, 0);
+  EXPECT_EQ(2U, Alloc.GetNumSlabs());
+}
+
+// Test some allocations at varying alignments.
+TEST(AllocatorTest, TestAlignment) {
+  BumpPtrAllocator Alloc;
+  uintptr_t a;
+  a = (uintptr_t)Alloc.Allocate(1, 2);
+  EXPECT_EQ(0U, a & 1);
+  a = (uintptr_t)Alloc.Allocate(1, 4);
+  EXPECT_EQ(0U, a & 3);
+  a = (uintptr_t)Alloc.Allocate(1, 8);
+  EXPECT_EQ(0U, a & 7);
+  a = (uintptr_t)Alloc.Allocate(1, 16);
+  EXPECT_EQ(0U, a & 15);
+  a = (uintptr_t)Alloc.Allocate(1, 32);
+  EXPECT_EQ(0U, a & 31);
+  a = (uintptr_t)Alloc.Allocate(1, 64);
+  EXPECT_EQ(0U, a & 63);
+  a = (uintptr_t)Alloc.Allocate(1, 128);
+  EXPECT_EQ(0U, a & 127);
+}
+
+// Test allocating just over the slab size.  This tests a bug where before the
+// allocator incorrectly calculated the buffer end pointer.
+TEST(AllocatorTest, TestOverflow) {
+  BumpPtrAllocator Alloc(4096, 4096);
+
+  // Fill the slab right up until the end pointer.
+  Alloc.Allocate(4096 - sizeof(MemSlab), 0);
+  EXPECT_EQ(1U, Alloc.GetNumSlabs());
+
+  // If we dont't allocate a new slab, then we will have overflowed.
+  Alloc.Allocate(1, 0);
+  EXPECT_EQ(2U, Alloc.GetNumSlabs());
+}
+
+}  // anonymous namespace