gc.cpp

#include <iostream>
#include <type_traits>
#include <vector>
#include <cstring>

#include "gc.h"


template<class T>
static inline T* align(T* ptr, std::size_t align) {
    std::uintptr_t p = reinterpret_cast<std::uintptr_t>(ptr);
    return reinterpret_cast<T*>((p + align - 1) & ~(align - 1));
}

namespace psil {
    std::vector<Type>& get_type_info() {
        static std::vector<Type> type_info;
        return type_info;
    }

    const Type& get_type(TypeId id) { return get_type_info()[id]; }

    TypeId register_type(std::string name, std::size_t size,
                         std::function<void(void*)> visitor,
                         std::function<void(void*, void*)> mover,
                         std::function<void(void*)> destructor) {
        TypeId id = get_type_info().size();
        get_type_info().push_back(Type{id, name, size, visitor, mover, destructor});
        return id;
    }


    const Type& GCObject::type() const { return get_type(id); }
    TypeId GCObject::type_id() const { return id; }

    void* GCObject::data() { return reinterpret_cast<char*>(this) + sizeof(GCObject); }

    GCObject* GCObject::from_data(void* data) {
        return reinterpret_cast<GCObject*>(reinterpret_cast<char*>(data) - sizeof(GCObject));
    }

    std::size_t GCObject::size_on_heap() const {
        const char* base = reinterpret_cast<const char*>(this);
        return align(base + sizeof(GCObject) + type().size, alignof(GCObject)) - base;
    }


    struct GC {
        GC() : heap(128), heap_ptr(heap.data()), alive_size(0), updating_locations(false) { }

        ~GC() { 
            char* iter_heap_ptr = heap.data();
            GCObject* obj;
            while ((obj = next_on_heap(iter_heap_ptr, heap_ptr))) {
                auto type = obj->type();
                if (type.destructor) type.destructor(obj->data());
            }
        }

        GCObject* alloc(TypeId type_id) {
            auto type = get_type(type_id);

            heap_ptr = align(heap_ptr, alignof(GCObject));
            std::size_t heap_free = heap.data() + heap.size() - heap_ptr;
            if (heap_free < sizeof(GCObject) + type.size) {
                collect(alignof(GCObject) + sizeof(GCObject) + type.size + 128);
                heap_ptr = align(heap_ptr, alignof(GCObject));
            }

            auto obj_location = heap_ptr;
            heap_ptr += sizeof(GCObject) + type.size;
            return new (obj_location) GCObject(type_id);
        }

        void visit(GCObject*& obj, bool weak) {
            if (updating_locations) {
                if (obj) obj = obj->new_location;
                return;
            }

            if (!obj || obj->marked || weak) return;
            obj->marked = true;

            alive_size += obj->size_on_heap();
            auto visitor = obj->type().visitor;
            if (visitor) visitor(obj->data());
        }

        void collect(std::size_t extra_space) {
            // Mark alive objects and get total size of alive objects.
            alive_size = 0;
            for (auto ref : refs) visit(ref->obj, false);

            // Allocate new heap.
            std::vector<char> new_heap(alignof(GCObject) + alive_size * 1.5 + extra_space);
            
            // Loop over all objects, copy to new heap if alive, destruct if dead.
            char* iter_heap_ptr = heap.data();
            char* new_heap_ptr = new_heap.data();
            GCObject* obj;
            while ((obj = next_on_heap(iter_heap_ptr, heap_ptr))) move_if_alive(obj, new_heap_ptr);

            // Use new heap.
            heap_ptr = new_heap_ptr;
            heap.swap(new_heap); // Swap to keep old heap alive until we update locations.

            update_locations();
        }

    private:
        GCObject* next_on_heap(char*& heap_ptr, char* heap_end) {
            if (heap_ptr >= heap_end) return nullptr;
            heap_ptr = align(heap_ptr, alignof(GCObject));
            auto obj = reinterpret_cast<GCObject*>(heap_ptr);
            heap_ptr += sizeof(GCObject) + obj->type().size;
            return obj;
        }

        void move_if_alive(GCObject* obj, char*& new_heap_ptr) {
            auto type = obj->type();

            GCObject* new_obj = nullptr;
            if (obj->marked) {
                new_heap_ptr = align(new_heap_ptr, alignof(GCObject));
                new_obj = reinterpret_cast<GCObject*>(new_heap_ptr);
                new (new_obj) GCObject(obj->id);
                new_heap_ptr += sizeof(GCObject);
                if (type.mover) type.mover(new_heap_ptr, obj->data());
                else std::memcpy(new_heap_ptr, obj->data(), type.size);
                new_heap_ptr += type.size;
            }

            if (type.destructor) type.destructor(obj->data());
            obj->new_location = new_obj;
        }

