[libc][stdlib] Add the FreelistHeap (#95066)

This is the actual freelist allocator which utilizes the generic
FreeList and the Block classes. We will eventually wrap the malloc
interface around this.

This is a part of #94270 to land in smaller patches.

4 files changed, 449 insertions, 0 deletions

diff --git a/libc/src/stdlib/CMakeLists.txt b/libc/src/stdlib/CMakeLists.txt
index d4aa50a..971b39b 100644
--- a/libc/src/stdlib/CMakeLists.txt
+++ b/libc/src/stdlib/CMakeLists.txt
@@ -402,6 +402,21 @@ else()
       libc.src.__support.CPP.array
       libc.src.__support.CPP.span
   )
+  add_header_library(
+    freelist_heap
+    HDRS
+      freelist_heap.h
+    DEPENDS
+      .block
+      .freelist
+      libc.src.__support.CPP.cstddef
+      libc.src.__support.CPP.array
+      libc.src.__support.CPP.optional
+      libc.src.__support.CPP.span
+      libc.src.__support.libc_assert
+      libc.src.string.memory_utils.inline_memcpy
+      libc.src.string.memory_utils.inline_memset
+  )
   add_entrypoint_external(
     malloc
   )
diff --git a/libc/src/stdlib/freelist_heap.h b/libc/src/stdlib/freelist_heap.h
new file mode 100644
index 0000000..b65d361
--- /dev/null
+++ b/libc/src/stdlib/freelist_heap.h
@@ -0,0 +1,182 @@
+//===-- Interface for freelist_heap ---------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_STDLIB_FREELIST_HEAP_H
+#define LLVM_LIBC_SRC_STDLIB_FREELIST_HEAP_H
+
+#include <stddef.h>
+
+#include "block.h"
+#include "freelist.h"
+#include "src/__support/CPP/optional.h"
+#include "src/__support/CPP/span.h"
+#include "src/__support/libc_assert.h"
+#include "src/string/memory_utils/inline_memcpy.h"
+#include "src/string/memory_utils/inline_memset.h"
+
+namespace LIBC_NAMESPACE {
+
+using cpp::optional;
+using cpp::span;
+
+static constexpr cpp::array<size_t, 6> DEFAULT_BUCKETS{16,  32,  64,
+                                                       128, 256, 512};
+
+template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
+public:
+  using BlockType = Block<>;
+
+  struct HeapStats {
+    size_t total_bytes;
+    size_t bytes_allocated;
+    size_t cumulative_allocated;
+    size_t cumulative_freed;
+    size_t total_allocate_calls;
+    size_t total_free_calls;
+  };
+  FreeListHeap(span<cpp::byte> region);
+
+  void *allocate(size_t size);
+  void free(void *ptr);
+  void *realloc(void *ptr, size_t size);
+  void *calloc(size_t num, size_t size);
+
+  const HeapStats &heap_stats() const { return heap_stats_; }
+
+private:
+  span<cpp::byte> block_to_span(BlockType *block) {
+    return span<cpp::byte>(block->usable_space(), block->inner_size());
+  }
+
+  span<cpp::byte> region_;
+  FreeList<NUM_BUCKETS> freelist_;
+  HeapStats heap_stats_;
+};
+
+template <size_t NUM_BUCKETS>
+FreeListHeap<NUM_BUCKETS>::FreeListHeap(span<cpp::byte> region)
+    : region_(region), freelist_(DEFAULT_BUCKETS), heap_stats_() {
+  auto result = BlockType::init(region);
+  BlockType *block = *result;
+
+  freelist_.add_chunk(block_to_span(block));
+
+  heap_stats_.total_bytes = region.size();
+}
+
+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);
+
+  if (chunk.data() == nullptr)
+    return nullptr;
+  freelist_.remove_chunk(chunk);
+
+  BlockType *chunk_block = BlockType::from_usable_space(chunk.data());
+
+  // 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));
+
+  chunk_block->mark_used();
+
+  heap_stats_.bytes_allocated += size;
+  heap_stats_.cumulative_allocated += size;
+  heap_stats_.total_allocate_calls += 1;
+
+  return chunk_block->usable_space();
+}
+
+template <size_t NUM_BUCKETS> void FreeListHeap<NUM_BUCKETS>::free(void *ptr) {
+  cpp::byte *bytes = static_cast<cpp::byte *>(ptr);
+
+  LIBC_ASSERT(bytes >= region_.data() && bytes < region_.data() + region_.size() && "Invalid pointer");
+
+  BlockType *chunk_block = BlockType::from_usable_space(bytes);
+
+  size_t size_freed = chunk_block->inner_size();
+  LIBC_ASSERT(chunk_block->used() && "The block is not in-use");
+  chunk_block->mark_free();
+
+  // Can we combine with the left or right blocks?
+  BlockType *prev = chunk_block->prev();
+  BlockType *next = nullptr;
+
+  if (!chunk_block->last())
+    next = chunk_block->next();
+
+  if (prev != nullptr && !prev->used()) {
+    // Remove from freelist and merge
+    freelist_.remove_chunk(block_to_span(prev));
+    chunk_block = chunk_block->prev();
+    BlockType::merge_next(chunk_block);
+  }
+
+  if (next != nullptr && !next->used()) {
+    freelist_.remove_chunk(block_to_span(next));
+    BlockType::merge_next(chunk_block);
+  }
+  // Add back to the freelist
+  freelist_.add_chunk(block_to_span(chunk_block));
+
+  heap_stats_.bytes_allocated -= size_freed;
+  heap_stats_.cumulative_freed += size_freed;
+  heap_stats_.total_free_calls += 1;
+}
+
+// Follows contract of the C standard realloc() function
+// If ptr is free'd, will return nullptr.
+template <size_t NUM_BUCKETS>
+void *FreeListHeap<NUM_BUCKETS>::realloc(void *ptr, size_t size) {
+  if (size == 0) {
+    free(ptr);
+    return nullptr;
+  }
+
+  // If the pointer is nullptr, allocate a new memory.
+  if (ptr == nullptr)
+    return allocate(size);
+
+  cpp::byte *bytes = static_cast<cpp::byte *>(ptr);
+
+  if (bytes < region_.data() || bytes >= region_.data() + region_.size())
+    return nullptr;
+
+  BlockType *chunk_block = BlockType::from_usable_space(bytes);
+  if (!chunk_block->used())
+    return nullptr;
+  size_t old_size = chunk_block->inner_size();
+
+  // Do nothing and return ptr if the required memory size is smaller than
+  // the current size.
+  if (old_size >= size)
+    return ptr;
+
+  void *new_ptr = allocate(size);
+  // Don't invalidate ptr if allocate(size) fails to initilize the memory.
+  if (new_ptr == nullptr)
+    return nullptr;
+  LIBC_NAMESPACE::inline_memcpy(new_ptr, ptr, old_size);
+
+  free(ptr);
+  return new_ptr;
+}
+
+template <size_t NUM_BUCKETS>
+void *FreeListHeap<NUM_BUCKETS>::calloc(size_t num, size_t size) {
+  void *ptr = allocate(num * size);
+  if (ptr != nullptr)
+    LIBC_NAMESPACE::inline_memset(ptr, 0, num * size);
+  return ptr;
+}
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_FREELIST_HEAP_H
diff --git a/libc/test/src/stdlib/CMakeLists.txt b/libc/test/src/stdlib/CMakeLists.txt
index d3954f0..f3033a4 100644
--- a/libc/test/src/stdlib/CMakeLists.txt
+++ b/libc/test/src/stdlib/CMakeLists.txt
@@ -79,6 +79,19 @@ add_libc_test(
     libc.src.__support.CPP.span
 )
 
