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 {