        void update_locations() {
            updating_locations = true;
            // All references are still alive, we can just assign without checks.
            for (auto ref : refs) ref->obj = ref->obj->new_location;

            // Weak references might get removed, iterate over copy.
            std::vector<GCWeakRef*> weakrefs_copy = weakrefs;
            for (auto ref : weakrefs_copy) ref->reset(ref->obj->new_location);

            auto iter_heap_ptr = heap.data();
            GCObject* obj;
            while ((obj = next_on_heap(iter_heap_ptr, heap_ptr))) {
                auto type = obj->type();
                if (type.visitor) type.visitor(obj->data());
            }
            updating_locations = false;
        }

        std::vector<char> heap;
        char* heap_ptr;
        std::size_t alive_size;
        bool updating_locations;
        std::vector<GCRef*> refs;
        std::vector<GCWeakRef*> weakrefs;

        friend class GCRef;
        friend class GCWeakRef;
    };

    GC& get_gc() {
        static GC gc;
        return gc;
    }

    void GCRef::ref() {
        auto& refs = get_gc().refs;
        ref_id = refs.size();
        refs.push_back(this);
    }

    void GCRef::unref() {
        auto& refs = get_gc().refs;
        refs.back()->ref_id = ref_id;
        refs[ref_id] = std::move(refs.back());
        refs.pop_back();
    }

    void GCWeakRef::ref() {
        auto& weakrefs = get_gc().weakrefs;
        ref_id = weakrefs.size();
        weakrefs.push_back(this);
    }

    void GCWeakRef::unref() {
        auto& weakrefs = get_gc().weakrefs;
        weakrefs.back()->ref_id = ref_id;
        weakrefs[ref_id] = std::move(weakrefs.back());
        weakrefs.pop_back();
    }

    void GCVisitRef::visit(bool weak) const { get_gc().visit(obj, weak); }


    GCObject* gc_alloc(TypeId type_id) { return get_gc().alloc(type_id); }
    void gc_collect(std::size_t extra_space) { return get_gc().collect(extra_space); }
}

gc.h

#ifndef PSIL_GC_H
#define PSIL_GC_H

#include <cstdint>
#include <type_traits>
#include <string>
#include <functional>
#include <unordered_set>
#include <vector>

namespace psil {
    // Type information for GC objects.
    typedef uint16_t TypeId;

    struct Type {
        const TypeId id;                              
        const std::string name;                       
        const std::size_t size;                       
        const std::function<void(void*)> visitor;     
        const std::function<void(void*, void*)> mover;
        const std::function<void(void*)> destructor;  
    };
    
    const Type& get_type(TypeId id);

    TypeId register_type(std::string name, std::size_t size,
                         std::function<void(void*)> visitor,
                         std::function<void(void*, void*)> mover,
                         std::function<void(void*)> destructor);

    template<class T> TypeId register_class_as_type(std::string name);

    // Base GC object.
    struct GCObject {
        GCObject(TypeId id) : id(id), marked(false) { }
        
        const Type& type() const;
        TypeId type_id() const;
        void* data();
        static GCObject* from_data(void* data);

    private:
        std::size_t size_on_heap() const;

        union {
            GCObject* new_location;
            struct { TypeId id; bool marked; };
        };

        friend class GC;
    };

    // Reference types.
    struct GCRef; struct GCWeakRef; struct GCVisitRef;

    // Code deduplication without virtual pointers.
#define PSIL_REF_BODY(C)                                                      \
        C() : obj(nullptr) { }                                                \
        C(GCObject* obj) : obj(obj) { if (obj) ref(); }                       \
        C(const GCRef& other);                                                \
        C(const GCWeakRef& other);                                            \
        C(const GCVisitRef& other);                                           \
        C& operator=(const C& other) { reset(other.obj); return *this; }      \
        C(C&& other) : obj(other.obj) { other.reset(); if (obj) ref(); }      \
        ~C() { reset(); }                                                     \
                                                                              \
        void swap(C& other) {                                                 \
            std::swap(obj, other.obj);                                        \
            if (bool(obj) == bool(other.obj)) return;                         \
            if (obj) { ref(); other.unref(); } else { unref(); other.ref(); } \
        }                                                                     \
                                                                              \
        void reset() { if (obj) { unref(); obj = nullptr; } }                 \
        void reset(GCObject* obj) {                                           \
            if (this->obj && !obj) unref();                                   \
            if (!this->obj && obj) ref();                                     \
            this->obj = obj;                                                  \
        }                                                                     \
                                                                              \
        GCObject* get() const { return obj; }                                 \
        explicit operator bool() const { return obj != nullptr; }             \
        GCObject& operator*() const { return *obj; }                          \
        GCObject* operator->() const { return obj; }                          \
                                                                              \
    private:                                                                  \
        mutable GCObject* obj;                                                \
                                                                              \
        friend class GC;
    
