[libc][stdlib] Add freelist class

[PiJoules]

This implements a traditional freelist to be used by the freelist allocator. It operates on spans of bytes which can be anything. The freelist allocator will store Blocks inside them.

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

[llvmbot]

@llvm/pr-subscribers-libc

Author: None (PiJoules)

Changes

This implements a traditional freelist to be used by the freelist allocator. It operates on spans of bytes which can be anything. The freelist allocator will store Blocks inside them.

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

---

Full diff: https://github.com/llvm/llvm-project/pull/95041.diff

4 Files Affected:

• (modified) libc/src/stdlib/CMakeLists.txt (+10)
• (added) libc/src/stdlib/freelist.h (+199)
• (modified) libc/test/src/stdlib/CMakeLists.txt (+12)
• (added) libc/test/src/stdlib/freelist_test.cpp (+166)

diff --git a/libc/src/stdlib/CMakeLists.txt b/libc/src/stdlib/CMakeLists.txt
index afb2d6d91cba4..d4aa50a43d186 100644
--- a/libc/src/stdlib/CMakeLists.txt
+++ b/libc/src/stdlib/CMakeLists.txt
@@ -392,6 +392,16 @@ else()
       libc.src.__support.CPP.span
       libc.src.__support.CPP.type_traits
   )
+  add_header_library(
+    freelist
+    HDRS
+      freelist.h
+    DEPENDS
+      libc.src.__support.fixedvector
+      libc.src.__support.CPP.cstddef
+      libc.src.__support.CPP.array
+      libc.src.__support.CPP.span
+  )
   add_entrypoint_external(
     malloc
   )
