From 4f240010fdea1fcbe7dac3c8e9ebad082e4036c2 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans@hanshq.net>
Date: Sun, 7 Sep 2014 04:24:31 +0000
Subject: [PATCH] BumpPtrAllocator: do the size check without moving any
 pointers

Instead of aligning and moving the CurPtr forward, and then comparing
with End, simply calculate how much space is needed, and compare that
to how much is available.

Hopefully this avoids any doubts about comparing addresses possibly
derived from past the end of the slab array, overflowing, etc.

Also add a test where aligning CurPtr would move it past End.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@217330 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/llvm/Support/Allocator.h    | 12 ++++++------
 include/llvm/Support/MathExtras.h   |  8 ++++++++
 unittests/Support/AllocatorTest.cpp | 12 ++++++++++++
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/include/llvm/Support/Allocator.h b/include/llvm/Support/Allocator.h
index 05dccec94de..ff08622eee5 100644
--- a/include/llvm/Support/Allocator.h
+++ b/include/llvm/Support/Allocator.h
@@ -209,12 +209,12 @@ public:
     // Keep track of how many bytes we've allocated.
     BytesAllocated += Size;
 
-    // Allocate the aligned space, going forwards from CurPtr.
-    uintptr_t AlignedAddr = alignAddr(CurPtr, Alignment);
+    size_t Adjustment = alignmentAdjustment(CurPtr, Alignment);
+    assert(Adjustment + Size >= Size && "Adjustment + Size must not overflow");
 
-    // Check if we can hold it.
-    if (AlignedAddr + Size <= (uintptr_t)End) {
-      char *AlignedPtr = (char*)AlignedAddr;
+    // Check if we have enough space.
+    if (Adjustment + Size <= size_t(End - CurPtr)) {
+      char *AlignedPtr = CurPtr + Adjustment;
       CurPtr = AlignedPtr + Size;
       // Update the allocation point of this memory block in MemorySanitizer.
       // Without this, MemorySanitizer messages for values originated from here
@@ -238,7 +238,7 @@ public:
 
     // Otherwise, start a new slab and try again.
     StartNewSlab();
-    AlignedAddr = alignAddr(CurPtr, Alignment);
+    uintptr_t AlignedAddr = alignAddr(CurPtr, Alignment);
     assert(AlignedAddr + Size <= (uintptr_t)End &&
            "Unable to allocate memory!");
     char *AlignedPtr = (char*)AlignedAddr;
diff --git a/include/llvm/Support/MathExtras.h b/include/llvm/Support/MathExtras.h
index adfca78fbea..6a104e738fa 100644
--- a/include/llvm/Support/MathExtras.h
+++ b/include/llvm/Support/MathExtras.h
@@ -558,9 +558,17 @@ inline uintptr_t alignAddr(void *Addr, size_t Alignment) {
   assert(Alignment && isPowerOf2_64((uint64_t)Alignment) &&
          "Alignment is not a power of two!");
 
+  assert((uintptr_t)Addr + Alignment - 1 >= (uintptr_t)Addr);
+
   return (((uintptr_t)Addr + Alignment - 1) & ~(uintptr_t)(Alignment - 1));
 }
 
+/// \brief Returns the necessary adjustment for aligning \c Ptr to \c Alignment
+/// bytes, rounding up.
+inline size_t alignmentAdjustment(void *Ptr, size_t Alignment) {
+  return alignAddr(Ptr, Alignment) - (uintptr_t)Ptr;
+}
+
 /// NextPowerOf2 - Returns the next power of two (in 64-bits)
 /// that is strictly greater than A.  Returns zero on overflow.
 inline uint64_t NextPowerOf2(uint64_t A) {
diff --git a/unittests/Support/AllocatorTest.cpp b/unittests/Support/AllocatorTest.cpp
index 7789df5dc6f..616f8af3415 100644
--- a/unittests/Support/AllocatorTest.cpp
+++ b/unittests/Support/AllocatorTest.cpp
@@ -115,6 +115,18 @@ TEST(AllocatorTest, TestSmallSlabSize) {
   EXPECT_EQ(1U, Alloc.GetNumSlabs());
 }
 
+// Test requesting alignment that goes past the end of the current slab.
+TEST(AllocatorTest, TestAlignmentPastSlab) {
+  BumpPtrAllocator Alloc;
+  Alloc.Allocate(1234, 1);
+
+  // Any attempt to align the pointer in the current slab would move it beyond
+  // the end of that slab.
+  Alloc.Allocate(1024, 8192);
+
+  EXPECT_EQ(2U, Alloc.GetNumSlabs());
+}
+
 // Mock slab allocator that returns slabs aligned on 4096 bytes.  There is no
 // easy portable way to do this, so this is kind of a hack.
 class MockSlabAllocator {
-- 
2.34.1