Toy implementation of mark&sweep GC

mark_sweep.cc

#include <stdio.h>
#include <cassert>
#include <stdexcept>

enum class ObjectType
{
  Int,
  Pair
};

class VM;
class Object
{
public:
  ObjectType type;
  bool marked;
  Object* next;

  union {
    int value;
    struct {
      Object* head;
      Object* tail;
    };
  };

  Object(VM& vm, ObjectType type) noexcept;
};

class VM
{
public:
  static constexpr int STACK_MAX = 256;
  static constexpr int INITIAL_GC_THRESHOLD = 8;

  Object* stack[STACK_MAX];
  int stack_size;
  Object* first_obj;
  int num_objects;
  int max_objects;

  VM() noexcept;
  ~VM();
  void push(Object& value);
  Object& pop();
  void push_int(int value);
  Object& push_pair();
  void mark(Object& obj) noexcept;
  void mark_all() noexcept;
  void sweep() noexcept;
  void gc() noexcept;
  void clear() noexcept;
};

Object::Object(VM& vm, ObjectType type) noexcept
  : type(type)
  , marked()
  , next(vm.first_obj)
{
  if (vm.num_objects == vm.max_objects)
    vm.gc();
  vm.first_obj = this;
  ++vm.num_objects;
}

VM::VM() noexcept
  : stack_size()
  , first_obj()
  , num_objects()
  , max_objects(INITIAL_GC_THRESHOLD)
{
}

VM::~VM()
{
  clear();
}

void
VM::push(Object& value)
{
  if (stack_size >= STACK_MAX)
    throw std::overflow_error("Stack overflow!");
  stack[stack_size++] = &value;
}

Object&
VM::pop()
{
  if (stack_size <= 0)
    throw std::underflow_error("Stack underflow!");
  return *stack[--stack_size];
}

void
VM::push_int(int value)
{
  Object* obj = new Object(*this, ObjectType::Int);
  obj->value = value;
  push(*obj);
}

Object&
VM::push_pair()
{
  Object* obj = new Object(*this, ObjectType::Pair);
  obj->tail = &pop();
  obj->head = &pop();
  push(*obj);
  return *obj;
}

void
VM::mark(Object& obj) noexcept
{
  if (obj.marked)
    return;

  obj.marked = true;

  if (obj.type == ObjectType::Pair)
  {
    mark(*obj.head);
    mark(*obj.tail);
  }
}

void
VM::mark_all() noexcept
{
  for (int i = 0; i < stack_size; ++i)
    mark(*stack[i]);
}

void
VM::sweep() noexcept
{
  Object** obj = &first_obj;
  while (*obj)
  {
    if ((*obj)->marked)
    {
      (*obj)->marked = false;
      obj = &(*obj)->next;
    }
    else
    {
      Object* unreached = *obj;
      *obj = unreached->next;
      delete unreached;
      --num_objects;
    }
  }
}

void
VM::gc() noexcept
{
  int old_num_objs = num_objects;
  mark_all();
  sweep();

  max_objects = num_objects << 1;
  printf("Collected %d objects, %d remaining.\n",
         old_num_objs - num_objects, num_objects);
}

void
VM::clear() noexcept
{
  stack_size = 0;
  gc();
}

int main(int argc, char* argv[])
{
  {
    printf("Test 1: Objects on stack are preserved.\n");
    VM vm;
    vm.push_int(1);
    vm.push_int(2);
    vm.gc();
    assert(vm.num_objects == 2 && "Should have preserved objects.");
  }
  {
    printf("Test 2: Unreached objects are collected.\n");
    VM vm;
    vm.push_int(1);
    vm.push_int(2);
    vm.pop();
    vm.pop();
    vm.gc();
    assert(vm.num_objects == 0 && "Should have collected objects.");
  }
  {
    printf("Test 3: Reach nested objects.\n");
    VM vm;
    vm.push_int(1);
    vm.push_int(2);
    vm.push_pair();
    vm.push_int(3);
    vm.push_int(4);
    vm.push_pair();
    vm.push_pair();
    vm.gc();
    assert(vm.num_objects == 7 && "Should have reached objects.");
  }
  {
    printf("Test 4: Handle cycles.\n");
    VM vm;
    vm.push_int(1);
    vm.push_int(2);
    Object& a = vm.push_pair();
    vm.push_int(3);
    vm.push_int(4);
    Object& b = vm.push_pair();
    a.tail = &b;
    b.tail = &a;
    vm.gc();
    assert(vm.num_objects == 4 && "Should have collected objects.");
  }
  {
    printf("Performance Test.\n");
    VM vm;
    for (int i = 0; i < 1000; ++i)
    {
      for (int j = 0; j < 20; ++j)
        vm.push_int(j);
      for (int j = 0; j < 20; ++j)
        vm.pop();
    }
  }
}