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

[Paul Kirth via libc-commits]

================
@@ -567,3 +569,197 @@ TEST_FOR_EACH_BLOCK_TYPE(CanGetConstBlockFromUsableSpace) {
   const BlockType *block2 = BlockType::from_usable_space(ptr);
   EXPECT_EQ(block1, block2);
 }
+
+TEST_FOR_EACH_BLOCK_TYPE(CanAllocate) {
+  constexpr size_t kN = 1024;
+
+  // Ensure we can allocate everything up to the block size within this block.
+  for (size_t i = 0; i < kN - BlockType::BLOCK_OVERHEAD; ++i) {
+    alignas(BlockType::ALIGNMENT) array<byte, kN> bytes{};
+    auto result = BlockType::init(bytes);
+    ASSERT_TRUE(result.has_value());
+    BlockType *block = *result;
+
+    constexpr size_t kAlign = 1; // Effectively ignores alignment.
+    EXPECT_TRUE(block->can_allocate(kAlign, i));
+
+    // For each can_allocate, we should be able to do a successful call to
+    // allocate.
+    auto info = BlockType::allocate(block, kAlign, i);
+    EXPECT_NE(info.block, static_cast<BlockType *>(nullptr));
+  }
+
+  alignas(BlockType::ALIGNMENT) array<byte, kN> bytes{};
+  auto result = BlockType::init(bytes);
+  ASSERT_TRUE(result.has_value());
+  BlockType *block = *result;
+
+  // Given a block of size kN (assuming it's also a power of two), we should be
+  // able to allocate a block within it that's aligned to half its size. This is
+  // because regardless of where the buffer is located, we can always find a
+  // starting location within it that meets this alignment.
+  EXPECT_TRUE(block->can_allocate(kN / 2, 1));
+  auto info = BlockType::allocate(block, kN / 2, 1);
+  EXPECT_NE(info.block, static_cast<BlockType *>(nullptr));
+}
+
+TEST_FOR_EACH_BLOCK_TYPE(AllocateAlreadyAligned) {
+  constexpr size_t kN = 1024;
+
+  alignas(BlockType::ALIGNMENT) array<byte, kN> bytes{};
+  auto result = BlockType::init(bytes);
+  ASSERT_TRUE(result.has_value());
+  BlockType *block = *result;
+
+  // This should result in no new blocks.
+  constexpr size_t kAlignment = BlockType::ALIGNMENT;
+  constexpr size_t kExpectedSize = BlockType::ALIGNMENT;
+  EXPECT_TRUE(block->can_allocate(kAlignment, kExpectedSize));
+
+  auto [aligned_block, prev, next] =
+      BlockType::allocate(block, BlockType::ALIGNMENT, kExpectedSize);
+
+  // Since this is already aligned, there should be no previous block.
+  EXPECT_EQ(prev, static_cast<BlockType *>(nullptr));
+
+  // Ensure we the block is aligned and the size we expect.
+  EXPECT_NE(aligned_block, static_cast<BlockType *>(nullptr));
+  EXPECT_TRUE(aligned_block->is_usable_space_aligned(BlockType::ALIGNMENT));
+  EXPECT_EQ(aligned_block->inner_size(), kExpectedSize);
+
+  // Check the next block.
+  EXPECT_NE(next, static_cast<BlockType *>(nullptr));
+  EXPECT_EQ(aligned_block->next(), next);
+  EXPECT_EQ(next->next(), static_cast<BlockType *>(nullptr));
+  EXPECT_EQ(reinterpret_cast<byte *>(next) + next->outer_size(), &*bytes.end());
+}
+
+TEST_FOR_EACH_BLOCK_TYPE(AllocateNeedsAlignment) {
+  constexpr size_t kN = 1024;
+
+  alignas(kN) array<byte, kN> bytes{};
+  auto result = BlockType::init(bytes);
+  ASSERT_TRUE(result.has_value());
+  BlockType *block = *result;
+
+  // Ensure first the usable_data is only aligned to the block alignment.
+  ASSERT_EQ(block->usable_space(), bytes.data() + BlockType::BLOCK_OVERHEAD);
+  ASSERT_EQ(block->prev(), static_cast<BlockType *>(nullptr));
+
+  // Now pick an alignment such that the usable space is not already aligned to
+  // it. We want to explicitly test that the block will split into one before
+  // it.
+  constexpr size_t kAlignment = bit_ceil(BlockType::BLOCK_OVERHEAD) * 8;
+  ASSERT_FALSE(block->is_usable_space_aligned(kAlignment));
+
+  constexpr size_t kSize = BlockType::ALIGNMENT;
+  EXPECT_TRUE(block->can_allocate(kAlignment, kSize));
+
+  auto [aligned_block, prev, next] =
+      BlockType::allocate(block, kAlignment, kSize);
----------------
ilovepi wrote:

```suggestion
  EXPECT_TRUE(block->can_allocate(kAlignment,  BlockType::ALIGNMENT));

  auto [aligned_block, prev, next] =
      BlockType::allocate(block, kAlignment,  BlockType::ALIGNMENT);
```
It's not clear why you need the variable here, when you can use the original constant.

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