diff --git a/libc/src/stdlib/freelist.h b/libc/src/stdlib/freelist.h
new file mode 100644
index 0000000000000..a5c214ebf6649
--- /dev/null
+++ b/libc/src/stdlib/freelist.h
@@ -0,0 +1,199 @@
+//===-- Interface 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_STDLIB_FREELIST_H
+#define LLVM_LIBC_SRC_STDLIB_FREELIST_H
+
+#include "src/__support/CPP/array.h"
+#include "src/__support/CPP/cstddef.h"
+#include "src/__support/CPP/span.h"
+#include "src/__support/fixedvector.h"
+
+namespace LIBC_NAMESPACE {
+
+using cpp::span;
+
+/// Basic [freelist](https://en.wikipedia.org/wiki/Free_list) implementation
+/// for an allocator. This implementation buckets by chunk size, with a list
+/// of user-provided buckets. Each bucket is a linked list of storage chunks.
+/// Because this freelist uses the added chunks themselves as list nodes, there
+/// is a lower bound of `sizeof(FreeList.FreeListNode)` bytes for chunks which
+/// can be added to this freelist. There is also an implicit bucket for
+/// "everything else", for chunks which do not fit into a bucket.
+///
+/// Each added chunk will be added to the smallest bucket under which it fits.
+/// If it does not fit into any user-provided bucket, it will be added to the
+/// default bucket.
+///
+/// As an example, assume that the `FreeList` is configured with buckets of
+/// sizes {64, 128, 256, and 512} bytes. The internal state may look like the
+/// following:
+///
+/// @code{.unparsed}
+/// bucket[0] (64B) --> chunk[12B] --> chunk[42B] --> chunk[64B] --> NULL
+/// bucket[1] (128B) --> chunk[65B] --> chunk[72B] --> NULL
+/// bucket[2] (256B) --> NULL
+/// bucket[3] (512B) --> chunk[312B] --> chunk[512B] --> chunk[416B] --> NULL
+/// bucket[4] (implicit) --> chunk[1024B] --> chunk[513B] --> NULL
+/// @endcode
+///
+/// Note that added chunks should be aligned to a 4-byte boundary.
+template <size_t NUM_BUCKETS = 6> class FreeList {
+public:
+  // Remove copy/move ctors
+  FreeList(const FreeList &other) = delete;
+  FreeList(FreeList &&other) = delete;
+  FreeList &operator=(const FreeList &other) = delete;
+  FreeList &operator=(FreeList &&other) = delete;
+
+  /// Adds a chunk to this freelist.
+  bool add_chunk(cpp::span<cpp::byte> chunk);
+
+  /// Finds an eligible chunk for an allocation of size `size`.
+  ///
+  /// @note This returns the first allocation possible within a given bucket;
+  /// It does not currently optimize for finding the smallest chunk.
+  ///
+  /// @returns
+  /// * On success - A span representing the chunk.
+  /// * On failure (e.g. there were no chunks available for that allocation) -
+  ///   A span with a size of 0.
+  cpp::span<cpp::byte> find_chunk(size_t size) const;
+
+  /// 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.
+  unsigned short 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)
+      : chunks_(NUM_BUCKETS + 1, 0), sizes_(sizes.begin(), sizes.end()) {}
+
+  FixedVector<FreeList::FreeListNode *, NUM_BUCKETS + 1> chunks_;
+  FixedVector<size_t, NUM_BUCKETS> sizes_;
+};
+
+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;
+
+  union {
+    FreeListNode *node;
+    cpp::byte *bytes;
+  } aliased;
+
+  aliased.bytes = chunk.data();
+
+  unsigned short 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;
+
+  return true;
+}
+
+template <size_t NUM_BUCKETS>
+span<cpp::byte> FreeList<NUM_BUCKETS>::find_chunk(size_t size) const {
+  if (size == 0)
+    return span<cpp::byte>();
+
+  unsigned short chunk_ptr = find_chunk_ptr_for_size(size, true);
+
+  // Check that there's data. This catches the case where we run off the
+  // end of the array
+  if (chunks_[chunk_ptr] == nullptr)
+    return span<cpp::byte>();
+
+  // Now iterate up the buckets, walking each list to find a good candidate
+  for (size_t i = chunk_ptr; i < chunks_.size(); i++) {
+    union {
+      FreeListNode *node;
+      cpp::byte *data;
+    } aliased;
+    aliased.node = chunks_[static_cast<unsigned short>(i)];
+
+    while (aliased.node != nullptr) {
+      if (aliased.node->size >= size)
+        return span<cpp::byte>(aliased.data, aliased.node->size);
+
+      aliased.node = aliased.node->next;
+    }
+  }
+
+  // If we get here, we've checked every block in every bucket. There's
+  // nothing that can support this allocation.
+  return span<cpp::byte>();
+}
+
+template <size_t NUM_BUCKETS>
+bool FreeList<NUM_BUCKETS>::remove_chunk(span<cpp::byte> chunk) {
+  unsigned short chunk_ptr = find_chunk_ptr_for_size(chunk.size(), true);
+
+  // Walk that list, finding the chunk.
+  union {
+    FreeListNode *node;
+    cpp::byte *data;
+  } aliased, aliased_next;
+
+  // Check head first.
+  if (chunks_[chunk_ptr] == nullptr)
+    return false;
+
+  aliased.node = chunks_[chunk_ptr];
+  if (aliased.data == chunk.data()) {
+    chunks_[chunk_ptr] = aliased.node->next;
+    return true;
+  }
+
+  // No? Walk the nodes.
+  aliased.node = chunks_[chunk_ptr];
+
+  while (aliased.node->next != nullptr) {
+    aliased_next.node = aliased.node->next;
+    if (aliased_next.data == chunk.data()) {
+      // Found it, remove this node out of the chain
+      aliased.node->next = aliased_next.node->next;
+      return true;
+    }
+
+    aliased.node = aliased.node->next;
+  }
+
+  return false;
+}
+
+template <size_t NUM_BUCKETS>
+unsigned short
+FreeList<NUM_BUCKETS>::find_chunk_ptr_for_size(size_t size,
+                                               bool non_null) const {
+  unsigned short chunk_ptr = 0;
+  for (chunk_ptr = 0u; chunk_ptr < sizes_.size(); chunk_ptr++) {
+    if (sizes_[chunk_ptr] >= size &&
+        (!non_null || chunks_[chunk_ptr] != nullptr)) {
+      break;
+    }
+  }
+
+  return chunk_ptr;
+}
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_FREELIST_H
diff --git a/libc/test/src/stdlib/CMakeLists.txt b/libc/test/src/stdlib/CMakeLists.txt
index f122cd56a6060..d3954f077a219 100644
--- a/libc/test/src/stdlib/CMakeLists.txt
+++ b/libc/test/src/stdlib/CMakeLists.txt
@@ -67,6 +67,18 @@ add_libc_test(
     libc.src.string.memcpy
 )
 
