Use Arena allocator

CMakeLists.txt

@@ -7,12 +7,16 @@ project(
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
+add_compile_options(-fno-exceptions -fno-rtti)
+add_link_options(-nodefaultlibs -lc)
+
 add_library(conflict_set ConflictSet.cpp ConflictSet.h)
 
 add_executable(conflict_set_test ConflictSet.cpp ConflictSet.h)
 target_compile_definitions(conflict_set_test PRIVATE ENABLE_TESTS)
 # keep asserts for test
 target_compile_options(conflict_set_test PRIVATE -UNDEBUG)
+# Only emit compile warnings for test
 target_compile_options(conflict_set_test PRIVATE -Wall -Wextra -Wpedantic -Wunreachable-code)
 
 include(CTest)

ConflictSet.cpp

@@ -20,6 +20,172 @@ static auto operator<=>(const Key &lhs, const Key &rhs) {
   return c != 0 ? c <=> 0 : lhs.len <=> rhs.len;
 }
 
+// ==================== BEGIN ARENA IMPL ====================
+
+/// Group allocations with similar lifetimes to amortize the cost of malloc/free
+struct Arena {
+  explicit Arena(int initialSize = 0);
+  /// O(log n) in the number of allocations
+  ~Arena();
+  struct ArenaImpl;
+  Arena(const Arena &) = delete;
+  Arena &operator=(const Arena &) = delete;
+  Arena(Arena &&other) noexcept;
+  Arena &operator=(Arena &&other) noexcept;
+
+private:
+  ArenaImpl *impl = nullptr;
+  friend void *operator new(size_t size, std::align_val_t align, Arena &arena);
+};
+
+inline void operator delete(void *, std::align_val_t, Arena &) {}
+void *operator new(size_t size, std::align_val_t align, Arena &arena);
+void *operator new(size_t size, std::align_val_t align, Arena *arena) = delete;
+
+inline void operator delete(void *, Arena &) {}
+inline void *operator new(size_t size, Arena &arena) {
+  return operator new(size, std::align_val_t(alignof(std::max_align_t)), arena);
+}
+inline void *operator new(size_t size, Arena *arena) = delete;
+
+inline void operator delete[](void *, Arena &) {}
+inline void *operator new[](size_t size, Arena &arena) {
+  return operator new(size, arena);
+}
+inline void *operator new[](size_t size, Arena *arena) = delete;
+
+inline void operator delete[](void *, std::align_val_t, Arena &) {}
+inline void *operator new[](size_t size, std::align_val_t align, Arena &arena) {
+  return operator new(size, align, arena);
+}
+inline void *operator new[](size_t size, std::align_val_t align,
+                            Arena *arena) = delete;
+
+/// align must be a power of two
+template <class T> T *align_up(T *t, size_t align) {
+  auto unaligned = uintptr_t(t);
+  auto aligned = (unaligned + align - 1) & ~(align - 1);
+  return reinterpret_cast<T *>(reinterpret_cast<char *>(t) + aligned -
+                               unaligned);
+}
+
+/// align must be a power of two
+constexpr inline int align_up(uint32_t unaligned, uint32_t align) {
+  return (unaligned + align - 1) & ~(align - 1);
+}
+
+/// Returns the smallest power of two >= x
+constexpr inline uint32_t nextPowerOfTwo(uint32_t x) {
+  return x <= 1 ? 1 : 1 << (32 - __builtin_clz(x - 1));
+}
+
+/// \private
+struct Arena::ArenaImpl {
+  Arena::ArenaImpl *prev;
+  int capacity;
+  int used;
+  uint8_t *begin() { return reinterpret_cast<uint8_t *>(this + 1); }
+};
+
+static_assert(sizeof(Arena::ArenaImpl) == 16);
+static_assert(alignof(Arena::ArenaImpl) == 8);
+
+Arena::Arena(int initialSize) : impl(nullptr) {
+  if (initialSize > 0) {
+    auto allocationSize = align_up(initialSize + sizeof(ArenaImpl), 16);
+    impl = (Arena::ArenaImpl *)malloc(allocationSize);
+    impl->prev = nullptr;
+    impl->capacity = allocationSize - sizeof(ArenaImpl);
+    impl->used = 0;
+  }
+}
+
+namespace {
+void onDestroy(Arena::ArenaImpl *impl) {
+  while (impl) {
+    auto *prev = impl->prev;
+    free(impl);
+    impl = prev;
+  }
+}
+} // namespace
+
+Arena::Arena(Arena &&other) noexcept
+    : impl(std::exchange(other.impl, nullptr)) {}
+Arena &Arena::operator=(Arena &&other) noexcept {
+  onDestroy(impl);
+  impl = std::exchange(other.impl, nullptr);
+  return *this;
+}
+
+Arena::~Arena() { onDestroy(impl); }
+
+void *operator new(size_t size, std::align_val_t align, Arena &arena) {
+  int64_t aligned_size = size + size_t(align) - 1;
+  if (arena.impl == nullptr ||
+      (arena.impl->capacity - arena.impl->used) < aligned_size) {
+    auto allocationSize = align_up(
+        sizeof(Arena::ArenaImpl) +
+            std::max<int>(aligned_size,
+                          (arena.impl ? std::max<int>(sizeof(Arena::ArenaImpl),
+                                                      arena.impl->capacity * 2)
+                                      : 0)),
+        16);
+    auto *impl = (Arena::ArenaImpl *)malloc(allocationSize);
+    impl->prev = arena.impl;
+    impl->capacity = allocationSize - sizeof(Arena::ArenaImpl);
+    impl->used = 0;
+    arena.impl = impl;
+  }
+  auto *result =
+      align_up(arena.impl->begin() + arena.impl->used, size_t(align));
+  auto usedDelta = (result - arena.impl->begin()) + size - arena.impl->used;
+  arena.impl->used += usedDelta;
+  return result;
+}
+
+/// STL-friendly allocator using an arena
+template <class T> struct ArenaAlloc {
+  typedef T value_type;
+
+  ArenaAlloc() = delete;
+  explicit ArenaAlloc(Arena *arena) : arena(arena) {}
+
+  Arena *arena;
+
+  template <class U> constexpr ArenaAlloc(const ArenaAlloc<U> &other) noexcept {
+    arena = other.arena;
+  }
+
+  [[nodiscard]] T *allocate(size_t n) {
+    if (n > 0xfffffffffffffffful / sizeof(T)) { // NOLINT
+      fprintf(stderr, "Requested bad alloc! sizeof(T): %zu, n: %zu\n",
+              sizeof(T), n); // NOLINT
+      fflush(stderr);
+      abort();
+    }
+
+    return static_cast<T *>((void *)new (std::align_val_t(alignof(T)), *arena)
+                                uint8_t[n * sizeof(T)]); // NOLINT
+  }
+
+  void deallocate(T *, size_t) noexcept {}
+
+private:
+};
+
+template <class T, class U>
+bool operator==(const ArenaAlloc<T> &lhs, const ArenaAlloc<U> &rhs) {
+  return lhs.arena == rhs.arena;
+}
+
+template <class T, class U>
+bool operator!=(const ArenaAlloc<T> &lhs, const ArenaAlloc<U> &rhs) {
+  return !(lhs == rhs);
+}
+
+// ==================== END ARENA IMPL ====================
+
 namespace {
 // A node in the tree representing write conflict history. This tree maintains
 // several invariants:
@@ -171,15 +337,15 @@ struct StepwiseLastLeq {
   }
 };
 
-void lastLeqMulti(Node *root, std::span<Key> keys, Iterator *results) {
+void lastLeqMulti(Arena &arena, Node *root, std::span<Key> keys,
+                  Iterator *results) {
   assert(std::is_sorted(keys.begin(), keys.end()));
 
   if (keys.size() == 0) {
     return;
   }
 
-  auto stepwiseLastLeqs =
+  auto *stepwiseLastLeqs = new (arena) StepwiseLastLeq[keys.size()];
-      std::unique_ptr<StepwiseLastLeq[]>{new StepwiseLastLeq[keys.size()]};
 
   // Descend until queries for front and back diverge
   Node *current = root;
@@ -201,13 +367,12 @@ void lastLeqMulti(Node *root, std::span<Key> keys, Iterator *results) {
 
   int index = 0;
   {
-    auto iter = stepwiseLastLeqs.get();
+    auto iter = stepwiseLastLeqs;
     for (const auto &k : keys) {
       *iter++ = StepwiseLastLeq(current, resultP, k, index++);
     }
   }
-  auto stepwiseSpan =
+  auto stepwiseSpan = std::span<StepwiseLastLeq>(stepwiseLastLeqs, keys.size());
-      std::span<StepwiseLastLeq>(stepwiseLastLeqs.get(), keys.size());
   runInterleaved(stepwiseSpan);
   for (const auto &stepwise : stepwiseSpan) {
     results[stepwise.index] = Iterator{stepwise.result, stepwise.resultC};
@@ -300,14 +465,12 @@ void lastLeqMulti(Node *root, std::span<Key> keys, Iterator *results) {
   fprintf(file, "}\n");
 }
 
-[[maybe_unused]] Key toKey(int n) {
+[[maybe_unused]] Key toKey(Arena &arena, int n) {
   constexpr int kMaxLength = 4;
   // TODO use arena allocation
-  static std::vector<std::vector<uint8_t>> *results =
-      new std::vector<std::vector<uint8_t>>;
   int i = kMaxLength;
-  results->push_back(std::vector<uint8_t>(kMaxLength, '0'));
+  uint8_t *itoaBuf = new (arena) uint8_t[kMaxLength];
-  uint8_t *itoaBuf = results->back().data();
+  memset(itoaBuf, '0', kMaxLength);
   do {
     itoaBuf[--i] = "0123456789abcdef"[n % 16];
     n /= 16;
@@ -401,7 +564,9 @@ int64_t checkMaxVersion(Node *node, bool &success) {
 bool checkInvariants(Node *node) {
   bool success = true;
   // Check bst invariant
-  std::vector<std::string_view> keys;
+  Arena arena;
+  std::vector<std::string_view, ArenaAlloc<std::string_view>> keys{
+      ArenaAlloc<std::string_view>(&arena)};
   for (auto iter = extrema(node, false); iter != nullptr;
        iter = next(iter, true)) {
     keys.push_back(std::string_view((char *)(iter + 1), iter->len));
@@ -438,12 +603,13 @@ struct ConflictSet::Impl {
   }
 
   void check(const ReadRange *reads, Result *results, int count) const {
-    auto iters = std::unique_ptr<Iterator[]>{new Iterator[count]};
+    Arena arena;
-    auto begins = std::unique_ptr<Key[]>{new Key[count]};
+    auto *iters = new (arena) Iterator[count];
+    auto *begins = new (arena) Key[count];
     for (int i = 0; i < count; ++i) {
-      begins.get()[i] = reads[i].begin;
+      begins[i] = reads[i].begin;
     }
-    lastLeqMulti(root, std::span<Key>(begins.get(), count), iters.get());
+    lastLeqMulti(arena, root, std::span<Key>(begins, count), iters);
     // TODO check non-singleton reads lol
     for (int i = 0; i < count; ++i) {
       assert(reads[i].end.len == 0);
@@ -501,8 +667,8 @@ struct ConflictSet::Impl {
   };
 
   void addWrites(const WriteRange *writes, int count) {
-    auto stepwiseInserts =
+    Arena arena;
-        std::unique_ptr<StepwiseInsert[]>{new StepwiseInsert[count]};
+    auto *stepwiseInserts = new (arena) StepwiseInsert[count];
     for (int i = 0; i < count; ++i) {
       // TODO handle non-singleton writes lol
       assert(writes[i].end.len == 0);
@@ -518,11 +684,12 @@ struct ConflictSet::Impl {
     // probably fine.
     // TODO better/faster RNG?
     std::mt19937 g(fastRand());
-    std::shuffle(stepwiseInserts.get(), stepwiseInserts.get() + count, g);
+    std::shuffle(stepwiseInserts, stepwiseInserts + count, g);
 
-    runInterleaved(std::span<StepwiseInsert>(stepwiseInserts.get(), count));
+    runInterleaved(std::span<StepwiseInsert>(stepwiseInserts, count));
 
-    std::vector<Node *> workList;
+    std::vector<Node *, ArenaAlloc<Node *>> workList{
+        ArenaAlloc<Node *>(&arena)};
     workList.reserve(count);
     for (int i = 0; i < count; ++i) {
       workList.push_back(*stepwiseInserts[i].current);
@@ -568,7 +735,8 @@ struct ConflictSet::Impl {
   }
 
   ~Impl() {
-    std::vector<Node *> toFree;
+    Arena arena;
+    std::vector<Node *, ArenaAlloc<Node *>> toFree{ArenaAlloc<Node *>(&arena)};
     if (root != nullptr) {
       toFree.push_back(root);
     }
@@ -600,9 +768,12 @@ void ConflictSet::setOldestVersion(int64_t oldestVersion) {
 }
 
 ConflictSet::ConflictSet(int64_t oldestVersion)
-    : impl(new Impl{oldestVersion}) {}
+    : impl(new(malloc(sizeof(Impl))) Impl{oldestVersion}) {}
 
-ConflictSet::~ConflictSet() { delete impl; }
+ConflictSet::~ConflictSet() {
+  impl->~Impl();
+  free(impl);
+}
 
 ConflictSet::ConflictSet(ConflictSet &&other) noexcept
     : impl(std::exchange(other.impl, nullptr)) {}
@@ -618,8 +789,9 @@ int main(void) {
   ConflictSet::Impl cs{writeVersion};
   constexpr int kNumKeys = 5;
   ConflictSet::WriteRange write[kNumKeys];
+  Arena arena;
   for (int i = 0; i < kNumKeys; ++i) {
-    write[i].begin = toKey(i);
+    write[i].begin = toKey(arena, i);
     write[i].end.len = 0;
     write[i].writeVersion = ++writeVersion;
   }