[libc-commits] [libc] [libc] Add aligned_alloc (PR #96586)

[Nick Desaulniers via libc-commits]

================
@@ -214,4 +224,70 @@ TEST_FOR_EACH_ALLOCATOR(CallocTooLarge, 2048) {
   EXPECT_EQ(allocator.calloc(1, ALLOC_SIZE), static_cast<void *>(nullptr));
 }
 
+TEST_FOR_EACH_ALLOCATOR(AlignedAlloc, 2048) {
+  constexpr size_t ALIGNMENTS[] = {1, 2, 4, 8, 16, 32, 64, 128, 256};
+  constexpr size_t SIZE_SCALES[] = {1, 2, 3, 4, 5};
+
+  for (size_t alignment : ALIGNMENTS) {
+    for (size_t scale : SIZE_SCALES) {
+      size_t size = alignment * scale;
+      void *ptr = allocator.aligned_allocate(alignment, size);
+      EXPECT_NE(ptr, static_cast<void *>(nullptr));
+      EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % alignment, size_t(0));
+      allocator.free(ptr);
+    }
+  }
+}
+
+// This test is not part of the TEST_FOR_EACH_ALLOCATOR since we want to
+// explicitly ensure that the buffer can still return aligned allocations even
+// if the underlying buffer is at most aligned to the BlockType alignment. This
+// is so we can check that we can still get aligned allocations even if the
+// underlying buffer is not aligned to the alignments we request.
+TEST(LlvmLibcFreeListHeap, AlignedAllocOnlyBlockTypeAligned) {
+  constexpr size_t BUFFER_SIZE = 4096;
+  constexpr size_t BUFFER_ALIGNMENT = alignof(FreeListHeap<>::BlockType) * 2;
+  alignas(BUFFER_ALIGNMENT) cpp::byte buf[BUFFER_SIZE] = {cpp::byte(0)};
+
+  // Ensure the underlying buffer is at most aligned to the block type.
+  FreeListHeap<> allocator(
+      span<cpp::byte>(buf).subspan(alignof(FreeListHeap<>::BlockType)));
+
+  constexpr size_t ALIGNMENTS[] = {1, 2, 4, 8, 16, 32, 64, 128, 256};
+  constexpr size_t SIZE_SCALES[] = {1, 2, 3, 4, 5};
+
+  for (size_t alignment : ALIGNMENTS) {
+    for (size_t scale : SIZE_SCALES) {
+      size_t size = alignment * scale;
+      void *ptr = allocator.aligned_allocate(alignment, size);
+      EXPECT_NE(ptr, static_cast<void *>(nullptr));
+      EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % alignment, size_t(0));
+      allocator.free(ptr);
+    }
+  }
+}
+
+TEST_FOR_EACH_ALLOCATOR(InvalidAlignedAllocAlignment, 2048) {
+  // Must be a power of 2.
+  constexpr size_t ALIGNMENTS[] = {4, 8, 16, 32, 64, 128, 256};
+  for (size_t alignment : ALIGNMENTS) {
+    void *ptr = allocator.aligned_allocate(alignment - 1, alignment - 1);
+    EXPECT_EQ(ptr, static_cast<void *>(nullptr));
----------------
nickdesaulniers wrote:

Should you free ptr? Here and below.

https://github.com/llvm/llvm-project/pull/96586