[libc] Lazily initialize freelist malloc using symbols

This allows the user to customize the size of the heap provided by
overriding the weak symbol __libc_heap_limit at link time. The heap
begins at _end and ends __libc_heap_limit bytes afterwards. It also
prevents a completely unused heap from requiring any space. (It's
reasonably common to link in malloc/free to handle exceptional
situations that can never actually be encountered.)

I'd think this should eventually be replaced with an implemenation based
on sbrk(), with sbrk() implemented generically in terms of brk(), which
would be a system call on e.g. Linux and a bump pointer from _end up to
__libc_heap_limit on baremetal. This would allow the allocator to be
more flexible (e.g., an implementation could swap out brk() and do
dynamic memory availability checks), to manage the heap better by
keeping better track of "wilderness" and to work properly as a malloc on
Linux.

Incidentally, this sets the default heap limit to 128KiB, from 102400
bytes, which had been reported as "1GiB".

libc/config/baremetal/config.json

@@ -15,7 +15,7 @@
   },
   "malloc": {
     "LIBC_CONF_FREELIST_MALLOC_BUFFER_SIZE": {
-      "value": 102400
+      "value": 131072
     }
   }
 }

libc/src/__support/freelist_heap.h

@@ -20,6 +20,9 @@
 #include "src/string/memory_utils/inline_memcpy.h"
 #include "src/string/memory_utils/inline_memset.h"
 
+extern char _end;
+extern char __libc_heap_limit;
+
 namespace LIBC_NAMESPACE_DECL {
 
 using cpp::optional;
@@ -47,17 +50,10 @@ template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()> class FreeListHeap {
     size_t total_free_calls;
   };
 
-  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;
+  constexpr FreeListHeap(span<cpp::byte> region)
+      : is_initialized_(false), region_(region), freelist_(DEFAULT_BUCKETS),
+        heap_stats_{} {
+    heap_stats_.total_bytes = region.size();
   }
 
   void *allocate(size_t size);
@@ -69,61 +65,68 @@ 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_);
-  }
+  cpp::span<cpp::byte> region() const { return region_; }
 
-protected:
-  constexpr void set_freelist_node(typename FreeListType::FreeListNode &node,
-                                   cpp::span<cpp::byte> chunk) {
-    freelist_.set_freelist_node(node, chunk);
-  }
+private:
+  void init();
 
   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());
   }
 
   bool is_valid_ptr(void *ptr) {
-    return ptr >= block_region_start_ && ptr < block_region_end_;
+    return ptr >= region_.begin() && ptr < region_.end();
   }
 
-  void *block_region_start_;
-  void *block_region_end_;
+  bool is_initialized_;
+  cpp::span<cpp::byte> region_;
   FreeListType freelist_;
   HeapStats heap_stats_;
 };
 
 template <size_t BUFF_SIZE, size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()>
