[libc-commits] [libc] [WIP][libc] Add freelist malloc (PR #94270)

via libc-commits libc-commits at lists.llvm.org
Thu Jun 13 11:56:23 PDT 2024


https://github.com/PiJoules updated https://github.com/llvm/llvm-project/pull/94270

>From 843ea00e03dda3cab7615e968d904c49176ff324 Mon Sep 17 00:00:00 2001
From: Leonard Chan <leonardchan at google.com>
Date: Fri, 31 May 2024 14:27:16 -0700
Subject: [PATCH] [WIP][libc] Add freelist malloc

---
 libc/config/baremetal/riscv/entrypoints.txt   |   1 +
 libc/src/__support/threads/thread.cpp         |   3 +-
 libc/src/stdlib/CMakeLists.txt                |   8 +-
 libc/src/stdlib/block.h                       |  10 +-
 libc/src/stdlib/freelist.h                    |  41 ++--
 libc/src/stdlib/freelist_heap.h               |  69 +++++--
 libc/src/stdlib/freelist_malloc.cpp           |  36 ++++
 libc/test/src/CMakeLists.txt                  |   2 +-
 libc/test/src/stdlib/CMakeLists.txt           |  33 ++--
 libc/test/src/stdlib/freelist_heap_test.cpp   | 186 ++++++++----------
 libc/test/src/stdlib/freelist_malloc_test.cpp |  56 ++++++
 11 files changed, 290 insertions(+), 155 deletions(-)
 create mode 100644 libc/src/stdlib/freelist_malloc.cpp
 create mode 100644 libc/test/src/stdlib/freelist_malloc_test.cpp

diff --git a/libc/config/baremetal/riscv/entrypoints.txt b/libc/config/baremetal/riscv/entrypoints.txt
index b769b43f03a2c..363a762909c3a 100644
--- a/libc/config/baremetal/riscv/entrypoints.txt
+++ b/libc/config/baremetal/riscv/entrypoints.txt
@@ -170,6 +170,7 @@ set(TARGET_LIBC_ENTRYPOINTS
     libc.src.stdlib.ldiv
     libc.src.stdlib.llabs
     libc.src.stdlib.lldiv
+    libc.src.stdlib.malloc
     libc.src.stdlib.qsort
     libc.src.stdlib.rand
     libc.src.stdlib.srand
diff --git a/libc/src/__support/threads/thread.cpp b/libc/src/__support/threads/thread.cpp
index d7bedb8ca67d0..d39d1ed431d04 100644
--- a/libc/src/__support/threads/thread.cpp
+++ b/libc/src/__support/threads/thread.cpp
@@ -115,7 +115,8 @@ class ThreadAtExitCallbackMgr {
 public:
   constexpr ThreadAtExitCallbackMgr()
       : mtx(/*timed=*/false, /*recursive=*/false, /*robust=*/false,
-            /*pshared=*/false) {}
+            /*pshared=*/false),
+        callback_list() {}
 
   int add_callback(AtExitCallback *callback, void *obj) {
     cpp::lock_guard lock(mtx);
diff --git a/libc/src/stdlib/CMakeLists.txt b/libc/src/stdlib/CMakeLists.txt
index 971b39bb900de..664de616319eb 100644
--- a/libc/src/stdlib/CMakeLists.txt
+++ b/libc/src/stdlib/CMakeLists.txt
@@ -417,8 +417,14 @@ else()
       libc.src.string.memory_utils.inline_memcpy
       libc.src.string.memory_utils.inline_memset
   )
-  add_entrypoint_external(
+  add_entrypoint_object(
     malloc
+    SRCS
+      freelist_malloc.cpp
+    HDRS
+      malloc.h
+    DEPENDS
+      .freelist_heap
   )
   add_entrypoint_external(
     free
diff --git a/libc/src/stdlib/block.h b/libc/src/stdlib/block.h
index afb18c1ef738f..b0462a12afb39 100644
--- a/libc/src/stdlib/block.h
+++ b/libc/src/stdlib/block.h
@@ -245,7 +245,7 @@ class Block {
   void mark_free() { info_.used = 0; }
 
   /// Marks this block as the last one in the chain.
-  void mark_last() { info_.last = 1; }
+  constexpr void mark_last() { info_.last = 1; }
 
   /// Clears the last bit from this block.
   void clear_last() { info_.last = 1; }
@@ -259,6 +259,8 @@ class Block {
     return check_status() == internal::BlockStatus::VALID;
   }
 
+  constexpr Block(size_t prev_outer_size, size_t outer_size);
+
 private:
   /// Consumes the block and returns as a span of bytes.
   static ByteSpan as_bytes(Block *&&block);
@@ -266,8 +268,6 @@ class Block {
   /// Consumes the span of bytes and uses it to construct and return a block.
   static Block *as_block(size_t prev_outer_size, ByteSpan bytes);
 
-  Block(size_t prev_outer_size, size_t outer_size);
-
   /// Returns a `BlockStatus` that is either VALID or indicates the reason why
   /// the block is invalid.
   ///
@@ -442,7 +442,9 @@ Block<OffsetType, kAlign> *Block<OffsetType, kAlign>::prev() const {
 // Private template method implementations.
 
 template <typename OffsetType, size_t kAlign>
-Block<OffsetType, kAlign>::Block(size_t prev_outer_size, size_t outer_size) {
+constexpr Block<OffsetType, kAlign>::Block(size_t prev_outer_size,
+                                           size_t outer_size)
+    : info_{} {
   prev_ = prev_outer_size / ALIGNMENT;
   next_ = outer_size / ALIGNMENT;
   info_.used = 0;
diff --git a/libc/src/stdlib/freelist.h b/libc/src/stdlib/freelist.h
index c01ed6eddb7d4..54e08b4343578 100644
--- a/libc/src/stdlib/freelist.h
+++ b/libc/src/stdlib/freelist.h
@@ -68,43 +68,53 @@ template <size_t NUM_BUCKETS = 6> class FreeList {
   /// Removes a chunk from this freelist.
   bool remove_chunk(cpp::span<cpp::byte> chunk);
 
-private:
-  // For a given size, find which index into chunks_ the node should be written
-  // to.
-  size_t find_chunk_ptr_for_size(size_t size, bool non_null) const;
+  /// For a given size, find which index into chunks_ the node should be written
+  /// to.
+  constexpr size_t find_chunk_ptr_for_size(size_t size, bool non_null) const;
 
   struct FreeListNode {
     FreeListNode *next;
     size_t size;
   };
 
-public:
-  explicit FreeList(cpp::array<size_t, NUM_BUCKETS> sizes)
+  constexpr void set_freelist_node(FreeListNode *node,
+                                   cpp::span<cpp::byte> chunk);
+
+  constexpr explicit FreeList(cpp::array<size_t, NUM_BUCKETS> sizes)
       : chunks_(NUM_BUCKETS + 1, 0), sizes_(sizes.begin(), sizes.end()) {}
 
+private:
   FixedVector<FreeList::FreeListNode *, NUM_BUCKETS + 1> chunks_;
   FixedVector<size_t, NUM_BUCKETS> sizes_;
 };
 
+template <size_t NUM_BUCKETS>
+constexpr void FreeList<NUM_BUCKETS>::set_freelist_node(FreeListNode *node,
+                                                        span<cpp::byte> chunk) {
+  // Add it to the correct list.
+  size_t chunk_ptr = find_chunk_ptr_for_size(chunk.size(), false);
+  node->size = chunk.size();
+  node->next = chunks_[chunk_ptr];
+  chunks_[chunk_ptr] = node;
+}
+
 template <size_t NUM_BUCKETS>
 bool FreeList<NUM_BUCKETS>::add_chunk(span<cpp::byte> chunk) {
   // Check that the size is enough to actually store what we need
   if (chunk.size() < sizeof(FreeListNode))
     return false;
 
+  // FIXME: This is UB since type punning is not allowed in C++. THe idea here
+  // is that we write the FreeListNode `size` and `next` onto the start of the
+  // buffer. Unless the original underlying bytes were a `FreeListNode`, the
+  // only safe way to do this is with `memcpy`.
   union {
     FreeListNode *node;
     cpp::byte *bytes;
   } aliased;
 
   aliased.bytes = chunk.data();
-
-  size_t chunk_ptr = find_chunk_ptr_for_size(chunk.size(), false);
-
-  // Add it to the correct list.
-  aliased.node->size = chunk.size();
-  aliased.node->next = chunks_[chunk_ptr];
-  chunks_[chunk_ptr] = aliased.node;
+  set_freelist_node(aliased.node, chunk);
 
   return true;
 }
@@ -180,8 +190,9 @@ bool FreeList<NUM_BUCKETS>::remove_chunk(span<cpp::byte> chunk) {
 }
 
 template <size_t NUM_BUCKETS>
-size_t FreeList<NUM_BUCKETS>::find_chunk_ptr_for_size(size_t size,
-                                                      bool non_null) const {
+constexpr size_t
+FreeList<NUM_BUCKETS>::find_chunk_ptr_for_size(size_t size,
+                                               bool non_null) const {
   size_t chunk_ptr = 0;
   for (chunk_ptr = 0u; chunk_ptr < sizes_.size(); chunk_ptr++) {
     if (sizes_[chunk_ptr] >= size &&
diff --git a/libc/src/stdlib/freelist_heap.h b/libc/src/stdlib/freelist_heap.h
index b65d361e9ca73..f9f1ce4e723f1 100644
--- a/libc/src/stdlib/freelist_heap.h
+++ b/libc/src/stdlib/freelist_heap.h
@@ -30,6 +30,7 @@ static constexpr cpp::array<size_t, 6> DEFAULT_BUCKETS{16,  32,  64,
 template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
 public:
   using BlockType = Block<>;
+  using FreeListType = FreeList<NUM_BUCKETS>;
 
   struct HeapStats {
     size_t total_bytes;
@@ -39,7 +40,19 @@ template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
     size_t total_allocate_calls;
     size_t total_free_calls;
   };
-  FreeListHeap(span<cpp::byte> region);
+
+  FreeListHeap(span<cpp::byte> region)
+      : FreeListHeap(&*region.begin(), &*region.end(), region.size()) {
+    auto result = BlockType::init(region);
+    BlockType *block = *result;
+    freelist_.add_chunk(block_to_span(block));
+  }
+
+  constexpr FreeListHeap(void *start, cpp::byte *end, size_t total_bytes)
+      : block_region_start_(start), block_region_end_(end),
+        freelist_(DEFAULT_BUCKETS), heap_stats_{} {
+    heap_stats_.total_bytes = total_bytes;
+  }
 
   void *allocate(size_t size);
   void free(void *ptr);
@@ -47,27 +60,53 @@ template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
   void *calloc(size_t num, size_t size);
 
   const HeapStats &heap_stats() const { return heap_stats_; }
+  void reset_heap_stats() { heap_stats_ = {}; }
+
+  void *region_start() const { return block_region_start_; }
+  size_t region_size() const {
+    return reinterpret_cast<uintptr_t>(block_region_end_) -
+           reinterpret_cast<uintptr_t>(block_region_start_);
+  }
+
+protected:
+  constexpr void set_freelist_node(typename FreeListType::FreeListNode *node,
+                                   cpp::span<cpp::byte> chunk) {
+    freelist_.set_freelist_node(node, chunk);
+  }
 
 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_;
+  bool is_valid_ptr(void *ptr) {
+    return ptr >= block_region_start_ && ptr < block_region_end_;
+  }
+
+  void *block_region_start_;
+  void *block_region_end_;
+  FreeListType 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;
+template <size_t BUFF_SIZE, size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()>
+struct FreeListHeapBuffer : public FreeListHeap<NUM_BUCKETS> {
+  using parent = FreeListHeap<NUM_BUCKETS>;
+  using FreeListNode = typename parent::FreeListType::FreeListNode;
 
-  freelist_.add_chunk(block_to_span(block));
+  constexpr FreeListHeapBuffer()
+      : FreeListHeap<NUM_BUCKETS>(&block, buffer + sizeof(buffer), BUFF_SIZE),
+        block(0, BUFF_SIZE), node{}, buffer{} {
+    block.mark_last();
 
-  heap_stats_.total_bytes = region.size();
-}
+    cpp::span<cpp::byte> chunk(buffer, sizeof(buffer));
+    parent::set_freelist_node(&node, chunk);
+  }
+
+  typename parent::BlockType block;
+  FreeListNode node;
+  cpp::byte buffer[BUFF_SIZE - sizeof(block) - sizeof(node)];
+};
 
 template <size_t NUM_BUCKETS>
 void *FreeListHeap<NUM_BUCKETS>::allocate(size_t size) {
@@ -97,7 +136,7 @@ void *FreeListHeap<NUM_BUCKETS>::allocate(size_t size) {
 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");
+  LIBC_ASSERT(is_valid_ptr(bytes) && "Invalid pointer");
 
   BlockType *chunk_block = BlockType::from_usable_space(bytes);
 
@@ -131,7 +170,7 @@ template <size_t NUM_BUCKETS> void FreeListHeap<NUM_BUCKETS>::free(void *ptr) {
   heap_stats_.total_free_calls += 1;
 }
 
-// Follows contract of the C standard realloc() function
+// Follows constract 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) {
@@ -146,7 +185,7 @@ void *FreeListHeap<NUM_BUCKETS>::realloc(void *ptr, size_t size) {
 
   cpp::byte *bytes = static_cast<cpp::byte *>(ptr);
 
-  if (bytes < region_.data() || bytes >= region_.data() + region_.size())
+  if (!is_valid_ptr(bytes))
     return nullptr;
 
   BlockType *chunk_block = BlockType::from_usable_space(bytes);
@@ -177,6 +216,8 @@ void *FreeListHeap<NUM_BUCKETS>::calloc(size_t num, size_t size) {
   return ptr;
 }
 
+extern FreeListHeap<> *freelist_heap;
+
 } // namespace LIBC_NAMESPACE
 
 #endif // LLVM_LIBC_SRC_STDLIB_FREELIST_HEAP_H
diff --git a/libc/src/stdlib/freelist_malloc.cpp b/libc/src/stdlib/freelist_malloc.cpp
new file mode 100644
index 0000000000000..e6e2e7f808413
--- /dev/null
+++ b/libc/src/stdlib/freelist_malloc.cpp
@@ -0,0 +1,36 @@
+//===-- Implementation for freelist_malloc --------------------------------===//
+//
+// 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 "freelist_heap.h"
+
+#include <stddef.h>
+
+namespace LIBC_NAMESPACE {
+
+namespace {
+// TODO: We should probably have something akin to what scudo/sanitizer
+// allocators do where each platform defines this.
+constexpr size_t SIZE = 0x40000000ULL; // 1GB
+LIBC_CONSTINIT FreeListHeapBuffer<SIZE> freelist_heap_buffer;
+} // namespace
+
+FreeListHeap<> *freelist_heap = &freelist_heap_buffer;
+
+void *malloc(size_t size) { return freelist_heap->allocate(size); }
+
+void free(void *ptr) { freelist_heap->free(ptr); }
+
+void *calloc(size_t num, size_t size) {
+  return freelist_heap->calloc(num, size);
+}
+
+void *realloc(void *ptr, size_t size) {
+  return freelist_heap->realloc(ptr, size);
+}
+
+} // namespace LIBC_NAMESPACE
diff --git a/libc/test/src/CMakeLists.txt b/libc/test/src/CMakeLists.txt
index a5e7a2a4dee72..935feb59ecdf6 100644
--- a/libc/test/src/CMakeLists.txt
+++ b/libc/test/src/CMakeLists.txt
@@ -61,7 +61,7 @@ add_subdirectory(inttypes)
 if(${LIBC_TARGET_OS} STREQUAL "linux")
   add_subdirectory(fcntl)
   add_subdirectory(sched)
-  add_subdirectory(sys)
+  #add_subdirectory(sys)
   add_subdirectory(termios)
   add_subdirectory(unistd)
 endif()
diff --git a/libc/test/src/stdlib/CMakeLists.txt b/libc/test/src/stdlib/CMakeLists.txt
index f3033a4d958bf..2521e46d81bcf 100644
--- a/libc/test/src/stdlib/CMakeLists.txt
+++ b/libc/test/src/stdlib/CMakeLists.txt
@@ -88,6 +88,7 @@ add_libc_test(
   DEPENDS
     libc.src.__support.CPP.span
     libc.src.stdlib.freelist_heap
+    libc.src.stdlib.malloc
     libc.src.string.memcmp
     libc.src.string.memcpy
 )
@@ -437,19 +438,21 @@ if(LLVM_LIBC_FULL_BUILD)
       libc.src.stdlib.quick_exit
   )
 
-  # Only the GPU has an in-tree 'malloc' implementation.
-  if(LIBC_TARGET_OS_IS_GPU)
-    add_libc_test(
-      malloc_test
-      HERMETIC_TEST_ONLY
-      SUITE
-        libc-stdlib-tests
-      SRCS
-        malloc_test.cpp
-      DEPENDS
-        libc.include.stdlib
-        libc.src.stdlib.malloc
-        libc.src.stdlib.free
-    )
-  endif()
+  add_libc_test(
+    malloc_test
+    SUITE
+      libc-stdlib-tests
+    SRCS
+      block_test.cpp
+      malloc_test.cpp
+      freelist_malloc_test.cpp
+      freelist_heap_test.cpp
+      freelist_test.cpp
+    DEPENDS
+      libc.include.stdlib
+      libc.src.stdlib.malloc
+      libc.src.stdlib.free
+      libc.src.string.memcmp
+      libc.src.string.memcpy
+  )
 endif()
diff --git a/libc/test/src/stdlib/freelist_heap_test.cpp b/libc/test/src/stdlib/freelist_heap_test.cpp
index b89f47f9a2def..c9f61fa08db55 100644
--- a/libc/test/src/stdlib/freelist_heap_test.cpp
+++ b/libc/test/src/stdlib/freelist_heap_test.cpp
@@ -14,66 +14,80 @@
 
 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);
+using LIBC_NAMESPACE::freelist_heap;
 
-  void *ptr = allocator.allocate(ALLOC_SIZE);
+// Similar to `LlvmLibcBlockTest` in block_test.cpp, we'd like to run the same
+// tests independently for different parameters. In this case, we'd like to test
+// functionality for a `FreeListHeap` and the global `freelist_heap` which was
+// constinit'd. Functionally, it should operate the same if the FreeListHeap
+// were initialized locally at runtime or at compile-time.
+//
+// Note that calls to `allocate` for each test case here don't always explicitly
+// `free` them afterwards, so when testing the global allocator, allocations
+// made in tests leak and aren't free'd. This is fine for the purposes of this
+// test file.
+#define TEST_FOR_EACH_ALLOCATOR(TestCase, BufferSize)                          \
+  class LlvmLibcFreeListHeapTest##TestCase : public testing::Test {            \
+  public:                                                                      \
+    void RunTest(FreeListHeap<> &allocator, [[maybe_unused]] size_t N);        \
+  };                                                                           \
+  TEST_F(LlvmLibcFreeListHeapTest##TestCase, TestCase) {                       \
+    alignas(FreeListHeap<>::BlockType)                                         \
+        cpp::byte buf[BufferSize] = {cpp::byte(0)};                            \
+    FreeListHeap<> allocator(buf);                                             \
+    RunTest(allocator, BufferSize);                                            \
+    RunTest(*freelist_heap, freelist_heap->region_size());                     \
+  }                                                                            \
+  void LlvmLibcFreeListHeapTest##TestCase::RunTest(FreeListHeap<> &allocator,  \
+                                                   size_t N)
+
+TEST_FOR_EACH_ALLOCATOR(CanAllocate, 2048) {
+  constexpr size_t kAllocSize = 512;
+
+  void *ptr = allocator.allocate(kAllocSize);
 
   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));
+                     reinterpret_cast<cpp::byte *>(allocator.region_start()) +
+                     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)};
+TEST_FOR_EACH_ALLOCATOR(AllocationsDontOverlap, 2048) {
+  constexpr size_t kAllocSize = 512;
 
-  FreeListHeap<> allocator(buf);
-
-  void *ptr1 = allocator.allocate(ALLOC_SIZE);
-  void *ptr2 = allocator.allocate(ALLOC_SIZE);
+  void *ptr1 = allocator.allocate(kAllocSize);
+  void *ptr2 = allocator.allocate(kAllocSize);
 
   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 ptr1_end = ptr1_start + kAllocSize;
   uintptr_t ptr2_start = reinterpret_cast<uintptr_t>(ptr2);
 
   EXPECT_GT(ptr2_start, ptr1_end);
 }
 
-TEST(LlvmLibcFreeListHeap, CanFreeAndRealloc) {
+TEST_FOR_EACH_ALLOCATOR(CanFreeAndRealloc, 2048) {
   // 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);
+  constexpr size_t kAllocSize = 512;
 
-  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  void *ptr1 = allocator.allocate(kAllocSize);
   allocator.free(ptr1);
-  void *ptr2 = allocator.allocate(ALLOC_SIZE);
+  void *ptr2 = allocator.allocate(kAllocSize);
 
   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);
-
+TEST_FOR_EACH_ALLOCATOR(ReturnsNullWhenAllocationTooLarge, 2048) {
   EXPECT_EQ(allocator.allocate(N), static_cast<void *>(nullptr));
 }
 
+// NOTE: This doesn't use TEST_FOR_EACH_ALLOCATOR because the first `allocate`
+// here will likely actually return a nullptr since the same global allocator
+// is used for other test cases and we don't explicitly free them.
 TEST(LlvmLibcFreeListHeap, ReturnsNullWhenFull) {
   constexpr size_t N = 2048;
   alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(0)};
@@ -85,12 +99,7 @@ TEST(LlvmLibcFreeListHeap, ReturnsNullWhenFull) {
   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);
-
+TEST_FOR_EACH_ALLOCATOR(ReturnedPointersAreAligned, 2048) {
   void *ptr1 = allocator.allocate(1);
 
   // Should be aligned to native pointer alignment
@@ -105,84 +114,64 @@ TEST(LlvmLibcFreeListHeap, ReturnedPointersAreAligned) {
   EXPECT_EQ(ptr2_start % alignment, static_cast<size_t>(0));
 }
 
-TEST(LlvmLibcFreeListHeap, CanRealloc) {
-  constexpr size_t N = 2048;
-  constexpr size_t ALLOC_SIZE = 512;
+TEST_FOR_EACH_ALLOCATOR(CanRealloc, 2048) {
+  constexpr size_t kAllocSize = 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 *ptr1 = allocator.allocate(kAllocSize);
   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);
+TEST_FOR_EACH_ALLOCATOR(ReallocHasSameContent, 2048) {
+  constexpr size_t kAllocSize = 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];
+  cpp::byte data1[kAllocSize];
   // Data inside the reallocated block.
-  cpp::byte data2[ALLOC_SIZE];
-
-  FreeListHeap<> allocator(buf);
+  cpp::byte data2[kAllocSize];
 
-  int *ptr1 = reinterpret_cast<int *>(allocator.allocate(ALLOC_SIZE));
+  int *ptr1 = reinterpret_cast<int *>(allocator.allocate(kAllocSize));
   *ptr1 = 42;
-  memcpy(data1, ptr1, ALLOC_SIZE);
+  LIBC_NAMESPACE::memcpy(data1, ptr1, kAllocSize);
   int *ptr2 = reinterpret_cast<int *>(allocator.realloc(ptr1, kNewAllocSize));
-  memcpy(data2, ptr2, ALLOC_SIZE);
+  LIBC_NAMESPACE::memcpy(data2, ptr2, kAllocSize);
 
   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);
+  EXPECT_EQ(LIBC_NAMESPACE::memcmp(data1, data2, kAllocSize), 0);
 }
 
-TEST(LlvmLibcFreeListHeap, ReturnsNullReallocFreedPointer) {
-  constexpr size_t N = 2048;
-  constexpr size_t ALLOC_SIZE = 512;
+TEST_FOR_EACH_ALLOCATOR(ReturnsNullReallocFreedPointer, 2048) {
+  constexpr size_t kAllocSize = 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 *ptr1 = allocator.allocate(kAllocSize);
   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;
+TEST_FOR_EACH_ALLOCATOR(ReallocSmallerSize, 2048) {
+  constexpr size_t kAllocSize = 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 *ptr1 = allocator.allocate(kAllocSize);
   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)};
+TEST_FOR_EACH_ALLOCATOR(ReallocTooLarge, 2048) {
+  constexpr size_t kAllocSize = 512;
+  size_t kNewAllocSize = N * 2; // Large enough to fail.
 
-  FreeListHeap<> allocator(buf);
-
-  void *ptr1 = allocator.allocate(ALLOC_SIZE);
+  void *ptr1 = allocator.allocate(kAllocSize);
   void *ptr2 = allocator.realloc(ptr1, kNewAllocSize);
 
   // realloc() will not invalidate the original pointer if realloc() fails
@@ -190,50 +179,39 @@ TEST(LlvmLibcFreeListHeap, ReallocTooLarge) {
   EXPECT_EQ(static_cast<void *>(nullptr), ptr2);
 }
 
-TEST(LlvmLibcFreeListHeap, CanCalloc) {
-  constexpr size_t N = 2048;
-  constexpr size_t ALLOC_SIZE = 128;
+TEST_FOR_EACH_ALLOCATOR(CanCalloc, 2048) {
+  constexpr size_t kAllocSize = 128;
   constexpr size_t kNum = 4;
-  constexpr int size = kNum * ALLOC_SIZE;
-  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(1)};
+  constexpr int size = kNum * kAllocSize;
   constexpr cpp::byte zero{0};
 
-  FreeListHeap<> allocator(buf);
-
   cpp::byte *ptr1 =
-      reinterpret_cast<cpp::byte *>(allocator.calloc(kNum, ALLOC_SIZE));
+      reinterpret_cast<cpp::byte *>(allocator.calloc(kNum, kAllocSize));
 
   // calloc'd content is zero.
-  for (int i = 0; i < size; i++)
+  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;
+TEST_FOR_EACH_ALLOCATOR(CanCallocWeirdSize, 2048) {
+  constexpr size_t kAllocSize = 143;
   constexpr size_t kNum = 3;
-  constexpr int size = kNum * ALLOC_SIZE;
-  alignas(FreeListHeap<>::BlockType) cpp::byte buf[N] = {cpp::byte(132)};
+  constexpr int size = kNum * kAllocSize;
   constexpr cpp::byte zero{0};
 
-  FreeListHeap<> allocator(buf);
-
   cpp::byte *ptr1 =
-      reinterpret_cast<cpp::byte *>(allocator.calloc(kNum, ALLOC_SIZE));
+      reinterpret_cast<cpp::byte *>(allocator.calloc(kNum, kAllocSize));
 
   // calloc'd content is zero.
-  for (int i = 0; i < size; i++)
+  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));
+TEST_FOR_EACH_ALLOCATOR(CallocTooLarge, 2048) {
+  size_t kAllocSize = N + 1;
+  EXPECT_EQ(allocator.calloc(1, kAllocSize), static_cast<void *>(nullptr));
 }
 
 } // namespace LIBC_NAMESPACE
diff --git a/libc/test/src/stdlib/freelist_malloc_test.cpp b/libc/test/src/stdlib/freelist_malloc_test.cpp
new file mode 100644
index 0000000000000..0f6f7bcacee9a
--- /dev/null
+++ b/libc/test/src/stdlib/freelist_malloc_test.cpp
@@ -0,0 +1,56 @@
+//===-- Unittests for freelist_malloc -------------------------------------===//
+//
+// 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/stdlib/calloc.h"
+#include "src/stdlib/free.h"
+#include "src/stdlib/freelist_heap.h"
+#include "src/stdlib/malloc.h"
+#include "test/UnitTest/Test.h"
+
+using LIBC_NAMESPACE::freelist_heap;
+
+TEST(LlvmLibcFreeListMalloc, ReplacingMalloc) {
+  constexpr size_t kAllocSize = 256;
+  constexpr size_t kCallocNum = 4;
+  constexpr size_t kCallocSize = 64;
+
+  freelist_heap->reset_heap_stats(); // Do this because other tests might've
+                                     // called the same global allocator.
+
+  void *ptr1 = LIBC_NAMESPACE::malloc(kAllocSize);
+
+  const auto &freelist_heap_stats = freelist_heap->heap_stats();
+
+  ASSERT_NE(ptr1, static_cast<void *>(nullptr));
+  EXPECT_EQ(freelist_heap_stats.bytes_allocated, kAllocSize);
+  EXPECT_EQ(freelist_heap_stats.cumulative_allocated, kAllocSize);
+  EXPECT_EQ(freelist_heap_stats.cumulative_freed, size_t(0));
+
+  LIBC_NAMESPACE::free(ptr1);
+  EXPECT_EQ(freelist_heap_stats.bytes_allocated, size_t(0));
+  EXPECT_EQ(freelist_heap_stats.cumulative_allocated, kAllocSize);
+  EXPECT_EQ(freelist_heap_stats.cumulative_freed, kAllocSize);
+
+  void *ptr2 = LIBC_NAMESPACE::calloc(kCallocNum, kCallocSize);
+  ASSERT_NE(ptr2, static_cast<void *>(nullptr));
+  EXPECT_EQ(freelist_heap_stats.bytes_allocated, kCallocNum * kCallocSize);
+  EXPECT_EQ(freelist_heap_stats.cumulative_allocated,
+            kAllocSize + kCallocNum * kCallocSize);
+  EXPECT_EQ(freelist_heap_stats.cumulative_freed, kAllocSize);
+
+  for (size_t i = 0; i < kCallocNum * kCallocSize; ++i) {
+    EXPECT_EQ(reinterpret_cast<uint8_t *>(ptr2)[i], uint8_t(0));
+  }
+
+  LIBC_NAMESPACE::free(ptr2);
+  EXPECT_EQ(freelist_heap_stats.bytes_allocated, size_t(0));
+  EXPECT_EQ(freelist_heap_stats.cumulative_allocated,
+            kAllocSize + kCallocNum * kCallocSize);
+  EXPECT_EQ(freelist_heap_stats.cumulative_freed,
+            kAllocSize + kCallocNum * kCallocSize);
+}



More information about the libc-commits mailing list