[libc] Add a resizable container with constexpr constructor and destructor.

The new container is used to store atexit callbacks. This way, we avoid
the possibility of the destructor of the container itself getting added
as an at exit callback.

Reviewed By: abrachet

Differential Revision: https://reviews.llvm.org/D121350

7 files changed, 257 insertions, 13 deletions

diff --git a/libc/src/__support/CPP/CMakeLists.txt b/libc/src/__support/CPP/CMakeLists.txt
index dfc28ea4aead..a4ffea3a2838 100644
--- a/libc/src/__support/CPP/CMakeLists.txt
+++ b/libc/src/__support/CPP/CMakeLists.txt
@@ -67,3 +67,9 @@ add_header_library(
   HDRS
     atomic.h
 )
+
+add_header_library(
+  blockstore
+  HDRS
+    blockstore.h
+)
diff --git a/libc/src/__support/CPP/blockstore.h b/libc/src/__support/CPP/blockstore.h
new file mode 100644
index 000000000000..b4c599881482
--- /dev/null
+++ b/libc/src/__support/CPP/blockstore.h
@@ -0,0 +1,160 @@
+//===-- A data structure which stores data in blocks  -----------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SUPPORT_CPP_BLOCKSTORE_H
+#define LLVM_LIBC_SUPPORT_CPP_BLOCKSTORE_H
+
+#include <stddef.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+namespace __llvm_libc {
+namespace cpp {
+
+// The difference between BlockStore a traditional vector types is that,
+// when more capacity is desired, a new block is added instead of allocating
+// a larger sized array and copying over existing items to the new allocation.
+// Also, the initial block does not need heap allocation. Hence, a BlockStore is
+// suitable for global objects as it does not require explicit construction.
+// Also, the destructor of this class does nothing, which eliminates the need
+// for an atexit global object destruction. But, it also means that the global
+// object should be explicitly cleaned up at the appropriate time.
+//
+// If REVERSE_ORDER is true, the iteration of elements will in the reverse
+// order. Also, since REVERSE_ORDER is a constexpr, conditionals branching
+// on its value will be optimized out in the code below.
+template <typename T, size_t BLOCK_SIZE, bool REVERSE_ORDER = false>
+class BlockStore {
+protected:
+  struct Block {
+    alignas(T) uint8_t data[BLOCK_SIZE * sizeof(T)] = {0};
+    Block *next = nullptr;
+  };
+
+  Block first;
+  Block *current = &first;
+  size_t fill_count = 0;
+
+public:
+  constexpr BlockStore() = default;
+  ~BlockStore() = default;
+
+  class iterator {
+    Block *block;
+    size_t index;
+
+  public:
+    constexpr iterator(Block *b, size_t i) : block(b), index(i) {}
+
+    iterator &operator++() {
+      if (REVERSE_ORDER) {
+        if (index == 0)
+          return *this;
+
+        --index;
+        if (index == 0 && block->next != nullptr) {
+          index = BLOCK_SIZE;
+          block = block->next;
+        }
+      } else {
+        if (index == BLOCK_SIZE)
+          return *this;
+
+        ++index;
+        if (index == BLOCK_SIZE && block->next != nullptr) {
+          index = 0;
+          block = block->next;
+        }
+      }
+
+      return *this;
+    }
+
+    T &operator*() {
+      size_t true_index = REVERSE_ORDER ? index - 1 : index;
+      return *reinterpret_cast<T *>(block->data + sizeof(T) * true_index);
+    }
+
+    bool operator==(const iterator &rhs) const {
+      return block == rhs.block && index == rhs.index;
+    }
+
+    bool operator!=(const iterator &rhs) const {
+      return block != rhs.block || index != rhs.index;
+    }
+  };
+
+  static void destroy(BlockStore<T, BLOCK_SIZE, REVERSE_ORDER> *block_store);
+
+  T *new_obj() {
+    if (fill_count == BLOCK_SIZE) {
+      auto new_block = reinterpret_cast<Block *>(::malloc(sizeof(Block)));
+      if (REVERSE_ORDER) {
+        new_block->next = current;
+      } else {
+        new_block->next = nullptr;
+        current->next = new_block;
+      }
+      current = new_block;
+      fill_count = 0;
+    }
+    T *obj = reinterpret_cast<T *>(current->data + fill_count * sizeof(T));
+    ++fill_count;
+    return obj;
+  }
+
+  void push_back(const T &value) {
+    T *ptr = new_obj();
+    *ptr = value;
+  }
+
+  iterator begin() {
+    if (REVERSE_ORDER)
+      return iterator(current, fill_count);
+    else
+      return iterator(&first, 0);
+  }
+
+  iterator end() {
+    if (REVERSE_ORDER)
+      return iterator(&first, 0);
+    else
+      return iterator(current, fill_count);
+  }
+};
+
+template <typename T, size_t BLOCK_SIZE, bool REVERSE_ORDER>
+void BlockStore<T, BLOCK_SIZE, REVERSE_ORDER>::destroy(
+    BlockStore<T, BLOCK_SIZE, REVERSE_ORDER> *block_store) {
+  if (REVERSE_ORDER) {
+    auto current = block_store->current;
+    while (current->next != nullptr) {
+      auto temp = current;
+      current = current->next;
+      free(temp);
+    }
+  } else {
+    auto current = block_store->first.next;
+    while (current != nullptr) {
+      auto temp = current;
+      current = current->next;
+      free(temp);
+    }
+  }
+  block_store->current = nullptr;
+  block_store->fill_count = 0;
+}
+
+// A convenience type for reverse order block stores.
+template <typename T, size_t BLOCK_SIZE>
+using ReverseOrderBlockStore = BlockStore<T, BLOCK_SIZE, true>;
+
+} // namespace cpp
+} // namespace __llvm_libc
+
+#endif // LLVM_LIBC_SUPPORT_CPP_BLOCKSTORE_H
diff --git a/libc/src/stdlib/CMakeLists.txt b/libc/src/stdlib/CMakeLists.txt
index 692134727164..e1bf45f5a3bb 100644
--- a/libc/src/stdlib/CMakeLists.txt
+++ b/libc/src/stdlib/CMakeLists.txt
@@ -276,8 +276,10 @@ add_entrypoint_object(
     atexit.cpp
   HDRS
     atexit.h
+  CXX_STANDARD
+    20 # For constinit of the atexit callback list.
   DEPENDS
-    libc.src.__support.CPP.vector
+    libc.src.__support.CPP.blockstore
     libc.src.__support.threads.thread
 )
 