-struct FreeListHeapBuffer : public FreeListHeap<NUM_BUCKETS> {
+class FreeListHeapBuffer : public FreeListHeap<NUM_BUCKETS> {
   using parent = FreeListHeap<NUM_BUCKETS>;
   using FreeListNode = typename parent::FreeListType::FreeListNode;
 
+public:
   constexpr FreeListHeapBuffer()
-      : FreeListHeap<NUM_BUCKETS>(&block, buffer + sizeof(buffer), BUFF_SIZE),
-        block(0, BUFF_SIZE), node{}, buffer{} {
-    block.mark_last();
+      : FreeListHeap<NUM_BUCKETS>{buffer}, buffer{} {}
 
-    cpp::span<cpp::byte> chunk(buffer, sizeof(buffer));
-    parent::set_freelist_node(node, chunk);
-  }
+private:
+  cpp::byte buffer[BUFF_SIZE];
+};
 
-  typename parent::BlockType block;
-  FreeListNode node;
-  cpp::byte buffer[BUFF_SIZE - sizeof(block) - sizeof(node)];
+template <size_t NUM_BUCKETS = DEFAULT_BUCKETS.size()>
+class FreeListHeapSymbols : public FreeListHeap<NUM_BUCKETS> {
+  using parent = FreeListHeap<NUM_BUCKETS>;
+  using FreeListNode = typename parent::FreeListType::FreeListNode;
+
+public:
+  constexpr FreeListHeapSymbols()
+      : FreeListHeap<NUM_BUCKETS>{{(cpp::byte*)&_end, (size_t)&__libc_heap_limit}} {}
 };
 
+template <size_t NUM_BUCKETS>
+void FreeListHeap<NUM_BUCKETS>::init() {
+  LIBC_ASSERT(!is_initialized_ && "duplicate initialization");
+  auto result = BlockType::init(region_);
+  BlockType *block = *result;
+  freelist_.add_chunk(block_to_span(block));
+  is_initialized_ = true;
+}
+
 template <size_t NUM_BUCKETS>
 void *FreeListHeap<NUM_BUCKETS>::allocate_impl(size_t alignment, size_t size) {
   if (size == 0)
     return nullptr;
 
+  if (!is_initialized_)
+    init();
+
   // Find a chunk in the freelist. Split it if needed, then return.
   auto chunk =
       freelist_.find_chunk_if([alignment, size](span<cpp::byte> chunk) {

libc/src/stdlib/CMakeLists.txt

@@ -377,6 +377,7 @@ if(NOT LIBC_TARGET_OS_IS_GPU)
       freelist_malloc
       SRCS
         freelist_malloc.cpp
+        heap_limit.S
       HDRS
         malloc.h
       DEPENDS

libc/src/stdlib/freelist_malloc.cpp

@@ -19,16 +19,13 @@
 namespace LIBC_NAMESPACE_DECL {
 
 namespace {
-#ifdef LIBC_FREELIST_MALLOC_SIZE
-// This is set via the LIBC_CONF_FREELIST_MALLOC_BUFFER_SIZE configuration.
-constexpr size_t SIZE = LIBC_FREELIST_MALLOC_SIZE;
-#else
+#ifndef LIBC_FREELIST_MALLOC_SIZE
 #error "LIBC_FREELIST_MALLOC_SIZE was not defined for this build."
 #endif
-LIBC_CONSTINIT FreeListHeapBuffer<SIZE> freelist_heap_buffer;
+FreeListHeapSymbols<> freelist_heap_symbols;
 } // namespace
 
-FreeListHeap<> *freelist_heap = &freelist_heap_buffer;
+FreeListHeap<> *freelist_heap = &freelist_heap_symbols;
 
 LLVM_LIBC_FUNCTION(void *, malloc, (size_t size)) {
   return freelist_heap->allocate(size);

libc/src/stdlib/heap_limit.S

@@ -0,0 +1,2 @@
+.weak __libc_heap_limit
+.set __libc_heap_limit, LIBC_FREELIST_MALLOC_SIZE

libc/test/src/__support/freelist_heap_test.cpp

@@ -30,14 +30,19 @@ using LIBC_NAMESPACE::freelist_heap;
 #define TEST_FOR_EACH_ALLOCATOR(TestCase, BufferSize)                          \
   class LlvmLibcFreeListHeapTest##TestCase : public testing::Test {            \
   public:                                                                      \
+    FreeListHeapBuffer<BufferSize> fake_global_buffer;                         \
+    void SetUp() override {                                                    \
+      freelist_heap =                                                          \
+          new (&fake_global_buffer) FreeListHeapBuffer<BufferSize>;            \
+    }                                                                          \
     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());                     \
+    RunTest(*freelist_heap, freelist_heap->region().size());                   \
   }                                                                            \
   void LlvmLibcFreeListHeapTest##TestCase::RunTest(FreeListHeap<> &allocator,  \
                                                    size_t N)

libc/test/src/__support/freelist_malloc_test.cpp

@@ -14,15 +14,22 @@
 #include "test/UnitTest/Test.h"
 
 using LIBC_NAMESPACE::freelist_heap;
+using LIBC_NAMESPACE::FreeListHeapBuffer;
 
-TEST(LlvmLibcFreeListMalloc, MallocStats) {
+class LlvmLibcFreeListMalloc : public LIBC_NAMESPACE::testing::Test {
+protected:
+  void SetUp() override {
+    freelist_heap = new (&buffer) FreeListHeapBuffer<1024>;
+  }
+
+  FreeListHeapBuffer<1024> buffer;
+};
+
+TEST_F(LlvmLibcFreeListMalloc, MallocStats) {
   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();