main.cpp

#include <assert.h>
#include <stdio.h>

#include <set>
#include <vector>

// INTERFACE

class GcObject;

class GarbageCollector
{
    friend class GcObject;
    friend class GcPinner;

    // Pointers to variables representing entrypoints to the object graph.
    std::set<GcObject**> rootPointers;

    // Pointers to the actual objects themshelves.
    std::set<GcObject*> gcObjects;

    // The flag value used for flagging object in the current run.
    int currentFlag;
public:
    GarbageCollector()
    {
        currentFlag = 0;
    }

    // Collect all lost objects using mark and sweep method.
    void collectGarbage();

    // Just queries some stats: the count of objects.
    int getObjectCount() {return gcObjects.size();}

    ~GarbageCollector();
};


/**
 * Initially all GC managed objects exist alone, not connected into the graph.
 * A GC sweep would clean them up right away.
 * In order to prevent that. Use this RAII class to register those stack based pointers as roots.
 */

class GcPinner
{
    GarbageCollector *gc;
    GcObject **obj;
public:
    // Don't copy this.
    GcPinner(const GcPinner&) = delete;
    GcPinner& operator=(const GcPinner&) = delete;

    // Register variable as root.
    GcPinner(GarbageCollector *gc, GcObject *&obj)
        : obj(&obj)
    {
        this->gc = gc;
        gc->rootPointers.insert(this->obj);
    }

    // Unregister variable when cleared up.
    ~GcPinner()
    {
        gc->rootPointers.erase(obj);
    }
};


class GcObject
{
    friend class GarbageCollector;

    int flag;
    GarbageCollector &gc;

protected:
    /* This function should be used by the implementations of markReferred().
     Call this function on all GcObject* members or its derived versions to mark them alive during garbage collection. */
    void mark();
    // Derived classes must implement this to be able to mark their members that are referred by them. (This assistance is needed from the user.)
    virtual void markReferred() = 0;
public:
    // Don't copy this.
    GcObject(const GcObject&) = delete;
    GcObject& operator=(const GcObject&) = delete;

    // Registers the objects in the given GC upon creation.
    GcObject(GarbageCollector &collector) :
        gc(collector)
    {
        flag = collector.currentFlag;
        collector.gcObjects.insert(this);
    }

    // Upon destruction, unregister the object from the GC:
    virtual ~GcObject()
    {
        gc.gcObjects.erase(this);
    };
};

// IMPLEMENTATION

void GarbageCollector::collectGarbage()
{
    // Set a new unique flag for this sweep.
    currentFlag++;

    // Starting from the roots, mark all objects as reachable.
    for (auto x : rootPointers)
    {
        if (*x == nullptr) continue;
        (*x)->mark();
    }

    // Finally scan the memory for garbage (unmarked objects)
    std::vector<GcObject*> toClear;
    for (auto x : gcObjects)
    {
        if (x->flag != currentFlag) toClear.push_back(x);
    }

    // Bury the corpses.
    for (auto x : toClear)
    {
        delete x;
    }
}

void GcObject::mark()
{
    if (flag != gc.currentFlag)
    {
        flag = gc.currentFlag;
        markReferred();
    }
}

GarbageCollector::~GarbageCollector()
{
    // Purge all managed objects
    std::vector<GcObject*> toClear;
    /* We need to put object pointers into another array, because their destructors would remove them from the gcObjects collections, 
       thus modifying the collection while enumeration, which is bad.*/
    for (auto x: gcObjects)
    {
        toClear.push_back(x);
    }

    for (auto x: toClear)
    {
        delete x;
    }
}

// TEST CODE

GarbageCollector GC;

class GraphNode : public GcObject
{
    int payload;
protected:
    virtual void markReferred()
    {
        for (size_t i = 0; i < referredNodes.size(); i++)
        {
            referredNodes[i]->mark();
        }
    }
public:
    std::vector<GraphNode*> referredNodes;

    GraphNode(int pl)
        : GcObject(GC)
    {
        payload = pl;
    }

    GraphNode(int pl, const std::vector<GraphNode*> children) 
        : GcObject(GC),
        referredNodes(children)
    {
        payload = pl;
    }

    ~GraphNode()
    {
        printf("Node: %d collected\n", payload);
    }
};

int main()
{
    {
        GcObject *root;
        GcPinner pin(&GC, root);
        
        // Make a graph.
        root = new GraphNode(1, {new GraphNode(2), new GraphNode(3), new GraphNode(4)});

        assert(GC.getObjectCount() == 4); // We have the 4 objects
        GC.collectGarbage();
        assert(GC.getObjectCount() == 4); // We still have the 4 objects.
        
        root = nullptr; // Lose the graph. 
        printf("Everything should be collected now.\n");
        GC.collectGarbage();
        assert(GC.getObjectCount() == 0); // Everything is collected.

        // Allocate some orphans
        new GraphNode(42);
        new GraphNode(666);
        assert(GC.getObjectCount() == 2); // We know about them.
        printf("Everything should be collected now.\n");
        GC.collectGarbage();
        assert(GC.getObjectCount() == 0); // Everything is collected.

        // Allocate a chain.
        root = new GraphNode(1, {new GraphNode(2, {new GraphNode(3, {new GraphNode(4)})})});
        assert(GC.getObjectCount() == 4);
        GC.collectGarbage();
        assert(GC.getObjectCount() == 4); // All is still alive.
        // Cut the chain
        GraphNode *gn = dynamic_cast<GraphNode*>(root);
        gn->referredNodes[0]->referredNodes.clear(); // Node 2 should lose it's referred stuff.
        GC.collectGarbage();
        assert(GC.getObjectCount() == 2); // Node 1 and 2 are exist, 3 and 4 are collected now.
    }

    GC.collectGarbage();
    assert(GC.getObjectCount() == 0); // The pin is gone, now it's OK to clean up everything.

    printf("All is OK!\n");
}