diff --git a/libc/src/stdlib/atexit.cpp b/libc/src/stdlib/atexit.cpp
index 31943e09b61d..1548ea15c35b 100644
--- a/libc/src/stdlib/atexit.cpp
+++ b/libc/src/stdlib/atexit.cpp
@@ -7,7 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/stdlib/atexit.h"
-#include "src/__support/CPP/vector.h"
+#include "src/__support/CPP/blockstore.h"
 #include "src/__support/common.h"
 #include "src/__support/threads/mutex.h"
 
@@ -17,29 +17,29 @@ namespace {
 
 Mutex handler_list_mtx(false, false, false);
 
-// TOOD should we make cpp::vector like llvm::SmallVector<T, N> where it will
-// allocate at least N before needing dynamic allocation?
-static cpp::vector<void (*)(void)> handlers;
+using AtExitCallback = void(void);
+using ExitCallbackList = cpp::ReverseOrderBlockStore<AtExitCallback *, 32>;
+constinit ExitCallbackList exit_callbacks;
 
 } // namespace
 
 namespace internal {
 
-void call_exit_handlers() {
+void call_exit_callbacks() {
   handler_list_mtx.lock();
-  // TODO: implement rbegin() + rend() for cpp::vector
-  for (int i = handlers.size() - 1; i >= 0; i--) {
+  for (auto callback : exit_callbacks) {
     handler_list_mtx.unlock();
-    handlers[i]();
+    callback();
     handler_list_mtx.lock();
   }
+  ExitCallbackList::destroy(&exit_callbacks);
 }
 
 } // namespace internal
 
-LLVM_LIBC_FUNCTION(int, atexit, (void (*function)())) {
+LLVM_LIBC_FUNCTION(int, atexit, (AtExitCallback * callback)) {
   handler_list_mtx.lock();
-  handlers.push_back(function);
+  exit_callbacks.push_back(callback);
   handler_list_mtx.unlock();
   return 0;
 }
diff --git a/libc/src/stdlib/exit.cpp b/libc/src/stdlib/exit.cpp
index 5a02a45d4f18..92faea4da6ad 100644
--- a/libc/src/stdlib/exit.cpp
+++ b/libc/src/stdlib/exit.cpp
@@ -13,11 +13,11 @@
 namespace __llvm_libc {
 
 namespace internal {
-void call_exit_handlers();
+void call_exit_callbacks();
 }
 
 LLVM_LIBC_FUNCTION(void, exit, (int status)) {
-  internal::call_exit_handlers();
+  internal::call_exit_callbacks();
   _Exit(status);
 }
 
diff --git a/libc/test/src/__support/CPP/CMakeLists.txt b/libc/test/src/__support/CPP/CMakeLists.txt
index 3464f68ccadd..9e8e5ae0d964 100644
--- a/libc/test/src/__support/CPP/CMakeLists.txt
+++ b/libc/test/src/__support/CPP/CMakeLists.txt
@@ -69,3 +69,13 @@ add_libc_unittest(
   DEPENDS
     libc.src.__support.CPP.atomic
 )
+
+add_libc_unittest(
+  blockstore_test
+  SUITE
+    libc_cpp_utils_unittests
+  SRCS
+    blockstore_test.cpp
+  DEPENDS
+    libc.src.__support.CPP.blockstore
+)
diff --git a/libc/test/src/__support/CPP/blockstore_test.cpp b/libc/test/src/__support/CPP/blockstore_test.cpp
new file mode 100644
index 000000000000..0fb40143e0d1
--- /dev/null
+++ b/libc/test/src/__support/CPP/blockstore_test.cpp
@@ -0,0 +1,66 @@
+//===-- Unittests for BlockStore ------------------------------------------===//
+//
+// 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/__support/CPP/blockstore.h"
+#include "utils/UnitTest/Test.h"
+
+struct Element {
+  int a;
+  long b;
+  unsigned c;
+};
+
+class LlvmLibcBlockStoreTest : public __llvm_libc::testing::Test {
+public:
+  template <size_t BLOCK_SIZE, size_t ELEMENT_COUNT, bool REVERSE>
+  void populate_and_iterate() {
+    __llvm_libc::cpp::BlockStore<Element, BLOCK_SIZE, REVERSE> block_store;
+    for (int i = 0; i < int(ELEMENT_COUNT); ++i)
+      block_store.push_back({i, 2 * i, 3 * unsigned(i)});
+    auto end = block_store.end();
+    int i = 0;
+    for (auto iter = block_store.begin(); iter != end; ++iter, ++i) {
+      Element &e = *iter;
+      if (REVERSE) {
+        int j = ELEMENT_COUNT - 1 - i;
+        ASSERT_EQ(e.a, j);
+        ASSERT_EQ(e.b, long(j * 2));
+        ASSERT_EQ(e.c, unsigned(j * 3));
+      } else {
+        ASSERT_EQ(e.a, i);
+        ASSERT_EQ(e.b, long(i * 2));
+        ASSERT_EQ(e.c, unsigned(i * 3));
+      }
+    }
+    ASSERT_EQ(i, int(ELEMENT_COUNT));
+  }
+};
+
+TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterate4) {
+  populate_and_iterate<4, 4, false>();
+}
+
+TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterate8) {
+  populate_and_iterate<4, 8, false>();
+}
+
+TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterate10) {
+  populate_and_iterate<4, 10, false>();
+}
+
+TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterateReverse4) {
+  populate_and_iterate<4, 4, true>();
+}
+
+TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterateReverse8) {
+  populate_and_iterate<4, 8, true>();
+}
+
+TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterateReverse10) {
+  populate_and_iterate<4, 10, true>();
+}