memgraph/poc/memory/block_allocator.hpp

#pragma once

#include <memory>
#include <vector>

#include "utils/on_scope_exit.hpp"

/* @brief Allocates blocks of block_size and stores
 * the pointers on allocated blocks inside a vector.
 */
template <size_t block_size>
class BlockAllocator {
  struct Block {
    Block() { data = malloc(block_size); }

    Block(void *ptr) { data = ptr; }

    void *data;
  };

 public:
  static constexpr size_t size = block_size;

  BlockAllocator(size_t capacity = 0) {
    for (size_t i = 0; i < capacity; ++i) unused_.emplace_back();
  }

  ~BlockAllocator() {
    for (auto block : unused_) free(block.data);
    unused_.clear();
    for (auto block : release_) free(block.data);
    release_.clear();
  }

  size_t unused_size() const { return unused_.size(); }

  size_t release_size() const { return release_.size(); }

  // Returns nullptr on no memory.
  void *acquire() {
    if (unused_.size() == 0) unused_.emplace_back();

    auto ptr = unused_.back().data;
    // TODO is it necessary to use OnScopeExit here? ptr is copied by value, right?
    utils::OnScopeExit on_exit([this]() { unused_.pop_back(); });
    return ptr;
  }

  void release(void *ptr) { release_.emplace_back(ptr); }

 private:
  // TODO: try implement with just one vector
  // but consecutive acquire release calls should work
  // TODO: measure first!
  std::vector<Block> unused_;
  std::vector<Block> release_;
};