    // Outside reference to GC owned object.
    struct GCRef {
        PSIL_REF_BODY(GCRef)
        std::uint32_t ref_id;
        void ref(); void unref();
    };

    // Outside weak reference to GC owned object.
    struct GCWeakRef {
        PSIL_REF_BODY(GCWeakRef)
        std::uint32_t ref_id;
        void ref(); void unref();
    };
    
    // Reference to use inside GC types. Tracked by calling visit().
    struct GCVisitRef {
        void visit(bool weak=false) const;
        PSIL_REF_BODY(GCVisitRef);
        void ref() { } void unref() { }
    };
#undef PSIL_REF_BODY

    // Garbage collector interface.
    GCObject* gc_alloc(TypeId type_id);
    template<class T, class... Args>
    GCRef gc_new(TypeId type_id, Args&&... args);
    void gc_collect(std::size_t extra_space=0);
}



// Implementation.
namespace psil {
    namespace detail {
        template<class T, class=void> struct has_visit : std::false_type { };
        template<class T>
        struct has_visit<T, decltype(std::declval<T>().visit())>
        : std::true_type { };

        template<class T>
        inline typename std::enable_if<has_visit<T>::value,
        std::function<void(void*)>>::type get_visitor()
            { return [](void* obj) { static_cast<T*>(obj)->visit(); }; }
        
        template<class T>
        inline typename std::enable_if<!has_visit<T>::value,
        std::function<void(void*)>>::type get_visitor() { return nullptr; }
    }

    template<class T>
    inline TypeId register_class_as_type(std::string name) {
        static_assert(alignof(T) <= alignof(GCObject),
                      "maximum alignment requirement is that of GCObject");
        
        std::function<void(void*, void*)> mover;
        std::function<void(void*)> destructor;

        if (!std::is_pod<T>::value) {
            mover = [](void* dest, void* src)
                { new (dest) T(std::move(*static_cast<T*>(src))); };
            destructor = [](void* obj) { static_cast<T*>(obj)->~T(); };
        }

        return register_type(name, sizeof(T),
                             detail::get_visitor<T>(), mover, destructor);
    }

    template<class T, class... Args>
    inline GCRef gc_new(TypeId type_id, Args&&... args) {
        GCRef obj = gc_alloc(type_id);
        new (obj->data()) T{std::forward<Args>(args)...};
        return obj;
    }

    // More code deduplication for the reference types.
#define PSIL_REF_EXTRAS(C)                                                     \
    inline C::C(const GCRef& other) : obj(other.get()) { if (obj) ref(); }     \
    inline C::C(const GCWeakRef& other) : obj(other.get()) { if (obj) ref(); } \
    inline C::C(const GCVisitRef& other) : obj(other.get()) { }                \
    inline bool operator==(const C& lhs, const C& rhs)                         \
        { return lhs.get() == rhs.get(); }                                     \
    inline bool operator==(const C& lhs, std::nullptr_t rhs)                   \
        { return lhs.get() == nullptr; }                                       \
    inline bool operator==(std::nullptr_t lhs, const C& rhs)                   \
        { return nullptr == rhs.get(); }                                       \
    inline bool operator!=(const C& lhs, const C& rhs)                         \
        { return !(lhs == rhs); }                                              \
    inline bool operator!=(const C& lhs, std::nullptr_t rhs)                   \
        { return !(lhs == rhs); }                                              \
    inline bool operator!=(std::nullptr_t lhs, const C& rhs)                   \
        { return !(lhs == rhs); }                                         
    
    PSIL_REF_EXTRAS(GCRef)
    PSIL_REF_EXTRAS(GCWeakRef)
    PSIL_REF_EXTRAS(GCVisitRef)
#undef PSIL_REF_EXTRAS
}

#endif