[libc] Add aligned_alloc

libc/src/__support/block.h

@@ -16,6 +16,7 @@
 #include "src/__support/CPP/optional.h"
 #include "src/__support/CPP/span.h"
 #include "src/__support/CPP/type_traits.h"
+#include "src/__support/libc_assert.h"
 
 #include <stdint.h>
 
@@ -261,6 +262,63 @@ class Block {
 
   constexpr Block(size_t prev_outer_size, size_t outer_size);
 
+  bool is_usable_space_aligned(size_t alignment) const {
+    return reinterpret_cast<uintptr_t>(usable_space()) % alignment == 0;
+  }
+
+  size_t padding_for_alignment(size_t alignment) const {
+    if (is_usable_space_aligned(alignment))
+      return 0;
+
+    // We need to ensure we can always split this block into a "padding" block
+    // and the aligned block. To do this, we need enough extra space for at
+    // least one block.
+    //
+    // |block   |usable_space                          |
+    // |........|......................................|
+    //                            ^
+    //                            Alignment requirement
+    //
+    //
+    // |block   |space   |block   |usable_space        |
+    // |........|........|........|....................|
+    //                            ^
+    //                            Alignment requirement
+    //
+    uintptr_t start = reinterpret_cast<uintptr_t>(usable_space());
+    alignment = cpp::max(alignment, ALIGNMENT);
+    return align_up(start + BLOCK_OVERHEAD, alignment) - start;
+  }
+
+  // Check that we can `allocate` a block with a given alignment and size from
+  // this existing block.
+  bool can_allocate(size_t alignment, size_t size) const;
+
+  // This is the return type for `allocate` which can split one block into up to
+  // three blocks.
+  struct BlockInfo {
+    // This is the newly aligned block. It will have the alignment requested by
+    // a call to `allocate` and at most `size`.
+    Block *block;
+
+    // If the usable_space in the new block was not aligned according to the
+    // `alignment` parameter, we will need to split into this block and the
+    // `block` to ensure `block` is properly aligned. In this case, `prev` will
+    // be a pointer to this new "padding" block. `prev` will be nullptr if no
+    // new block was created or we were able to merge the block before the
+    // original block with the "padding" block.
+    Block *prev;
+
+    // This is the remainder of the next block after splitting the `block`
+    // according to `size`. This can happen if there's enough space after the
+    // `block`.
+    Block *next;
+  };
+
+  // Divide a block into up to 3 blocks according to `BlockInfo`. This should
+  // only be called if `can_allocate` returns true.
+  static BlockInfo allocate(Block *block, size_t alignment, size_t size);
+
 private:
   /// Consumes the block and returns as a span of bytes.
   static ByteSpan as_bytes(Block *&&block);
@@ -357,6 +415,69 @@ void Block<OffsetType, kAlign>::free(Block *&block) {
   merge_next(block);
 }
 
+template <typename OffsetType, size_t kAlign>
+bool Block<OffsetType, kAlign>::can_allocate(size_t alignment,
+                                             size_t size) const {
+  if (is_usable_space_aligned(alignment) && inner_size() >= size)
+    return true; // Size and alignment constraints met.
+
+  // Either the alignment isn't met or we don't have enough size.
+  // If we don't meet alignment, we can always adjust such that we do meet the
+  // alignment. If we meet the alignment but just don't have enough size. This
+  // check will fail anyway.
+  size_t adjustment = padding_for_alignment(alignment);
+  return inner_size() >= size + adjustment;
+}
+
+template <typename OffsetType, size_t kAlign>
+typename Block<OffsetType, kAlign>::BlockInfo
+Block<OffsetType, kAlign>::allocate(Block *block, size_t alignment,
+                                    size_t size) {
+  LIBC_ASSERT(
+      block->can_allocate(alignment, size) &&
+      "Calls to this function for a given alignment and size should only be "
+      "done if `can_allocate` for these parameters returns true.");
+
+  BlockInfo info{block, /*prev=*/nullptr, /*next=*/nullptr};
+
+  if (!info.block->is_usable_space_aligned(alignment)) {
+    size_t adjustment = info.block->padding_for_alignment(alignment);
+    size_t new_inner_size = adjustment - BLOCK_OVERHEAD;
+    LIBC_ASSERT(new_inner_size % ALIGNMENT == 0 &&
+                "The adjustment calculation should always return a new size "
+                "that's a multiple of ALIGNMENT");
+
+    Block *original = info.block;
+    optional<Block *> maybe_aligned_block =
+        Block::split(original, adjustment - BLOCK_OVERHEAD);
+    LIBC_ASSERT(maybe_aligned_block.has_value() &&
+                "This split should always result in a new block. The check in "
+                "`can_allocate` ensures that we have enough space here to make "
+                "two blocks.");
+
+    if (Block *prev = original->prev()) {
+      // If there is a block before this, we can merge the current one with the
+      // newly created one.
+      merge_next(prev);
+    } else {
+      // Otherwise, this was the very first block in the chain. Now we can make
+      // it the new first block.
+      info.prev = original;
+    }
+
+    Block *aligned_block = *maybe_aligned_block;
+    LIBC_ASSERT(aligned_block->is_usable_space_aligned(alignment) &&
+                "The aligned block isn't aligned somehow.");
+    info.block = aligned_block;
+  }
+
+  // Now get a block for the requested size.
+  if (optional<Block *> next = Block::split(info.block, size))
+    info.next = *next;
+
+  return info;
+}
+
 template <typename OffsetType, size_t kAlign>
 optional<Block<OffsetType, kAlign> *>
 Block<OffsetType, kAlign>::split(Block *&block, size_t new_inner_size) {

libc/src/__support/freelist.h

@@ -66,6 +66,8 @@ template <size_t NUM_BUCKETS = 6> class FreeList {
   ///   A span with a size of 0.
   cpp::span<cpp::byte> find_chunk(size_t size) const;
 
+  template <typename Cond> cpp::span<cpp::byte> find_chunk_if(Cond op) const;
+
   /// Removes a chunk from this freelist.
   bool remove_chunk(cpp::span<cpp::byte> chunk);
 
@@ -111,6 +113,22 @@ bool FreeList<NUM_BUCKETS>::add_chunk(span<cpp::byte> chunk) {
   return true;
 }
 
+template <size_t NUM_BUCKETS>
+template <typename Cond>
+span<cpp::byte> FreeList<NUM_BUCKETS>::find_chunk_if(Cond op) const {
+  for (FreeListNode *node : chunks_) {
+    while (node != nullptr) {
+      span<cpp::byte> chunk(reinterpret_cast<cpp::byte *>(node), node->size);
+      if (op(chunk))
+        return chunk;
+
+      node = node->next;
+    }
+  }
+
+  return {};
+}
+
 template <size_t NUM_BUCKETS>
 span<cpp::byte> FreeList<NUM_BUCKETS>::find_chunk(size_t size) const {
   if (size == 0)

libc/src/__support/freelist_heap.h

@@ -24,6 +24,8 @@ namespace LIBC_NAMESPACE {
 using cpp::optional;
 using cpp::span;
 
+inline constexpr bool IsPow2(size_t x) { return x && (x & (x - 1)) == 0; }
+
 static constexpr cpp::array<size_t, 6> DEFAULT_BUCKETS{16,  32,  64,
                                                        128, 256, 512};
 
@@ -32,6 +34,9 @@ template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
   using BlockType = Block<>;
   using FreeListType = FreeList<NUM_BUCKETS>;
 
+  static constexpr size_t MIN_ALIGNMENT =
+      cpp::max(BlockType::ALIGNMENT, alignof(max_align_t));
+
   struct HeapStats {
     size_t total_bytes;
     size_t bytes_allocated;
@@ -55,6 +60,9 @@ template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
   }
 
   void *allocate(size_t size);
+  void *aligned_allocate(size_t alignment, size_t size);
+  // NOTE: All pointers passed to free must come from one of the other
+  // allocation functions: `allocate`, `aligned_allocate`, `realloc`, `calloc`.
   void free(void *ptr);
   void *realloc(void *ptr, size_t size);
   void *calloc(size_t num, size_t size);
@@ -74,6 +82,8 @@ template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
     freelist_.set_freelist_node(node, chunk);
   }
 
+  void *allocate_impl(size_t alignment, size_t size);
+
 private:
   span<cpp::byte> block_to_span(BlockType *block) {
     return span<cpp::byte>(block->usable_space(), block->inner_size());
@@ -109,20 +119,31 @@ struct FreeListHeapBuffer : public FreeListHeap<NUM_BUCKETS> {
 };
 
 template <size_t NUM_BUCKETS>
-void *FreeListHeap<NUM_BUCKETS>::allocate(size_t size) {
-  // Find a chunk in the freelist. Split it if needed, then return
-  auto chunk = freelist_.find_chunk(size);
+void *FreeListHeap<NUM_BUCKETS>::allocate_impl(size_t alignment, size_t size) {
+  if (size == 0)
+    return nullptr;
+
+  // Find a chunk in the freelist. Split it if needed, then return.
+  auto chunk =
+      freelist_.find_chunk_if([alignment, size](span<cpp::byte> chunk) {
+        BlockType *block = BlockType::from_usable_space(chunk.data());
+        return block->can_allocate(alignment, size);
+      });
 
   if (chunk.data() == nullptr)
     return nullptr;
   freelist_.remove_chunk(chunk);
 
   BlockType *chunk_block = BlockType::from_usable_space(chunk.data());
+  LIBC_ASSERT(!chunk_block->used());
 
   // Split that chunk. If there's a leftover chunk, add it to the freelist
-  optional<BlockType *> result = BlockType::split(chunk_block, size);
-  if (result)
-    freelist_.add_chunk(block_to_span(*result));
+  auto block_info = BlockType::allocate(chunk_block, alignment, size);
+  if (block_info.next)
+    freelist_.add_chunk(block_to_span(block_info.next));
+  if (block_info.prev)
+    freelist_.add_chunk(block_to_span(block_info.prev));
+  chunk_block = block_info.block;
 
   chunk_block->mark_used();
 
@@ -133,6 +154,25 @@ void *FreeListHeap<NUM_BUCKETS>::allocate(size_t size) {
   return chunk_block->usable_space();
 }
 
+template <size_t NUM_BUCKETS>
+void *FreeListHeap<NUM_BUCKETS>::allocate(size_t size) {
+  return allocate_impl(MIN_ALIGNMENT, size);
+}
+
+template <size_t NUM_BUCKETS>
+void *FreeListHeap<NUM_BUCKETS>::aligned_allocate(size_t alignment,
+                                                  size_t size) {
+  // The alignment must be an integral power of two.
+  if (!IsPow2(alignment))
+    return nullptr;
+
+  // The size parameter must be an integral multiple of alignment.
+  if (size % alignment != 0)
+    return nullptr;
+
+  return allocate_impl(alignment, size);
+}
+
 template <size_t NUM_BUCKETS> void FreeListHeap<NUM_BUCKETS>::free(void *ptr) {
   cpp::byte *bytes = static_cast<cpp::byte *>(ptr);
 

libc/src/stdlib/aligned_alloc.h

@@ -0,0 +1,20 @@
+//===-- Implementation header for aligned_alloc -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <stddef.h>
+
+#ifndef LLVM_LIBC_SRC_STDLIB_ALIGNED_ALLOC_H
+#define LLVM_LIBC_SRC_STDLIB_ALIGNED_ALLOC_H
+
+namespace LIBC_NAMESPACE {
+
+void *aligned_alloc(size_t alignment, size_t size);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_ALIGNED_ALLOC_H

libc/src/stdlib/freelist_malloc.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/__support/freelist_heap.h"
+#include "src/stdlib/aligned_alloc.h"
 #include "src/stdlib/calloc.h"
 #include "src/stdlib/free.h"
 #include "src/stdlib/malloc.h"
@@ -42,4 +43,8 @@ LLVM_LIBC_FUNCTION(void *, realloc, (void *ptr, size_t size)) {
   return freelist_heap->realloc(ptr, size);
 }
 
+LLVM_LIBC_FUNCTION(void *, aligned_alloc, (size_t alignment, size_t size)) {
+  return freelist_heap->aligned_allocate(alignment, size);
+}
+
 } // namespace LIBC_NAMESPACE

libc/test/src/__support/CMakeLists.txt

@@ -8,6 +8,7 @@ add_libc_test(
     block_test.cpp
   DEPENDS
     libc.src.__support.CPP.array
+    libc.src.__support.CPP.bit
     libc.src.__support.CPP.span
     libc.src.__support.block
     libc.src.string.memcpy