[WIP][libc] Add freelist malloc

Author: PiJoules

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

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);

libc/src/stdlib/CMakeLists.txt

@@ -401,9 +401,24 @@ else()
       libc.src.__support.CPP.cstddef
       libc.src.__support.CPP.array
       libc.src.__support.CPP.span
+      libc.src.stdio.printf
   )
-  add_entrypoint_external(
+  add_entrypoint_object(
     malloc
+    SRCS
+      freelist_malloc.cpp
+    HDRS
+      malloc.h
+    DEPENDS
+      .block
+      .freelist
+      libc.src.__support.CPP.new
+      libc.src.__support.CPP.optional
+      libc.src.__support.CPP.span
+      libc.src.__support.CPP.type_traits
+      libc.src.__support.fixedvector
+      libc.src.string.memcpy
+      libc.src.string.memset
   )
   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

@@ -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 &&

libc/src/stdlib/freelist_heap.h

@@ -0,0 +1,230 @@
+//===-- 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/string/memcpy.h"
+#include "src/string/memset.h"
+
+namespace LIBC_NAMESPACE {
+
+using cpp::optional;
+using cpp::span;
+
+static constexpr cpp::array<size_t, 6> defaultBuckets{16,  32,  64,
+                                                      128, 256, 512};
+
+template <size_t kNumBuckets = defaultBuckets.size()> class FreeListHeap {
+public:
+  using BlockType = Block<>;
+  using FreeListType = FreeList<kNumBuckets>;
+
+  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)
+      : FreeListHeap(&*region.begin(), &*region.end(), region.size()) {
+    auto result = BlockType::init(region);
+    BlockType *block = *result;
+    freelist_.add_chunk(BlockToSpan(block));
+  }
+
+  constexpr FreeListHeap(void *start, cpp::byte *end, size_t total_bytes)
+      : block_region_start_(start), block_region_end_(end),
+        freelist_(defaultBuckets), 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);
+
+  void LogHeapStats();
+  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:
+  constexpr span<cpp::byte> BlockToSpan(BlockType *block) {
+    return span<cpp::byte>(block->usable_space(), block->inner_size());
+  }
+
+  void InvalidFreeCrash() { __builtin_trap(); }
+
+  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 kBuffSize, size_t kNumBuckets = defaultBuckets.size()>
+struct FreeListHeapBuffer : public FreeListHeap<kNumBuckets> {
+  using parent = FreeListHeap<kNumBuckets>;
+  using FreeListNode = typename parent::FreeListType::FreeListNode;
+
+  constexpr FreeListHeapBuffer()
+      : FreeListHeap<kNumBuckets>(&block, buffer + sizeof(buffer), kBuffSize),
+        block(0, kBuffSize), node{}, buffer{} {
+    block.mark_last();
+
+    cpp::span<cpp::byte> chunk(buffer, sizeof(buffer));
+    parent::set_freelist_node(&node, chunk);
+  }
+
+  typename parent::BlockType block;
+  FreeListNode node;
+  cpp::byte buffer[kBuffSize - sizeof(block) - sizeof(node)];
+};
+
+template <size_t kNumBuckets>
+void *FreeListHeap<kNumBuckets>::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(BlockToSpan(*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 kNumBuckets> void FreeListHeap<kNumBuckets>::Free(void *ptr) {
+  cpp::byte *bytes = static_cast<cpp::byte *>(ptr);
+
+  if (!is_valid_ptr(bytes)) {
+    InvalidFreeCrash();
+    return;
+  }
+
+  BlockType *chunk_block = BlockType::from_usable_space(bytes);
+
+  size_t size_freed = chunk_block->inner_size();
+  // Ensure that the block is in-use
+  if (!chunk_block->used()) {
+    InvalidFreeCrash();
+    return;
+  }
+  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(BlockToSpan(prev));
+    chunk_block = chunk_block->prev();
+    BlockType::merge_next(chunk_block);
+  }
+
+  if (next != nullptr && !next->used()) {
+    freelist_.remove_chunk(BlockToSpan(next));
+    BlockType::merge_next(chunk_block);
+  }
+  // Add back to the freelist
+  freelist_.add_chunk(BlockToSpan(chunk_block));
+
+  heap_stats_.bytes_allocated -= size_freed;
+  heap_stats_.cumulative_freed += size_freed;
+  heap_stats_.total_free_calls += 1;
+}
+
+// Follows constract of the C standard realloc() function
+// If ptr is free'd, will return nullptr.
+template <size_t kNumBuckets>
+void *FreeListHeap<kNumBuckets>::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 (!is_valid_ptr(bytes))
+    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;
+  memcpy(new_ptr, ptr, old_size);
+
+  Free(ptr);
+  return new_ptr;
+}
+
+template <size_t kNumBuckets>
+void *FreeListHeap<kNumBuckets>::Calloc(size_t num, size_t size) {
+  void *ptr = Allocate(num * size);
+  if (ptr != nullptr)
+    memset(ptr, 0, num * size);
+  return ptr;
+}
+
+extern FreeListHeap<> *freelist_heap;
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_FREELIST_HEAP_H