valeri/src/arena.hpp

#pragma once

#include <concepts>

#include "error.hpp"
#include "pod.hpp"
#include "result.hpp"

class Arena;
class GcRootList;

class GcRootBase {
 public:
  GcRootBase(PodObject* ptr, GcRootList* node);

  ~GcRootBase();

 protected:
  PodObject* _ptr;
  GcRootList* _node;
};

template <class T>
  requires std::derived_from<T, PodObject>
class GcRoot : public GcRootBase {
 public:
  GcRoot() : GcRootBase(0, 0) {}
  GcRoot(T* ptr, GcRootList* node) : GcRootBase(ptr, node) {}
  GcRoot(GcRoot&& rhs);

  static Result<GcRoot<T>> create(T* ptr, Arena& arena);

  T& operator*() { return *(T*)_ptr; }
};

template <class T>
  requires std::derived_from<T, PodObject>
Result<GcRoot<T>> MkGcRoot(T* ptr, Arena& arena) {
  return GcRoot<T>::create(ptr, arena);
}

class GcRootList {
 public:
  GcRootList() : _prev(0), _next(0), _root(0) {}

  GcRootList(GcRootBase* root) : _prev(0), _next(0), _root(root) {}
  void insert(GcRootList* node) {
    GcRootList* next = _next;

    _next = node;
    node->_next = next;
    node->_prev = this;
  }

  void remove() {
    _prev = 0;
    _next = 0;
    _prev->_next = _next;
    _next->_prev = _prev;
  }

  void update(GcRootBase* root) { _root = root; }

 private:
  GcRootList* _prev;
  GcRootList* _next;
  GcRootBase* _root;
};

class ArenaHeap {
 public:
  ArenaHeap(uint8_t* buf, uint64_t bufsize)
      : buf(buf), bufsize(bufsize), boundary(0) {}

  Result<void*> alloc(uint64_t size) {
    // Always align to 64 bits
    if (size % 8 != 0) size += 8 - size % 8;

    if (boundary + size >= bufsize) return ErrorCode::OutOfMemory;
    void* res = buf + boundary;
    boundary += size;
    return res;
  }

 private:
  uint8_t* buf;
  uint64_t bufsize;
  uint64_t boundary;
};

class Arena {
 public:
  Arena(ArenaHeap* first, ArenaHeap* second)
      : _heaps{first, second}, _gcroot(), _current(0) {}

  template <class T>
  Result<T*> alloc(uint64_t extra = 0) {
    uint64_t objsize = sizeof(T) + extra;
    auto ptr = _heaps[_current]->alloc(objsize);

    if (ptr.has_value()) return (T*)ptr.value();

    // TODO: trigger GC

    return (T*)TRY(_heaps[_current]->alloc(objsize));
  }

  void add_root(GcRootList* node) { _gcroot.insert(node); }

 private:
  ArenaHeap* _heaps[2];
  GcRootList _gcroot;
  int _current;
};

template <uint64_t size>
class StaticArenaHeap {
 public:
  StaticArenaHeap() : _heap(_buf, heapsize) {}

  ArenaHeap* get() { return &_heap; }

 private:
  static const uint64_t heapsize = size - sizeof(ArenaHeap);
  ArenaHeap _heap;
  uint8_t _buf[heapsize]{0};
};

template <uint64_t size>
class StaticArena {
 public:
  StaticArena() : _arena(_heaps[0].get(), _heaps[1].get()) {}

 private:
  StaticArenaHeap<(size - sizeof(Arena)) / 2> _heaps[2];
  Arena _arena;
};

template <class T>
  requires std::derived_from<T, PodObject>
Result<GcRoot<T>> GcRoot<T>::create(T* ptr, Arena& arena) {
  auto lst = TRY(arena.alloc<GcRootList>());
  arena.add_root(lst);
  return GcRoot<T>(ptr, lst);
}

template <class T>
  requires std::derived_from<T, PodObject>
GcRoot<T>::GcRoot(GcRoot<T>&& rhs) {
  rhs._node->update(this);
}