+add_libc_test(
+  freelist_test
+  SUITE
+    libc-stdlib-tests
+  SRCS
+    freelist_test.cpp
+  DEPENDS
+    libc.src.stdlib.freelist
+    libc.src.__support.CPP.array
+    libc.src.__support.CPP.span
+)
+
 add_fp_unittest(
   strtod_test
   SUITE
diff --git a/libc/test/src/stdlib/freelist_test.cpp b/libc/test/src/stdlib/freelist_test.cpp
new file mode 100644
index 0000000000000..e25c74b47b852
--- /dev/null
+++ b/libc/test/src/stdlib/freelist_test.cpp
@@ -0,0 +1,166 @@
+//===-- Unittests for a freelist --------------------------------*- 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>
+
+#include "src/__support/CPP/array.h"
+#include "src/__support/CPP/span.h"
+#include "src/stdlib/freelist.h"
+#include "test/UnitTest/Test.h"
+
+using LIBC_NAMESPACE::FreeList;
+using LIBC_NAMESPACE::cpp::array;
+using LIBC_NAMESPACE::cpp::byte;
+using LIBC_NAMESPACE::cpp::span;
+
+static constexpr size_t SIZE = 8;
+static constexpr array<size_t, SIZE> example_sizes = {64,   128,  256,  512,
+                                                      1024, 2048, 4096, 8192};
+
+TEST(LlvmLibcFreeList, EmptyListHasNoMembers) {
+  FreeList<SIZE> list(example_sizes);
+
+  auto item = list.find_chunk(4);
+  EXPECT_EQ(item.size(), static_cast<size_t>(0));
+  item = list.find_chunk(128);
+  EXPECT_EQ(item.size(), static_cast<size_t>(0));
+}
+
+TEST(LlvmLibcFreeList, CanRetrieveAddedMember) {
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t N = 512;
+
+  byte data[N] = {byte(0)};
+
+  bool ok = list.add_chunk(span<byte>(data, N));
+  EXPECT_TRUE(ok);
+
+  auto item = list.find_chunk(N);
+  EXPECT_EQ(item.size(), N);
+  EXPECT_EQ(item.data(), data);
+}
+
+TEST(LlvmLibcFreeList, CanRetrieveAddedMemberForSmallerSize) {
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t N = 512;
+
+  byte data[N] = {byte(0)};
+
+  ASSERT_TRUE(list.add_chunk(span<byte>(data, N)));
+  auto item = list.find_chunk(N / 2);
+  EXPECT_EQ(item.size(), N);
+  EXPECT_EQ(item.data(), data);
+}
+
+TEST(LlvmLibcFreeList, CanRemoveItem) {
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t N = 512;
+
+  byte data[N] = {byte(0)};
+
+  ASSERT_TRUE(list.add_chunk(span<byte>(data, N)));
+  EXPECT_TRUE(list.remove_chunk(span<byte>(data, N)));
+
+  auto item = list.find_chunk(N);
+  EXPECT_EQ(item.size(), static_cast<size_t>(0));
+}
+
+TEST(LlvmLibcFreeList, FindReturnsSmallestChunk) {
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t kN1 = 512;
+  constexpr size_t kN2 = 1024;
+
+  byte data1[kN1] = {byte(0)};
+  byte data2[kN2] = {byte(0)};
+
+  ASSERT_TRUE(list.add_chunk(span<byte>(data1, kN1)));
+  ASSERT_TRUE(list.add_chunk(span<byte>(data2, kN2)));
+
+  auto chunk = list.find_chunk(kN1 / 2);
+  EXPECT_EQ(chunk.size(), kN1);
+  EXPECT_EQ(chunk.data(), data1);
+
+  chunk = list.find_chunk(kN1);
+  EXPECT_EQ(chunk.size(), kN1);
+  EXPECT_EQ(chunk.data(), data1);
+
+  chunk = list.find_chunk(kN1 + 1);
+  EXPECT_EQ(chunk.size(), kN2);
+  EXPECT_EQ(chunk.data(), data2);
+}
+
+TEST(LlvmLibcFreeList, FindReturnsCorrectChunkInSameBucket) {
+  // If we have two values in the same bucket, ensure that the allocation will
+  // pick an appropriately sized one.
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t kN1 = 512;
+  constexpr size_t kN2 = 257;
+
+  byte data1[kN1] = {byte(0)};
+  byte data2[kN2] = {byte(0)};
+
+  // List should now be 257 -> 512 -> NULL
+  ASSERT_TRUE(list.add_chunk(span<byte>(data1, kN1)));
+  ASSERT_TRUE(list.add_chunk(span<byte>(data2, kN2)));
+
+  auto chunk = list.find_chunk(kN2 + 1);
+  EXPECT_EQ(chunk.size(), kN1);
+}
+
+TEST(LlvmLibcFreeList, FindCanMoveUpThroughBuckets) {
+  // Ensure that finding a chunk will move up through buckets if no appropriate
+  // chunks were found in a given bucket
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t kN1 = 257;
+  constexpr size_t kN2 = 513;
+
+  byte data1[kN1] = {byte(0)};
+  byte data2[kN2] = {byte(0)};
+
+  // List should now be:
+  // bkt[3] (257 bytes up to 512 bytes) -> 257 -> NULL
+  // bkt[4] (513 bytes up to 1024 bytes) -> 513 -> NULL
+  ASSERT_TRUE(list.add_chunk(span<byte>(data1, kN1)));
+  ASSERT_TRUE(list.add_chunk(span<byte>(data2, kN2)));
+
+  // Request a 300 byte chunk. This should return the 513 byte one
+  auto chunk = list.find_chunk(kN1 + 1);
+  EXPECT_EQ(chunk.size(), kN2);
+}
+
+TEST(LlvmLibcFreeList, RemoveUnknownChunkReturnsNotFound) {
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t N = 512;
+
+  byte data[N] = {byte(0)};
+  byte data2[N] = {byte(0)};
+
+  ASSERT_TRUE(list.add_chunk(span<byte>(data, N)));
+  EXPECT_FALSE(list.remove_chunk(span<byte>(data2, N)));
+}
+
+TEST(LlvmLibcFreeList, CanStoreMultipleChunksPerBucket) {
+  FreeList<SIZE> list(example_sizes);
+  constexpr size_t N = 512;
+
+  byte data1[N] = {byte(0)};
+  byte data2[N] = {byte(0)};
+
+  ASSERT_TRUE(list.add_chunk(span<byte>(data1, N)));
+  ASSERT_TRUE(list.add_chunk(span<byte>(data2, N)));
+
+  auto chunk1 = list.find_chunk(N);
+  ASSERT_TRUE(list.remove_chunk(chunk1));
+  auto chunk2 = list.find_chunk(N);
+  ASSERT_TRUE(list.remove_chunk(chunk2));
+
+  // Ordering of the chunks doesn't matter
+  EXPECT_TRUE(chunk1.data() != chunk2.data());
+  EXPECT_TRUE(chunk1.data() == data1 || chunk1.data() == data2);
+  EXPECT_TRUE(chunk2.data() == data1 || chunk2.data() == data2);
+}

[nickdesaulniers]

This should be moved under src/__support/. I'd expect everything under stc/stdlib/ to be related to functions hope to find in stdlib.h.