[libc][stdlib] Make the FreeListHeap constant-initializable (#95453)

This refactors some of the FreeListHeap, FreeList, and Block classes to
have constexpr ctors so we can constinit a global allocator that does
not require running some global function or global ctor to initialize.
This is needed to prevent worrying about initialization order and any
other module-ctor can invoke malloc without worry.

Author: PiJoules

libc/src/__support/fixedvector.h

@@ -30,6 +30,13 @@ template <typename T, size_t CAPACITY> class FixedVector {
       push_back(*begin);
   }
 
+  using const_iterator = typename cpp::array<T, CAPACITY>::const_iterator;
+  constexpr FixedVector(const_iterator begin, const_iterator end)
+      : store{}, item_count{} {
+    for (; begin != end; ++begin)
+      push_back(*begin);
+  }
+
   constexpr FixedVector(size_t count, const T &value) : store{}, item_count{} {
     for (size_t i = 0; i < count; ++i)
       push_back(value);

libc/src/stdlib/CMakeLists.txt

@@ -418,8 +418,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

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,15 +259,15 @@ 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);
 
   /// 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;

libc/src/stdlib/freelist.h

@@ -69,36 +69,44 @@ 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(const 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;
 
-  // Add it to the correct list.
-  size_t chunk_ptr = find_chunk_ptr_for_size(chunk.size(), false);
-
-  FreeListNode *node =
-      ::new (chunk.data()) FreeListNode{chunks_[chunk_ptr], chunk.size()};
-  chunks_[chunk_ptr] = node;
+  FreeListNode *node = ::new (chunk.data()) FreeListNode;
+  set_freelist_node(*node, chunk);
 
   return true;
 }
@@ -163,8 +171,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 &&

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,35 +40,73 @@ 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);
   void *realloc(void *ptr, size_t size);
   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

libc/src/stdlib/freelist_malloc.cpp

@@ -0,0 +1,48 @@
+//===-- 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 "src/stdlib/calloc.h"
+#include "src/stdlib/free.h"
+#include "src/stdlib/malloc.h"
+#include "src/stdlib/realloc.h"
+
+#include <stddef.h>
+
+namespace LIBC_NAMESPACE {
+
+namespace {
+// Users can define LIBC_FREELIST_MALLOC_SIZE for setting the default buffer
+// size used by freelist malloc.
+#ifdef LIBC_FREELIST_MALLOC_SIZE
+constexpr size_t SIZE = LIBC_FREELIST_MALLOC_SIZE;
+#else
+// TODO: We should probably have something akin to what scudo/sanitizer
+// allocators do where each platform defines this.
+constexpr size_t SIZE = 0x40000000ULL; // 1GB
+#endif
+LIBC_CONSTINIT FreeListHeapBuffer<SIZE> freelist_heap_buffer;
+} // namespace
+
+FreeListHeap<> *freelist_heap = &freelist_heap_buffer;
+
+LLVM_LIBC_FUNCTION(void *, malloc, (size_t size)) {
+  return freelist_heap->allocate(size);
+}
+
+LLVM_LIBC_FUNCTION(void, free, (void *ptr)) { return freelist_heap->free(ptr); }
+
+LLVM_LIBC_FUNCTION(void *, calloc, (size_t num, size_t size)) {
+  return freelist_heap->calloc(num, size);
+}
+
+LLVM_LIBC_FUNCTION(void *, realloc, (void *ptr, size_t size)) {
+  return freelist_heap->realloc(ptr, size);
+}
+
+} // namespace LIBC_NAMESPACE

libc/src/stdlib/realloc.h

@@ -0,0 +1,20 @@
+//===-- Implementation header for realloc -----------------------*- 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_REALLOC_H
+#define LLVM_LIBC_SRC_STDLIB_REALLOC_H
+
+namespace LIBC_NAMESPACE {
+
+void *realloc(void *ptr, size_t size);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_REALLOC_H

libc/test/src/stdlib/CMakeLists.txt

@@ -85,9 +85,11 @@ add_libc_test(
     libc-stdlib-tests
   SRCS
     freelist_heap_test.cpp
+    freelist_malloc_test.cpp
   DEPENDS
     libc.src.__support.CPP.span
     libc.src.stdlib.freelist_heap
+    libc.src.stdlib.malloc
     libc.src.string.memcmp
     libc.src.string.memcpy
 )