+add_libc_test(
+  freelist_heap_test
+  SUITE
+    libc-stdlib-tests
+  SRCS
+    freelist_heap_test.cpp
+  DEPENDS
+    libc.src.__support.CPP.span
+    libc.src.stdlib.freelist_heap
+    libc.src.string.memcmp
+    libc.src.string.memcpy
+)
+
 add_fp_unittest(
   strtod_test
   SUITE
diff --git a/libc/test/src/stdlib/freelist_heap_test.cpp b/libc/test/src/stdlib/freelist_heap_test.cpp
new file mode 100644
index 0000000..b89f47f
--- /dev/null
+++ b/libc/test/src/stdlib/freelist_heap_test.cpp
@@ -0,0 +1,239 @@
+//===-- Unittests for freelist_heap ---------------------------------------===//
+//
+// 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 "src/__support/CPP/span.h"
+#include "src/stdlib/freelist_heap.h"
+#include "src/string/memcmp.h"
+#include "src/string/memcpy.h"
+#include "test/UnitTest/Test.h"
+
+namespace LIBC_NAMESPACE {
+
+TEST(LlvmLibcFreeListHeap, CanAllocate) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr = allocator.allocate(ALLOC_SIZE);
+
+  ASSERT_NE(ptr, static_cast<void *>(nullptr));
+  // In this case, the allocator should be returning us the start of the chunk.
+  EXPECT_EQ(ptr, static_cast<void *>(
+                     &buf[0] + FreeListHeap<>::BlockType::BLOCK_OVERHEAD));
+}
+
+TEST(LlvmLibcFreeListHeap, AllocationsDontOverlap) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  void *ptr2 = allocator.allocate(ALLOC_SIZE);
+
+  ASSERT_NE(ptr1, static_cast<void *>(nullptr));
+  ASSERT_NE(ptr2, static_cast<void *>(nullptr));
+
+  uintptr_t ptr1_start = reinterpret_cast<uintptr_t>(ptr1);
+  uintptr_t ptr1_end = ptr1_start + ALLOC_SIZE;
+  uintptr_t ptr2_start = reinterpret_cast<uintptr_t>(ptr2);
+
+  EXPECT_GT(ptr2_start, ptr1_end);
+}
+
+TEST(LlvmLibcFreeListHeap, CanFreeAndRealloc) {
+  // There's not really a nice way to test that free works, apart from to try
+  // and get that value back again.
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  allocator.free(ptr1);
+  void *ptr2 = allocator.allocate(ALLOC_SIZE);
+
+  EXPECT_EQ(ptr1, ptr2);
+}
+
+TEST(LlvmLibcFreeListHeap, ReturnsNullWhenAllocationTooLarge) {
+  constexpr size_t N = 2048;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  EXPECT_EQ(allocator.allocate(N), static_cast<void *>(nullptr));
+}
+
+TEST(LlvmLibcFreeListHeap, ReturnsNullWhenFull) {
+  constexpr size_t N = 2048;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  EXPECT_NE(allocator.allocate(N - FreeListHeap<>::BlockType::BLOCK_OVERHEAD),
+            static_cast<void *>(nullptr));
+  EXPECT_EQ(allocator.allocate(1), static_cast<void *>(nullptr));
+}
+
+TEST(LlvmLibcFreeListHeap, ReturnedPointersAreAligned) {
+  constexpr size_t N = 2048;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(1);
+
+  // Should be aligned to native pointer alignment
+  uintptr_t ptr1_start = reinterpret_cast<uintptr_t>(ptr1);
+  size_t alignment = alignof(void *);
+
+  EXPECT_EQ(ptr1_start % alignment, static_cast<size_t>(0));
+
+  void *ptr2 = allocator.allocate(1);
+  uintptr_t ptr2_start = reinterpret_cast<uintptr_t>(ptr2);
+
+  EXPECT_EQ(ptr2_start % alignment, static_cast<size_t>(0));
+}
+
+TEST(LlvmLibcFreeListHeap, CanRealloc) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  constexpr size_t kNewAllocSize = 768;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(1)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  void *ptr2 = allocator.realloc(ptr1, kNewAllocSize);
+
+  ASSERT_NE(ptr1, static_cast<void *>(nullptr));
+  ASSERT_NE(ptr2, static_cast<void *>(nullptr));
+}
+
+TEST(LlvmLibcFreeListHeap, ReallocHasSameContent) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = sizeof(int);
+  constexpr size_t kNewAllocSize = sizeof(int) * 2;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(1)};
+  // Data inside the allocated block.
+  cpp::byte data1[ALLOC_SIZE];
+  // Data inside the reallocated block.
+  cpp::byte data2[ALLOC_SIZE];
+
+  FreeListHeap<> allocator(buf);
+
+  int *ptr1 = reinterpret_cast<int *>(allocator.allocate(ALLOC_SIZE));
+  *ptr1 = 42;
+  memcpy(data1, ptr1, ALLOC_SIZE);
+  int *ptr2 = reinterpret_cast<int *>(allocator.realloc(ptr1, kNewAllocSize));
+  memcpy(data2, ptr2, ALLOC_SIZE);
+
+  ASSERT_NE(ptr1, static_cast<int *>(nullptr));
+  ASSERT_NE(ptr2, static_cast<int *>(nullptr));
+  // Verify that data inside the allocated and reallocated chunks are the same.
+  EXPECT_EQ(LIBC_NAMESPACE::memcmp(data1, data2, ALLOC_SIZE), 0);
+}
+
+TEST(LlvmLibcFreeListHeap, ReturnsNullReallocFreedPointer) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  constexpr size_t kNewAllocSize = 256;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  allocator.free(ptr1);
+  void *ptr2 = allocator.realloc(ptr1, kNewAllocSize);
+
+  EXPECT_EQ(static_cast<void *>(nullptr), ptr2);
+}
+
+TEST(LlvmLibcFreeListHeap, ReallocSmallerSize) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  constexpr size_t kNewAllocSize = 256;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  void *ptr2 = allocator.realloc(ptr1, kNewAllocSize);
+
+  // For smaller sizes, realloc will not shrink the block.
+  EXPECT_EQ(ptr1, ptr2);
+}
+
+TEST(LlvmLibcFreeListHeap, ReallocTooLarge) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 512;
+  constexpr size_t kNewAllocSize = 4096;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
+
+  FreeListHeap<> allocator(buf);
+
+  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  void *ptr2 = allocator.realloc(ptr1, kNewAllocSize);
+
+  // realloc() will not invalidate the original pointer if realloc() fails
+  EXPECT_NE(static_cast<void *>(nullptr), ptr1);
+  EXPECT_EQ(static_cast<void *>(nullptr), ptr2);
+}
+
+TEST(LlvmLibcFreeListHeap, CanCalloc) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 128;
+  constexpr size_t kNum = 4;
+  constexpr int size = kNum * ALLOC_SIZE;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(1)};
+  constexpr cpp::byte zero{0};
+
+  FreeListHeap<> allocator(buf);
+
+  cpp::byte *ptr1 =
+      reinterpret_cast<cpp::byte *>(allocator.calloc(kNum, ALLOC_SIZE));
+
+  // calloc'd content is zero.
+  for (int i = 0; i < size; i++)
+    EXPECT_EQ(ptr1[i], zero);
+}
+
+TEST(LlvmLibcFreeListHeap, CanCallocWeirdSize) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 143;
+  constexpr size_t kNum = 3;
+  constexpr int size = kNum * ALLOC_SIZE;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(132)};
+  constexpr cpp::byte zero{0};
+
+  FreeListHeap<> allocator(buf);
+
+  cpp::byte *ptr1 =
+      reinterpret_cast<cpp::byte *>(allocator.calloc(kNum, ALLOC_SIZE));
+
+  // calloc'd content is zero.
+  for (int i = 0; i < size; i++)
+    EXPECT_EQ(ptr1[i], zero);
+}
+
+TEST(LlvmLibcFreeListHeap, CallocTooLarge) {
+  constexpr size_t N = 2048;
+  constexpr size_t ALLOC_SIZE = 2049;
+  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(1)};
+
+  FreeListHeap<> allocator(buf);
+
+  EXPECT_EQ(allocator.calloc(1, ALLOC_SIZE), static_cast<void *>(nullptr));
+}
+
+} // namespace LIBC_NAMESPACE